unique_toolkit 1.46.10__tar.gz → 1.46.11__tar.gz

{unique_toolkit-1.46.10 → unique_toolkit-1.46.11}/CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.46.11] - 2026-02-20
+- Add OpenAI-proxy direct streaming with Python-side reference injection (UN-17264)
+- Add `stream_complete_with_references_openai()` using async OpenAI client and `chat.completions.stream()`
+- Add stream transformation pipeline: `StreamTransform` protocol, `TextTransformPipeline`, `ReferenceInjectionTransform` (reuses `add_references_to_message`), `NormalizationTransform` placeholder
+- Export `stream_complete_with_references_openai` and stream transform types from `unique_toolkit.language_model`
+
 ## [1.46.10] - 2026-02-19
 - Add flag-based cancellation support via `CancellationWatcher` and `TypedEventBus`
 - Expose `ChatService.cancellation` property for polling, event subscription, and `run_with_cancellation` helper

{unique_toolkit-1.46.10 → unique_toolkit-1.46.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 1.46.10
+Version: 1.46.11
 Summary: 
 License: Proprietary
 Author: Cedric Klinkert
@@ -127,6 +127,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.46.11] - 2026-02-20
+- Add OpenAI-proxy direct streaming with Python-side reference injection (UN-17264)
+- Add `stream_complete_with_references_openai()` using async OpenAI client and `chat.completions.stream()`
+- Add stream transformation pipeline: `StreamTransform` protocol, `TextTransformPipeline`, `ReferenceInjectionTransform` (reuses `add_references_to_message`), `NormalizationTransform` placeholder
+- Export `stream_complete_with_references_openai` and stream transform types from `unique_toolkit.language_model`
+
 ## [1.46.10] - 2026-02-19
 - Add flag-based cancellation support via `CancellationWatcher` and `TypedEventBus`
 - Expose `ChatService.cancellation` property for polling, event subscription, and `run_with_cancellation` helper

