Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions HISTORY.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,11 @@ Changelog
==========


14.7.3 (2026-08-03)
-------------------

* Initial release for DSS 14.7.3

14.7.2 (2026-07-13)
-------------------

Expand Down
7 changes: 5 additions & 2 deletions dataikuapi/dss/agent_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ def id(self):
"""
return self.tool_id

def get_descriptor(self):
def get_descriptor(self, context=None):
"""
Get the descriptor of the tool

Expand All @@ -73,7 +73,10 @@ def get_descriptor(self):
"""

if self._descriptor is None:
self._descriptor = self.client._perform_json("GET", "/projects/%s/agents/tools/%s/descriptor" % (self.project_key, self.tool_id))
if context is None:
self._descriptor = self.client._perform_json("GET", "/projects/%s/agents/tools/%s/descriptor" % (self.project_key, self.tool_id))
else:
self._descriptor = self.client._perform_json("POST", "/projects/%s/agents/tools/%s/descriptor" % (self.project_key, self.tool_id), body={"context": context})
return self._descriptor

def get_settings(self):
Expand Down
13 changes: 11 additions & 2 deletions dataikuapi/dss/langchain/embeddings.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,25 @@
import asyncio
import concurrent
import logging
import threading
import threading
import itertools

from typing import Callable, List, Any, Union

import pydantic

_thread_pool_executor_counter = itertools.count().__next__

def next_thread_pool_executor_prefix(prefix):
return "{}-{}".format(prefix, _thread_pool_executor_counter())

try:
from langchain_core.embeddings.embeddings import Embeddings
except ModuleNotFoundError:
from langchain.embeddings.base import Embeddings
from langchain_core.callbacks import BaseCallbackHandler, LLMManagerMixin


from dataikuapi.dss.llm_tracing import new_trace, SpanBuilder

from dataikuapi.dss.langchain.utils import must_use_deprecated_pydantic_config
Expand Down Expand Up @@ -121,7 +130,7 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]:

async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
loop = asyncio.get_event_loop()
with concurrent.futures.ThreadPoolExecutor() as executor:
with concurrent.futures.ThreadPoolExecutor(thread_name_prefix=next_thread_pool_executor_prefix("DKUEmbeddingsAsyncExecutor")) as executor:
result = await loop.run_in_executor(executor, self.embed_documents, texts)
return result

Expand Down
158 changes: 146 additions & 12 deletions dataikuapi/dss/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,8 @@ def add_image(self, image, text = None):
def new_guardrail(self, type):
"""
Start adding a guardrail to the request. You need to configure the returned object, and call add() to actually add it

:rtype: :class:`DSSLLMRequestGuardrailBuilder`
"""
return DSSLLMRequestGuardrailBuilder(self, type)

Expand Down Expand Up @@ -420,15 +422,27 @@ def with_structured_output(self, model_type, strict=None, compatible=None):


class DSSLLMRequestGuardrailBuilder(object):
"""
.. important::
Do not create this class directly, use :meth:`dataikuapi.dss.llm.DSSLLMCompletionQuery.new_guardrail`,
:meth:`dataikuapi.dss.llm.DSSLLMCompletionsQuery.new_guardrail`, :meth:`dataikuapi.dss.llm.DSSLLMEmbeddingsQuery.new_guardrail` or
:meth:`dataikuapi.dss.llm.DSSLLMImageGenerationQuery.new_guardrail`.
"""

def __init__(self, request, type):
self.request = request
self.guardrail = { "type" : type, "enabled": True, "params" : {}}
self.guardrail = {"type" : type, "enabled": True, "params" : {}}

@property
def params(self):
"""
:return: The parameters of this guardrail
:rtype: dict
"""
return self.guardrail["params"]

def add(self):
"""Add this guardrail to the completion query"""
if self.request._guardrails is None:
self.request._guardrails = {"guardrails" : []}
self.request._guardrails["guardrails"].append(self.guardrail)
Expand Down Expand Up @@ -537,6 +551,8 @@ def settings(self):
def new_guardrail(self, type):
"""
Start adding a guardrail to the request. You need to configure the returned object, and call add() to actually add it

