unique_toolkit 1.57__tar.gz → 1.58__tar.gz

{unique_toolkit-1.57 → unique_toolkit-1.58}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 1.57
+Version: 1.58
 Summary: 
 Author: Cedric Klinkert, Martin Fadler, Sadique Sheik, Fabian Schläpfer, Pascal Hauri, Ejiroghene Oki-Peter, Rami Azouz, Ricardo Visini, Sanzio Monti
 Author-email: Cedric Klinkert <cedric.klinkert@unique.ch>, Martin Fadler <martin.fadler@unique.ch>, Sadique Sheik <sadique@unique.ch>, Fabian Schläpfer <fabian@unique.ch>, Pascal Hauri <pascal@unique.ch>, Ejiroghene Oki-Peter <oki-peter.ejiroghene@unique.ai>, Rami Azouz <rami.ext@unique.ch>, Ricardo Visini <ricardo.visini@unique.ch>, Sanzio Monti <sanzio.monti@unique.ai>

{unique_toolkit-1.57 → unique_toolkit-1.58}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "1.57"
+version = "1.58"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12"

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/agentic/evaluation/evaluation_manager.py

@@ -1,4 +1,5 @@
 import asyncio
+import time
 from abc import ABC
 from logging import Logger
 
@@ -86,6 +87,7 @@ class EvaluationManager:
         self._chat_service = chat_service
         self._evaluations: dict[EvaluationMetricName, Evaluation] = {}
         self._evaluation_passed: bool = True
+        self._execution_times: dict[str, float] = {}
 
     def add_evaluation(self, evaluation: Evaluation):
         self._evaluations[evaluation.get_name()] = evaluation
@@ -99,6 +101,8 @@ class EvaluationManager:
         loop_response: LanguageModelStreamResponse,
         assistant_message_id: str,
     ) -> list[EvaluationMetricResult]:
+        self._execution_times = {}
+
         task_executor = SafeTaskExecutor(
             logger=self._logger,
         )
@@ -134,6 +138,9 @@ class EvaluationManager:
 
         return evaluation_results_unpacked
 