{unique_toolkit-1.46.10 → unique_toolkit-1.46.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unique_toolkit"
-version = "1.46.10"
+version = "1.46.11"
 description = ""
 authors = [
     "Cedric Klinkert <cedric.klinkert@unique.ch>",

{unique_toolkit-1.46.10 → unique_toolkit-1.46.11}/unique_toolkit/chat/cancellation.py

@@ -18,6 +18,7 @@ class CancellationEvent:
 
     message_id: str
 
+
 T = TypeVar("T")
 
 logger = logging.getLogger(__name__)

{unique_toolkit-1.46.10 → unique_toolkit-1.46.11}/unique_toolkit/language_model/__init__.py

@@ -5,6 +5,9 @@ from .default_language_model import (
 from .default_language_model import (
     DEFAULT_GPT_4o as DEFAULT_GPT_4o,
 )
+from .functions import (
+    stream_complete_with_references_openai as stream_complete_with_references_openai,
+)
 from .infos import LanguageModel as LanguageModel
 from .infos import LanguageModelName as LanguageModelName
 from .prompt import (
@@ -67,6 +70,18 @@ from .schemas import (
 from .service import (
     LanguageModelService as LanguageModelService,
 )
+from .stream_transform import (
+    NormalizationTransform as NormalizationTransform,
+)
+from .stream_transform import (
+    ReferenceInjectionTransform as ReferenceInjectionTransform,
+)
+from .stream_transform import (
+    StreamTransform as StreamTransform,
+)
+from .stream_transform import (
+    TextTransformPipeline as TextTransformPipeline,
+)
 from .utils import (
     convert_string_to_json as convert_string_to_json,
 )

{unique_toolkit-1.46.10 → unique_toolkit-1.46.11}/unique_toolkit/language_model/functions.py

@@ -2,7 +2,11 @@ import copy
 import logging
 import warnings
 from datetime import UTC, datetime
-from typing import Any, Sequence, cast
+from typing import TYPE_CHECKING, Any, Sequence, cast
+from uuid import uuid4
+
+if TYPE_CHECKING:
+    from unique_toolkit.app.unique_settings import UniqueSettings
 
 import humps
 import unique_sdk
@@ -12,15 +16,6 @@ from pydantic import BaseModel
 
 from unique_toolkit.chat.schemas import ChatMessage, ChatMessageRole
 from unique_toolkit.content.schemas import ContentChunk, ContentReference
-from unique_toolkit.language_model import (
-    LanguageModelMessageRole,
-    LanguageModelMessages,
-    LanguageModelResponse,
-    LanguageModelStreamResponse,
-    LanguageModelStreamResponseMessage,
-    LanguageModelTool,
-    LanguageModelToolDescription,
-)
 from unique_toolkit.language_model.infos import (
     LanguageModelInfo,
     LanguageModelName,
@@ -29,6 +24,16 @@ from unique_toolkit.language_model.infos import (
 from unique_toolkit.language_model.reference import (
     add_references_to_message,
 )
+from unique_toolkit.language_model.schemas import (
+    LanguageModelFunction,
+    LanguageModelMessageRole,
+    LanguageModelMessages,
+    LanguageModelResponse,
+    LanguageModelStreamResponse,
+    LanguageModelStreamResponseMessage,
+    LanguageModelTool,
+    LanguageModelToolDescription,
+)
 
 from .constants import (
     DEFAULT_COMPLETE_TEMPERATURE,
@@ -541,3 +546,127 @@ def _create_language_model_stream_response_with_references(
         message=stream_response_message,
         tool_calls=tool_calls,
     )
+
+
+async def stream_complete_with_references_openai(
+    messages: list[ChatCompletionMessageParam] | LanguageModelMessages,
+    model_name: LanguageModelName | str,
+    content_chunks: list[ContentChunk] | None = None,
+    temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+    start_text: str | None = None,
+    tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
+    tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
+    unique_settings: "UniqueSettings | None" = None,  # noqa: F821
+    message_id: str | None = None,
+) -> LanguageModelStreamResponse:
+    """Stream a chat completion via the OpenAI proxy, apply reference injection, and return a LanguageModelStreamResponse."""
+    from unique_toolkit.framework_utilities.openai.client import (
+        get_async_openai_client,
+    )
+    from unique_toolkit.language_model.stream_transform import (
+        ReferenceInjectionTransform,
+        TextTransformPipeline,
+    )
+
+    # prepare variables
+    client = get_async_openai_client(unique_settings=unique_settings)
+    model_name_: str = (
+        model_name.value if isinstance(model_name, LanguageModelName) else model_name
+    )
+    messages_: list[ChatCompletionMessageParam] = (
+        messages.model_dump(
+            exclude_none=True,
+            by_alias=False,
+        )
+        if isinstance(messages, LanguageModelMessages)
+        else list(messages)
+    )
+    tools_ = (
+        []
+        if tools is None
+        else [
+            tool.to_openai(mode="completions")
+            if isinstance(tool, LanguageModelToolDescription)
+            else {"type": "function", "function": tool.model_dump(exclude_none=True)}
+            for tool in tools
+        ]
+    )
+
+    # pipeline
+    pipeline = TextTransformPipeline()
+    # pipeline.add(NormalizationTransform()) # no-op for now
+    pipeline.add(
+        ReferenceInjectionTransform(
+            content_chunks=content_chunks or [], model=model_name_
+        )
+    )
+
+    stream_kwargs: dict[str, Any] = {
+        "messages": messages_,
+        "model": model_name_,
+        "temperature": temperature,
+        "stream": True,
+    }
+    if tool_choice is not None:
+        stream_kwargs["tool_choice"] = tool_choice
+    if tools_:
+        stream_kwargs["tools"] = tools_
+
+    full_text: str = start_text or ""
+    tool_calls_fragments: dict[int, dict[str, Any]] = {}
+    try:
+        stream = await client.chat.completions.create(**stream_kwargs)
+        async with stream:
+            async for chunk in stream:
+                if not chunk.choices:
+                    continue
+                choice_ = chunk.choices[0]
+                # text
+                delta_ = choice_.delta
+                if delta_.content is not None:
+                    full_text += delta_.content
+                    pipeline.feed_delta(delta_.content)  # does nothing right now
+                # tools
+                for tool_call in delta_.tool_calls or []:
+                    idx_ = tool_call.index
+                    if idx_ not in tool_calls_fragments.keys():
+                        tool_calls_fragments[idx_] = {
+                            "id": tool_call.id or "",
+                            "name": "",
+                            "arguments": "",
+                        }
+                    if tool_call.function is not None:
+                        tool_calls_fragments[idx_]["name"] += (
+                            tool_call.function.name or ""
+                        )
+                        tool_calls_fragments[idx_]["arguments"] += (
+                            tool_call.function.arguments or ""
+                        )
+        # run on full text
+        text_, references_ = pipeline.run(full_text)
+        # construct tool calls
+        tool_calls_list = None
+        if len(tool_calls_fragments) > 0:
+            tool_calls_list = [
+                LanguageModelFunction(
+                    id=tool_call["id"],
+                    name=tool_call["name"],
+                    arguments=tool_call["arguments"] or None,
+                )
+                for tool_call in tool_calls_fragments.values()
+            ]
+
+        return LanguageModelStreamResponse(
+            message=LanguageModelStreamResponseMessage(
+                id=message_id or uuid4().hex,
+                previous_message_id=None,
+                role=LanguageModelMessageRole.ASSISTANT,
+                text=text_,
+                original_text=full_text,
+                references=references_,
+            ),
+            tool_calls=tool_calls_list,
+        )
+    except Exception as e:
+        logger.error("Error streaming completion (model=%s): %s", model_name_, e)
+        raise

{unique_toolkit-1.46.10 → unique_toolkit-1.46.11}/unique_toolkit/language_model/reference.py

@@ -53,7 +53,9 @@ def _add_references(
 
     # Gemini 2.5 models have tendency to add multiple references for the same fact
     # This is a workaround to limit the number of references to 5
-    if model and model.startswith("litellm:gemini-2-5"):
+    if model and model.startswith(
+        "litellm:gemini-2-5"
+    ):  # TODO - check if this holds also for other models
         reduced_text = _limit_consecutive_source_references(with_footnotes)
 
         # Get the references that remain after reduction

unique_toolkit-1.46.11/unique_toolkit/language_model/stream_transform.py

@@ -0,0 +1,82 @@
+from datetime import UTC, datetime
+from typing import Protocol, runtime_checkable
+
+from unique_toolkit.chat.schemas import ChatMessage, ChatMessageRole
+from unique_toolkit.content.schemas import ContentChunk, ContentReference
+from unique_toolkit.language_model.reference import add_references_to_message
+
+
+@runtime_checkable
+class StreamTransform(Protocol):
+    def process_delta(self, delta: str) -> None: ...
+
+    def finalize(self, text: str) -> tuple[str, list[ContentReference]]: ...
+
+
+class ReferenceInjectionTransform:
+    """Transform to inject references into the text of a message."""
+
+    def __init__(self, content_chunks: list[ContentChunk], model: str | None = None):
+        self._content_chunks = content_chunks
+        self._model = (
+            model  # needed e.g. for gemini-2.5 models to limit the number of references
+        )
+
+    def process_delta(self, delta: str) -> None:
+        pass  # cannot manipulate delta right now, as it is a stream
+
+    def finalize(self, text: str) -> tuple[str, list[ContentReference]]:
+        message = ChatMessage(
+            id="irrelevant",
+            text=text,
+            role=ChatMessageRole.ASSISTANT,
+            created_at=datetime.now(UTC),
+            chat_id="irrelevant",
+        )
+
+        updated_message, __ = add_references_to_message(
+            message=message,
+            search_context=self._content_chunks,
+            model=self._model,
+        )
+
+        references: list[ContentReference] = [
+            ContentReference(**u.model_dump()) for u in updated_message.references or []
+        ]
+        return updated_message.content or "", references
+
+
+class NormalizationTransform:
+    """Transform to normalize the text of a message.
+    E.g. remove extra whitespace, convert to lowercase, etc.
+    For now, this is a no-op. Only trim trailing whitespace.
+    """
+
+    def process_delta(self, delta: str) -> None:
+        pass
+
+    def finalize(self, text: str) -> tuple[str, list[ContentReference]]:
+        return text.strip(), []
+
+
+class TextTransformPipeline:
+    """Pipeline of text transforms to be applied to a stream of text."""
+
+    def __init__(self):
+        self._transforms: list[StreamTransform] = []
+
+    def add(self, transform: StreamTransform) -> "TextTransformPipeline":
+        self._transforms.append(transform)
+        return self
+
+    def feed_delta(self, delta: str) -> None:
+        for transform in self._transforms:
+            transform.process_delta(delta)
+
+    def run(self, text: str) -> tuple[str, list[ContentReference]]:
+        refs_: list[ContentReference] = []
+        text_: str = text
+        for transform in self._transforms:
+            text_, new_refs = transform.finalize(text_)
+            refs_.extend(new_refs)
+        return text_, refs_