:rtype: :class:`DSSLLMRequestGuardrailBuilder`
"""
return DSSLLMRequestGuardrailBuilder(self, type)

Expand Down Expand Up @@ -638,6 +654,8 @@ def new_completion(self):
def new_guardrail(self, type):
"""
Start adding a guardrail to the request. You need to configure the returned object, and call add() to actually add it

:rtype: :class:`DSSLLMRequestGuardrailBuilder`
"""
return DSSLLMRequestGuardrailBuilder(self, type)

Expand All @@ -664,7 +682,7 @@ def execute(self):
return DSSLLMCompletionsResponse(ret["responses"], response_parser=self._response_parser)


class DSSLLMCompletionQueryMultipartBuilder(object):
class _DSSLLMCompletionQueryMultipartBuilder(object):
def __init__(self):
self.parts = []

Expand All @@ -681,6 +699,8 @@ def _encode_image(image):
def with_text(self, text):
"""
Add a text part to the multipart message

:param str text: The text to add
"""
self.parts.append({"type": "TEXT", "text": text})
return self
Expand All @@ -692,7 +712,7 @@ def with_inline_image(self, image, mime_type=None):
:param Union[str, bytes] image: The image
:param str mime_type: None for default
"""
img_b64 = DSSLLMCompletionQueryMultipartMessage._encode_image(image)
img_b64 = _DSSLLMCompletionQueryMultipartBuilder._encode_image(image)