+    def get_execution_times(self) -> dict[str, float]:
+        return self._execution_times.copy()
+
     async def execute_evaluation_call(
         self,
         evaluation_name: EvaluationMetricName,
@@ -145,13 +152,16 @@ class EvaluationManager:
         evaluation_instance = self.get_evaluation_by_name(evaluation_name)
 
         if evaluation_instance:
-            # Execute the evaluation
+            start = time.perf_counter()
             await self._create_assistant_message(
                 evaluation_instance, assistant_message_id
             )
             evaluation_metric_result: EvaluationMetricResult = (
                 await evaluation_instance.run(loop_response)
             )
+            self._execution_times[str(evaluation_name)] = round(
+                time.perf_counter() - start, 3
+            )
             return evaluation_metric_result
 
         return EvaluationMetricResult(

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/agentic/evaluation/tests/test_hallucination_utils.py

@@ -17,13 +17,13 @@ from unique_toolkit.agentic.evaluation.hallucination.utils import (
     context_text_from_stream_response,
 )
 from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricInput
+from unique_toolkit.chat.schemas import ChatMessage
 from unique_toolkit.content import ContentReference
 from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.language_model.schemas import (
     LanguageModelMessageRole,
     LanguageModelMessages,
     LanguageModelStreamResponse,
-    LanguageModelStreamResponseMessage,
 )
 
 
@@ -580,8 +580,9 @@ def test_context_text_from_stream_response__extracts_context__using_from_ids_mod
     """
     # Arrange
     response: LanguageModelStreamResponse = LanguageModelStreamResponse(
-        message=LanguageModelStreamResponseMessage(
+        message=ChatMessage(
             id="msg_1",
+            chat_id="",
             previous_message_id=None,
             role=LanguageModelMessageRole.ASSISTANT,
             text="Test",
@@ -612,8 +613,9 @@ def test_context_text_from_stream_response__extracts_context__using_from_order_m
     """
     # Arrange
     response: LanguageModelStreamResponse = LanguageModelStreamResponse(
-        message=LanguageModelStreamResponseMessage(
+        message=ChatMessage(
             id="msg_1",
+            chat_id="",
             previous_message_id=None,
             role=LanguageModelMessageRole.ASSISTANT,
             text="Test",
@@ -653,8 +655,9 @@ def test_context_text_from_stream_response__extracts_context__using_from_origina
         )
     ]
     response: LanguageModelStreamResponse = LanguageModelStreamResponse(
-        message=LanguageModelStreamResponseMessage(
+        message=ChatMessage(
             id="msg_1",
+            chat_id="",
             previous_message_id=None,
             role=LanguageModelMessageRole.ASSISTANT,
             text="Test",
@@ -690,8 +693,9 @@ def test_context_text_from_stream_response__falls_back_to_default__with_invalid_
     """
     # Arrange
     response: LanguageModelStreamResponse = LanguageModelStreamResponse(
-        message=LanguageModelStreamResponseMessage(
+        message=ChatMessage(
             id="msg_1",
+            chat_id="",
             previous_message_id=None,
             role=LanguageModelMessageRole.ASSISTANT,
             text="Test",
@@ -731,8 +735,9 @@ def test_context_text_from_stream_response__falls_back_to_default__on_extraction
         )
     ]
     response: LanguageModelStreamResponse = LanguageModelStreamResponse(
-        message=LanguageModelStreamResponseMessage(
+        message=ChatMessage(
             id="msg_1",
+            chat_id="",
             previous_message_id=None,
             role=LanguageModelMessageRole.ASSISTANT,
             text="Test",
@@ -763,8 +768,9 @@ def test_context_text_from_stream_response__returns_text_strings__not_chunk_obje
     """
     # Arrange
     response: LanguageModelStreamResponse = LanguageModelStreamResponse(
-        message=LanguageModelStreamResponseMessage(
+        message=ChatMessage(
             id="msg_1",
+            chat_id="",
             previous_message_id=None,
             role=LanguageModelMessageRole.ASSISTANT,
             text="Test",

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/agentic/history_manager/history_construction_with_contents.py

@@ -173,12 +173,12 @@ def get_full_history_with_contents(
 
     builder = LanguageModelMessages([]).builder()
     for c in grouped_elements:
-        # LanguageModelUserMessage has not field original content
-        text = c.original_content if c.original_content else c.content
+        # LanguageModelUserMessage has not field original_text
+        text = c.original_text if c.original_text else c.content
         if text is None:
             if c.role == ChatRole.USER:
                 raise ValueError(
-                    "Content or original_content of LanguageModelMessages should exist.",
+                    "Content or original_text of LanguageModelMessages should exist.",
                 )
             text = ""
 

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/agentic/postprocessor/postprocessor_manager.py

@@ -1,4 +1,5 @@
 import asyncio
+import time
 from abc import ABC
 from logging import Logger
 
@@ -84,6 +85,7 @@ class PostprocessorManager:
         self._logger = logger
         self._chat_service = chat_service
         self._postprocessors: list[Postprocessor | ResponsesApiPostprocessor] = []
+        self._execution_times: dict[str, float] = {}
 
         # Allow to add postprocessors that should be run before or after the others.
         self._first_postprocessor: Postprocessor | ResponsesApiPostprocessor | None = (
@@ -123,6 +125,8 @@ class PostprocessorManager:
         self,
         loop_response: LanguageModelStreamResponse,
     ) -> None:
+        self._execution_times = {}
+
         task_executor = SafeTaskExecutor(
             logger=self._logger,
         )
@@ -158,12 +162,19 @@ class PostprocessorManager:
                 references=loop_response.message.references,
             )
 
+    def get_execution_times(self) -> dict[str, float]:
+        return self._execution_times.copy()
+
     async def execute_postprocessors(
         self,
         loop_response: LanguageModelStreamResponse,
         postprocessor_instance: Postprocessor | ResponsesApiPostprocessor,
     ) -> None:
+        start = time.perf_counter()
         await postprocessor_instance.run(loop_response)  # type: ignore
+        self._execution_times[postprocessor_instance.name] = round(
+            time.perf_counter() - start, 3
+        )
 
     async def remove_from_text(
         self,

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/agentic/tools/tool_manager.py

@@ -1,4 +1,5 @@
 import asyncio
+import time
 from logging import Logger, getLogger
 from typing import Generic, Literal, TypeVar, overload
 
@@ -288,10 +289,16 @@ class _ToolManager(Generic[_ApiMode]):
             raise ValueError(f"Tool {tool_call.name} cannot be run")
 
         if tool_instance:
-            # Execute the tool
+            tool_start = time.perf_counter()
             tool_response: ToolCallResponse = await tool_instance.run(
                 tool_call=tool_call
             )
+            tool_execution_time = round(time.perf_counter() - tool_start, 3)
+
+            if tool_response.debug_info is None:
+                tool_response.debug_info = {}
+            tool_response.debug_info["execution_time_s"] = tool_execution_time
+
             evaluation_checks = tool_instance.evaluation_check_list()
             self._tool_evaluation_check_list.update(evaluation_checks)
 

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/chat/schemas.py

@@ -66,26 +66,67 @@ class ToolCall(BaseModel):
         )
 
 
+class ChatMessageAssessmentStatus(StrEnum):
+    PENDING = "PENDING"
+    DONE = "DONE"
+    ERROR = "ERROR"
+
+
+class ChatMessageAssessmentLabel(StrEnum):
+    RED = "RED"
+    YELLOW = "YELLOW"
+    GREEN = "GREEN"
+
+
+class ChatMessageAssessmentType(StrEnum):
+    HALLUCINATION = "HALLUCINATION"
+    COMPLIANCE = "COMPLIANCE"
+
+
+class ChatMessageAssessment(BaseModel):
+    model_config = model_config
+
+    id: str
+    object: str
+    message_id: str
+    status: ChatMessageAssessmentStatus
+    type: ChatMessageAssessmentType
+    title: str | None = None
+    explanation: str | None = None
+    label: ChatMessageAssessmentLabel | None = None
+    is_visible: bool
+    created_at: datetime | None = None
+    updated_at: datetime | None = None
+
+
 class ChatMessage(BaseModel):
-    # TODO: The below seems not to be True anymore @irina-unique. To be checked in separate PR
-    # This model should strictly meets https://github.com/Unique-AG/monorepo/blob/master/node/apps/node-chat/src/public-api/2023-12-06/dtos/message/public-message.dto.ts
+    # This model maps to PublicMessageDto from the backend public API:
+    # next/services/node-chat/src/public-api/2023-12-06/dtos/message/public-message.dto.ts
+    # Fields previous_message_id, tool_calls, tool_call_id are internal extensions not in the DTO.
     model_config = model_config
 
     id: str | None = None
     chat_id: str
     object: str | None = None
+    # alias="text" applies only to construction (model_validate / __init__).
+    # model_dump() uses the field name "content", not "text". Use model_dump(by_alias=True)
+    # to get "text" as the key, or access via the .text property.
     content: str | None = Field(default=None, alias="text")
-    original_content: str | None = Field(default=None, alias="originalText")
+    original_text: str | None = None
     role: ChatMessageRole
+    previous_message_id: str | None = None
     gpt_request: list[dict] | None = None
     tool_calls: list[ToolCall] | None = None
     tool_call_id: str | None = None
     debug_info: dict | None = {}
     created_at: datetime | None = None
     completed_at: datetime | None = None
+    started_streaming_at: datetime | None = None
+    stopped_streaming_at: datetime | None = None
     user_aborted_at: datetime | None = None
     updated_at: datetime | None = None
-    references: list[ContentReference] | None = None
+    references: list[ContentReference] | None = []
+    assessment: list[ChatMessageAssessment] | None = None
 
     # TODO make sdk return role consistently in lowercase
     # Currently needed as sdk returns role in uppercase
@@ -93,6 +134,14 @@ class ChatMessage(BaseModel):
     def set_role(cls, value: str):
         return value.lower()
 
+    @property
+    def text(self) -> str | None:
+        return self.content
+
+    @text.setter
+    def text(self, value: str | None) -> None:
+        self.content = value
+
     # Ensure tool_call_ids is required if role is 'tool'
     @model_validator(mode="after")
     def check_tool_call_ids_for_tool_role(self):
@@ -133,39 +182,6 @@ class ChatMessage(BaseModel):
                 raise NotImplementedError
 
 
-class ChatMessageAssessmentStatus(StrEnum):
-    PENDING = "PENDING"
-    DONE = "DONE"
-    ERROR = "ERROR"
-
-
-class ChatMessageAssessmentLabel(StrEnum):
-    RED = "RED"
-    YELLOW = "YELLOW"
-    GREEN = "GREEN"
-
-
-class ChatMessageAssessmentType(StrEnum):
-    HALLUCINATION = "HALLUCINATION"
-    COMPLIANCE = "COMPLIANCE"
-
-
-class ChatMessageAssessment(BaseModel):
-    model_config = model_config
-
-    id: str
-    object: str
-    message_id: str
-    status: ChatMessageAssessmentStatus
-    type: ChatMessageAssessmentType
-    title: str | None = None
-    explanation: str | None = None
-    label: ChatMessageAssessmentLabel | None = None
-    is_visible: bool
-    created_at: datetime | None = None
-    updated_at: datetime | None = None
-
-
 class MessageLogStatus(StrEnum):
     RUNNING = "RUNNING"
     COMPLETED = "COMPLETED"

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/content/schemas.py

@@ -117,7 +117,7 @@ class ContentReference(BaseModel):
     url: str
     original_index: list[int] = Field(
         default=[],
-        description="List of indices in the ChatMessage original_content this reference refers to. This is usually the id in the functionCallResponse. List type due to implementation in node-chat",
+        description="List of indices in the ChatMessage original_text this reference refers to. This is usually the id in the functionCallResponse. List type due to implementation in node-chat",
     )
 
     @classmethod

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/language_model/__init__.py

@@ -20,6 +20,7 @@ from .schemas import (
     LanguageModelStreamResponseMessage,
     LanguageModelSystemMessage,
     LanguageModelTokenLimits,
+    LanguageModelTokenUsage,
     LanguageModelTool,
     LanguageModelToolDescription,
     LanguageModelToolMessage,
@@ -57,6 +58,7 @@ __all__ = [
     "LanguageModelStreamResponseMessage",
     "LanguageModelSystemMessage",
     "LanguageModelTokenLimits",
+    "LanguageModelTokenUsage",
     "LanguageModelTool",
     "LanguageModelToolDescription",
     "LanguageModelToolMessage",

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/language_model/functions.py

@@ -30,7 +30,6 @@ from unique_toolkit.language_model.schemas import (
     LanguageModelMessages,
     LanguageModelResponse,
     LanguageModelStreamResponse,
-    LanguageModelStreamResponseMessage,
     LanguageModelTool,
     LanguageModelToolDescription,
 )
@@ -528,12 +527,13 @@ def _create_language_model_stream_response_with_references(
         search_context=content_chunks,
     )
 
-    stream_response_message = LanguageModelStreamResponseMessage(
+    stream_response_message = ChatMessage(
         id="stream_unknown",
         previous_message_id=None,
         role=LanguageModelMessageRole.ASSISTANT,
         text=message.content or "",
         original_text=content,
+        chat_id="stream_unknown",
         references=[
             ContentReference(**u.model_dump()) for u in message.references or []
         ],
@@ -657,8 +657,9 @@ async def stream_complete_with_references_openai(
             ]
 
         return LanguageModelStreamResponse(
-            message=LanguageModelStreamResponseMessage(
+            message=ChatMessage(
                 id=message_id or uuid4().hex,
+                chat_id="",
                 previous_message_id=None,
                 role=LanguageModelMessageRole.ASSISTANT,
                 text=text_,

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/language_model/schemas.py

@@ -1,6 +1,5 @@
 import json
 import math
-from enum import StrEnum
 from typing import Any, Literal, Self, TypeVar
 from uuid import uuid4
 
@@ -40,7 +39,8 @@ from pydantic import (
 )
 from typing_extensions import deprecated, overload
 
-from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.chat.schemas import ChatMessage
+from unique_toolkit.chat.schemas import ChatMessageRole as LanguageModelMessageRole
 from unique_toolkit.language_model._responses_api_utils import (
     convert_user_message_content_to_responses_api,
 )
@@ -54,33 +54,14 @@ model_config = ConfigDict(
 )
 
 
-# Equivalent to
-# from openai.types.chat.chat_completion_role import ChatCompletionRole
-class LanguageModelMessageRole(StrEnum):
-    USER = "user"
-    SYSTEM = "system"
-    ASSISTANT = "assistant"
-    TOOL = "tool"
-
-
-# This is tailored to the unique backend
-class LanguageModelStreamResponseMessage(BaseModel):
-    model_config = model_config
-
-    id: str
-    previous_message_id: (
-        str | None
-    )  # Stream response can return a null previous_message_id if an assisstant message is manually added
-    role: LanguageModelMessageRole
-    text: str
-    original_text: str | None = None
-    references: list[ContentReference] = []
-
-    # TODO make sdk return role in lowercase
-    # Currently needed as sdk returns role in uppercase
-    @field_validator("role", mode="before")
-    def set_role(cls, value: str):
-        return value.lower()
+# Backward compatibility alias — use ChatMessage directly.
+# LanguageModelStreamResponseMessage was a toolkit-only subclass with stricter types
+# (required id/content, non-null references). The backend DTOs (PublicMessageDto,
+# PublicStreamResultDto) use a single message type, so this distinction was unnecessary.
+# Note: id and content are now optional (str | None) instead of required (str). This is a
+# deliberate contract relaxation — the backend always returns non-null values in stream
+# responses, but callers that relied on the stricter types should add explicit null checks.
+LanguageModelStreamResponseMessage = ChatMessage
 
 
 class LanguageModelFunction(BaseModel):
@@ -163,12 +144,21 @@ class LanguageModelFunction(BaseModel):
             )
 
 
+class LanguageModelTokenUsage(BaseModel):
+    model_config = model_config
+
+    completion_tokens: int | None = None
+    prompt_tokens: int | None = None
+    total_tokens: int | None = None
+
+
 class LanguageModelStreamResponse(BaseModel):
     model_config = model_config
 
-    message: LanguageModelStreamResponseMessage
+    message: ChatMessage
     tool_calls: list[LanguageModelFunction] | None = None
     stopped_by_user: bool = False
+    usage: LanguageModelTokenUsage | None = None
 
     def is_empty(self) -> bool:
         """

{unique_toolkit-1.57 → unique_toolkit-1.58}/unique_toolkit/services/chat_service.py

@@ -140,6 +140,50 @@ class ChatService(ChatServiceDeprecated):
 
         return self._elicitation_service
 
+    def get_debug_info(self) -> dict[str, Any]:
+        """Retrieves the debug information from the current user message.
+
+        Returns:
+            dict[str, Any]: The debug information stored on the user message,
+                or an empty dict if retrieval fails or no debug info exists.
+
+        """
+        try:
+            raw_msg = unique_sdk.Message.retrieve(
+                user_id=self._user_id,
+                company_id=self._company_id,
+                id=self._user_message_id,
+                chatId=self._chat_id,
+            )
+            return getattr(raw_msg, "debugInfo", {}) or {}
+        except Exception:
+            logging.getLogger(__name__).warning(
+                "Failed to retrieve debug info from user message", exc_info=True
+            )
+            return {}
+
+    async def get_debug_info_async(self) -> dict[str, Any]:
+        """Retrieves the debug information from the current user message asynchronously.
+
+        Returns:
+            dict[str, Any]: The debug information stored on the user message,
+                or an empty dict if retrieval fails or no debug info exists.
+
+        """
+        try:
+            raw_msg = await unique_sdk.Message.retrieve_async(
+                user_id=self._user_id,
+                company_id=self._company_id,
+                id=self._user_message_id,
+                chatId=self._chat_id,
+            )
+            return getattr(raw_msg, "debugInfo", {}) or {}
+        except Exception:
+            logging.getLogger(__name__).warning(
+                "Failed to retrieve debug info from user message", exc_info=True
+            )
+            return {}
+
     async def update_debug_info_async(self, debug_info: dict):
         """Updates the debug information for the chat session.