part = {
"type": "IMAGE_INLINE",
Expand All @@ -713,7 +733,7 @@ def with_captioned_image_inline(self, caption, image, mime_type=None):
:param Union[str, bytes] image: The image
:param str mime_type: None for default
"""
img_b64 = DSSLLMCompletionQueryMultipartMessage._encode_image(image)
img_b64 = _DSSLLMCompletionQueryMultipartBuilder._encode_image(image)

image_part = {
"type": "IMAGE_INLINE",
Expand All @@ -736,14 +756,13 @@ def with_image_url(self, image):
"""
Add an image url part to the multipart message

:param image: str the image url
:param str image: the image url
"""

self.parts.append({"type": "IMAGE_URI", "imageUrl": image})
return self


class DSSLLMCompletionQueryMultipartMessage(DSSLLMCompletionQueryMultipartBuilder):
class DSSLLMCompletionQueryMultipartMessage(_DSSLLMCompletionQueryMultipartBuilder):
"""
.. important::
Do not create this class directly, use :meth:`dataikuapi.dss.llm.DSSLLMCompletionQuery.new_multipart_message` or
Expand All @@ -761,8 +780,43 @@ def add(self):
self.q.cq["messages"].append(self.msg)
return self.q

def with_text(self, text):
"""
Add a text part to the multipart message

:param str text: The text to add
"""
return super().with_text(text)

def with_inline_image(self, image, mime_type=None):
"""
Add an image part to the multipart message

:param Union[str, bytes] image: The image
:param str mime_type: None for default
"""
return super().with_inline_image(image, mime_type)

class DSSLLMCompletionQueryMultipartToolOutput(DSSLLMCompletionQueryMultipartBuilder):
def with_captioned_image_inline(self, caption, image, mime_type=None):
"""
Add a captioned image part to the multipart message

:param str caption: Image caption
:param Union[str, bytes] image: The image
:param str mime_type: None for default
"""
return super().with_captioned_image_inline(caption, image, mime_type)

def with_image_url(self, image):
"""
Add an image url part to the multipart message

:param str image: The image url
"""
return super().with_image_url(image)


class DSSLLMCompletionQueryMultipartToolOutput(_DSSLLMCompletionQueryMultipartBuilder):
"""
.. important::
Do not create this class directly, use :meth:`dataikuapi.dss.llm.DSSLLMCompletionQuery.new_multipart_tool_output` or
Expand All @@ -787,41 +841,107 @@ def add(self):
self.q.cq["messages"].append(self.msg)
return self.q

def with_text(self, text):
"""
Add a text part to the multipart tool output

:param str text: The text to add
"""
return super().with_text(text)

def with_inline_image(self, image, mime_type=None):
"""
Add an image part to the multipart tool output

:param Union[str, bytes] image: The image
:param str mime_type: None for default
"""
return super().with_inline_image(image, mime_type)

def with_captioned_image_inline(self, caption, image, mime_type=None):
"""
Add a captioned image part to the multipart tool output

:param str caption: Image caption
:param Union[str, bytes] image: The image
:param str mime_type: None for default
"""
return super().with_captioned_image_inline(caption, image, mime_type)

def with_image_url(self, image):
"""
Add an image url part to the multipart tool output

:param str image: The image url
"""
return super().with_image_url(image)


class DSSLLMStreamedCompletionChunk(object):
"""
A handle to interact with a streamed completion query chunk.

.. important::
Do not create this class directly, iterate over a :class:`dataikuapi.dss.llm.DSSLLMStreamedCompletionChunks` iterator instead to generate the chunks instead.
"""

def __init__(self, data):
self.data = data

@property
def type(self):
"""Type of this chunk, either "content" or "event" """
"""
:return: Type of this chunk, either "content" or "event"
:rtype: Literal["content", "event"]
"""
return self.data.get("type", "content")

@property
def text(self):
"""If this chunk is content and has text, the (partial) text"""
"""
:return: If this chunk is content and has text, the (partial) text
:rtype: bool
"""
return self.data.get("text", None)

@property
def event_kind(self):
"""If this chunk is an event, its kind"""
"""
:return: If this chunk is an event, its kind
:rtype: str
"""
return self.data.get("eventKind", None)

def __repr__(self):
return "<completion-chunk: %s>" % self.data


class DSSLLMStreamedCompletionFooter(object):
"""
A handle to interact with a streamed completion query footer.

.. important::
Do not create this class directly, iterate over a :class:`dataikuapi.dss.llm.DSSLLMStreamedCompletionChunks` iterator instead to generate the chunks instead.
"""

def __init__(self, data):
self.data = data

# Compatibility for code that just checks for "type""
@property
def type(self):
"""
:return: Type of this chunk, to distinguish it from :class:`dataikuapi.dss.llm.DSSLLMStreamedCompletionChunk` chunks. Can only be "footer"
:rtype: Literal["footer"]
"""
return "footer"

@property
def trace(self):
"""
:return: The trace of the completion query if available, None otherwise.
:rtype: Union[dict, None]
"""
return self.data.get("trace", None)

@property
Expand Down Expand Up @@ -890,7 +1010,11 @@ def iterevents(self):

class DSSLLMCompletionResponse(object):
"""
Response to a completion
A handle to interact with a completion query result.

.. important::
Do not create this class directly, use :meth:`dataikuapi.dss.llm.DSSLLMCompletionQuery.execute` or
:attr:`dataikuapi.dss.llm.DSSLLMCompletionsResponse.responses` or :attr:`dataikuapi.dss.llm.DSSLLMStreamedCompletionChunks.response` instead.
"""
def __init__(self, raw_resp=None, text=None, finish_reason=None, response_parser=None, trace=None, query=None):
if raw_resp is not None:
Expand Down Expand Up @@ -990,6 +1114,10 @@ def context_upsert(self):

@property
def trace(self):
"""
:return: The trace of the completion query if available, None otherwise.
:rtype: Union[dict, None]
"""
return self._raw.get("trace", None)

@property
Expand Down Expand Up @@ -1157,6 +1285,8 @@ def with_mask(self, mode, image=None):
def new_guardrail(self, type):
"""
Start adding a guardrail to the request. You need to configure the returned object, and call add() to actually add it

:rtype: :class:`DSSLLMRequestGuardrailBuilder`
"""
return DSSLLMRequestGuardrailBuilder(self, type)

Expand Down Expand Up @@ -1381,6 +1511,10 @@ def images(self):

@property
def trace(self):
"""
:return: The trace of the image generation query if available, None otherwise.
:rtype: Union[dict, None]
"""
return self._raw.get("trace", None)

@property
Expand Down
Loading