unique_toolkit 1.68.13__tar.gz → 1.69.1__tar.gz

{unique_toolkit-1.68.13 → unique_toolkit-1.69.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 1.68.13
+Version: 1.69.1
 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>
@@ -38,9 +38,11 @@ Requires-Dist: openai>=1.105.0,<2 ; extra == 'langchain'
 Requires-Dist: langchain>=1.0.0,<2 ; extra == 'langchain'
 Requires-Dist: langchain-openai>=0.3.28 ; extra == 'langchain'
 Requires-Dist: langchain-core>=1.0.0,<2 ; extra == 'langchain'
+Requires-Dist: prometheus-client>=0.24.1,<1 ; extra == 'monitoring'
 Requires-Python: >=3.12
 Provides-Extra: fastapi
 Provides-Extra: langchain
+Provides-Extra: monitoring
 Description-Content-Type: text/markdown
 
 # Unique Toolkit

{unique_toolkit-1.68.13 → unique_toolkit-1.69.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "1.68.13"
+version = "1.69.1"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12"
@@ -57,6 +57,9 @@ fastapi = [
     "starlette>=0.49.1",
     "uvicorn>=0.37.0,<0.38",
 ]
+monitoring = [
+    "prometheus_client>=0.24.1,<1",
+]
 
 [build-system]
 requires = ["uv_build>=0.8.14,<0.9.0"]

{unique_toolkit-1.68.13 → unique_toolkit-1.69.1}/unique_toolkit/agentic/history_manager/history_construction_with_contents.py

@@ -6,7 +6,7 @@ import mimetypes
 from datetime import datetime
 from enum import StrEnum
 from itertools import groupby
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Iterator
 
 if TYPE_CHECKING:
     from unique_toolkit.content.schemas import ContentChunk
@@ -77,7 +77,7 @@ class ChatHistoryWithContent(RootModel):
 
         return cls(root=grouped_elements)
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[ChatMessageWithContents]:
         return iter(self.root)
 
     def __getitem__(self, item):
@@ -94,7 +94,7 @@ def get_chat_history_with_contents(
         id=user_message.id,
         chat_id=chat_id,
         text=user_message.text,
-        originalText=user_message.original_text,
+        original_text=user_message.original_text,
         role=ChatRole.USER,
         gpt_request=None,
         created_at=datetime.fromisoformat(user_message.created_at),
@@ -118,6 +118,41 @@ def get_chat_history_with_contents(
     )
 
 
+async def get_chat_history_with_contents_async(
+    *,
+    user_message: ChatEventUserMessage,
+    chat_id: str,
+    chat_history: list[ChatMessage],
+    content_service: ContentService,
+) -> ChatHistoryWithContent:
+    last_user_message = ChatMessage(
+        id=user_message.id,
+        chat_id=chat_id,
+        text=user_message.text,
+        original_text=user_message.original_text,
+        role=ChatRole.USER,
+        gpt_request=None,
+        created_at=datetime.fromisoformat(user_message.created_at),
+    )
+    if len(chat_history) > 0 and last_user_message.id == chat_history[-1].id:
+        pass
+    else:
+        chat_history.append(last_user_message)
+
+    chat_contents = await content_service.search_contents_async(
+        where={
+            "ownerId": {
+                "equals": chat_id,
+            },
+        },
+    )
+
+    return ChatHistoryWithContent.from_chat_history_and_contents(
+        chat_history,
+        chat_contents,
+    )
+
+
 def download_encoded_images(
     contents: list[Content],
     content_service: ContentService,
@@ -141,6 +176,30 @@ def download_encoded_images(
     return base64_encoded_images
 
 
+async def download_encoded_images_async(
+    *,
+    contents: list[Content],
+    content_service: ContentService,
+    chat_id: str,
+) -> list[str]:
+    base64_encoded_images = []
+    for im in contents:
+        if FileUtils.is_image_content(im.key):
+            try:
+                file_bytes = await content_service.download_content_to_bytes_async(
+                    content_id=im.id,
+                    chat_id=chat_id,
+                )
+
+                mime_type, _ = mimetypes.guess_type(im.key)
+                encoded_string = base64.b64encode(file_bytes).decode("utf-8")
+                image_string = f"data:{mime_type};base64," + encoded_string
+                base64_encoded_images.append(image_string)
+            except Exception as e:
+                _LOGGER.warning("Failed to download image %s: %s", im.key, e)
+    return base64_encoded_images
+
+
 class FileContentSerialization(StrEnum):
     NONE = "none"
     FILE_NAME = "file_name"
@@ -212,6 +271,49 @@ def _append_element_to_builder(
         )
 
 
+async def _append_element_to_builder_async(
+    builder: MessagesBuilder,
+    *,
+    c: ChatMessageWithContents,
+    text: str,
+    include_images: ImageContentInclusion,
+    file_content_serialization_type: FileContentSerialization,
+    content_service: ContentService,
+    chat_id: str,
+) -> None:
+    if len(c.contents) > 0:
+        file_contents = [co for co in c.contents if FileUtils.is_file_content(co.key)]
+        image_contents = [co for co in c.contents if FileUtils.is_image_content(co.key)]
+        content = (
+            text
+            + "\n\n"
+            + file_content_serialization(
+                file_contents,
+                file_content_serialization_type,
+            )
+        ).strip()
+        if include_images and len(image_contents) > 0:
+            builder.image_message_append(
+                content=content,
+                images=await download_encoded_images_async(
+                    contents=image_contents,
+                    content_service=content_service,
+                    chat_id=chat_id,
+                ),
+                role=map_chat_llm_message_role[c.role],
+            )
+        else:
+            builder.message_append(
+                role=map_chat_llm_message_role[c.role],
+                content=content,
+            )
+    else:
+        builder.message_append(
+            role=map_chat_llm_message_role[c.role],
+            content=text,
+        )
+
+
 def get_full_history_with_contents(
     user_message: ChatEventUserMessage,
     chat_id: str,
@@ -228,7 +330,7 @@ def get_full_history_with_contents(
     )
 
     builder = LanguageModelMessages([]).builder()
-    for c in grouped_elements:
+    for c in grouped_elements.root:
         # LanguageModelUserMessage has no field original_text
         text = c.original_text if c.original_text else c.content
         if text is None:
@@ -249,6 +351,43 @@ def get_full_history_with_contents(
     return builder.build()
 
 
+async def get_full_history_with_contents_async(
+    *,
+    user_message: ChatEventUserMessage,
+    chat_id: str,
+    chat_service: ChatService,
+    content_service: ContentService,
+    include_images: ImageContentInclusion = ImageContentInclusion.ALL,
+    file_content_serialization_type: FileContentSerialization = FileContentSerialization.FILE_NAME,
+) -> LanguageModelMessages:
+    grouped_elements = await get_chat_history_with_contents_async(
+        user_message=user_message,
+        chat_id=chat_id,
+        chat_history=await chat_service.get_full_history_async(),
+        content_service=content_service,
+    )
+
+    builder = LanguageModelMessages([]).builder()
+    for c in grouped_elements.root:
+        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_text of LanguageModelMessages should exist.",
+                )
+            text = ""
+        await _append_element_to_builder_async(
+            builder=builder,
+            c=c,
+            text=text,
+            include_images=include_images,
+            file_content_serialization_type=file_content_serialization_type,
+            content_service=content_service,
+            chat_id=chat_id,
+        )
+    return builder.build()
+
+
 def get_full_history_with_contents_and_tool_calls(
     *,
     user_message: ChatEventUserMessage,
@@ -307,7 +446,7 @@ def get_full_history_with_contents_and_tool_calls(
     )
 
     builder = LanguageModelMessages([]).builder()
-    for c in grouped_elements:
+    for c in grouped_elements.root:
         text = c.original_text if c.original_text else c.content
         if text is None:
             if c.role == ChatRole.USER:
@@ -391,6 +530,122 @@ def get_full_history_with_contents_and_tool_calls(
     return builder.build(), max_source_number, source_map
 
 
+async def get_full_history_with_contents_and_tool_calls_async(
+    *,
+    user_message: ChatEventUserMessage,
+    chat_id: str,
+    chat_service: ChatService,
+    content_service: ContentService,
+    include_images: ImageContentInclusion = ImageContentInclusion.ALL,
+    file_content_serialization_type: FileContentSerialization = FileContentSerialization.FILE_NAME,
+) -> tuple[LanguageModelMessages, int, dict[int, "ContentChunk"]]:
+    """Async version of get_full_history_with_contents_and_tool_calls."""
+    from unique_toolkit.agentic.history_manager.utils import (
+        build_source_map_from_tool_calls,
+        compute_max_source_number_from_tool_calls,
+    )
+    from unique_toolkit.content.schemas import (
+        ContentChunk as ContentChunk,  # noqa: F811
+    )
+
+    chat_history = await chat_service.get_full_history_async()
+
+    assistant_message_ids = [
+        msg.id for msg in chat_history if msg.role == ChatRole.ASSISTANT and msg.id
+    ]
+    all_tool_calls: list[ChatMessageTool] = []
+    tool_calls_by_message: dict[str, list[ChatMessageTool]] = {}
+    if assistant_message_ids:
+        try:
+            all_tool_calls = await chat_service.get_message_tools_async(
+                message_ids=assistant_message_ids,
+            )
+            for tc in all_tool_calls:
+                if tc.message_id:
+                    tool_calls_by_message.setdefault(tc.message_id, []).append(tc)
+        except Exception:
+            _LOGGER.warning(
+                "Failed to batch-load tool calls, falling back to empty", exc_info=True
+            )
+    max_source_number = compute_max_source_number_from_tool_calls(all_tool_calls)
+    source_map: dict[int, ContentChunk] = build_source_map_from_tool_calls(
+        all_tool_calls
+    )
+
+    grouped_elements = await get_chat_history_with_contents_async(
+        user_message=user_message,
+        chat_id=chat_id,
+        chat_history=chat_history,
+        content_service=content_service,
+    )
+
+    builder = LanguageModelMessages([]).builder()
+    for c in grouped_elements.root:
+        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_text of LanguageModelMessages should exist.",
+                )
+            text = ""
+
+        if c.role == ChatRole.ASSISTANT and c.id and c.id in tool_calls_by_message:
+            tc_records = sorted(
+                tool_calls_by_message[c.id],
+                key=lambda tc: (tc.round_index, tc.sequence_index),
+            )
+            if tc_records:
+                for _round, round_group in groupby(
+                    tc_records, key=lambda tc: tc.round_index
+                ):
+                    round_tcs = list(round_group)
+                    round_tcs_with_response = [
+                        tc
+                        for tc in round_tcs
+                        if tc.response and tc.response.content is not None
+                    ]
+                    if not round_tcs_with_response:
+                        continue
+                    fns = [
+                        LanguageModelFunction(
+                            id=tc.external_tool_call_id,
+                            name=tc.function_name,
+                            arguments=tc.arguments,
+                        )
+                        for tc in round_tcs_with_response
+                    ]
+                    builder.messages.append(
+                        LanguageModelAssistantMessage.from_functions(tool_calls=fns)
+                    )
+                    for fn, tc in zip(fns, round_tcs_with_response):
+                        builder.messages.append(
+                            LanguageModelToolMessage(
+                                tool_call_id=fn.id,
+                                content=tc.response.content,  # type: ignore[union-attr]
+                                name=tc.function_name,
+                            )
+                        )
+
+        had_tool_calls = (
+            c.role == ChatRole.ASSISTANT
+            and c.id is not None
+            and c.id in tool_calls_by_message
+        )
+        if had_tool_calls and not text:
+            continue
+
+        await _append_element_to_builder_async(
+            builder=builder,
+            c=c,
+            text=text,
+            include_images=include_images,
+            file_content_serialization_type=file_content_serialization_type,
+            content_service=content_service,
+            chat_id=chat_id,
+        )
+    return builder.build(), max_source_number, source_map
+
+
 def get_full_history_as_llm_messages(
     chat_service: ChatService,
 ) -> LanguageModelMessages:

{unique_toolkit-1.68.13 → unique_toolkit-1.69.1}/unique_toolkit/agentic/history_manager/loop_token_reducer.py

@@ -10,8 +10,8 @@ from unique_toolkit._common.token.token_counting import (
 from unique_toolkit._common.validators import LMI
 from unique_toolkit.agentic.history_manager.history_construction_with_contents import (
     FileContentSerialization,
-    get_full_history_with_contents,
-    get_full_history_with_contents_and_tool_calls,
+    get_full_history_with_contents_and_tool_calls_async,
+    get_full_history_with_contents_async,
 )
 from unique_toolkit.agentic.history_manager.utils import serialize_tool_content_json
 from unique_toolkit.agentic.reference_manager.reference_manager import ReferenceManager
@@ -312,19 +312,21 @@ class LoopTokenReducer:
             else FileContentSerialization.FILE_NAME
         )
         if self._enable_tool_call_persistence:
-            full_history, max_src, src_map = (
-                get_full_history_with_contents_and_tool_calls(
-                    user_message=self._user_message,
-                    chat_id=self._chat_id,
-                    chat_service=self._chat_service,
-                    content_service=self._content_service,
-                    file_content_serialization_type=file_content_serialization_type,
-                )
+            (
+                full_history,
+                max_src,
+                src_map,
+            ) = await get_full_history_with_contents_and_tool_calls_async(
+                user_message=self._user_message,
+                chat_id=self._chat_id,
+                chat_service=self._chat_service,
+                content_service=self._content_service,
+                file_content_serialization_type=file_content_serialization_type,
             )
             self._max_db_source_number = max_src
             self._db_source_map = src_map
         else:
-            full_history = get_full_history_with_contents(
+            full_history = await get_full_history_with_contents_async(
                 user_message=self._user_message,
                 chat_id=self._chat_id,
                 chat_service=self._chat_service,

{unique_toolkit-1.68.13 → unique_toolkit-1.69.1}/unique_toolkit/agentic/loop_runner/runners/qwen/qwen_runner.py

@@ -65,7 +65,7 @@ class QwenLoopIterationRunner(BasicLoopIterationRunner):
             loop_runner_kwargs=kwargs,
             prepare_loop_runner_kwargs=_prepare,
         )
-        return self._process_response(response)
+        return await self._process_response(response)
 
     async def _handle_last_iteration(
         self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
@@ -78,22 +78,23 @@ class QwenLoopIterationRunner(BasicLoopIterationRunner):
             last_iteration_instruction=self._last_iteration_instruction,
         )
         response = await super()._handle_last_iteration(**kwargs)
-        return self._process_response(response)
+        return await self._process_response(response)
 
     async def _handle_normal_iteration(
         self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
     ) -> LanguageModelStreamResponse:
         response = await super()._handle_normal_iteration(**kwargs)
-        return self._process_response(response)
+        return await self._process_response(response)
 
-    def _process_response(
+    async def _process_response(
         self, response: LanguageModelStreamResponse
     ) -> LanguageModelStreamResponse:
-        if "</tool_call>" == response.message.text.strip():
+        message_text = response.message.text or ""
+        if "</tool_call>" == message_text.strip():
             _LOGGER.warning(
                 "Response contains only <tool_call>. This is not allowed. Returning empty response."
             )
-            self._chat_service.modify_assistant_message(content="")
+            await self._chat_service.modify_assistant_message_async(content="")
             response.message.text = ""
 
         return response

{unique_toolkit-1.68.13 → unique_toolkit-1.69.1}/unique_toolkit/agentic/tools/mcp/tool_wrapper.py

@@ -377,7 +377,10 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
     async def _call_mcp_tool_via_sdk(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
         """Call MCP tool via SDK to public API"""
         try:
-            result = unique_sdk.MCP.call_tool(
+            self.logger.info(
+                f"Calling MCP tool {self.name} with arguments: {arguments}"
+            )
+            result = await unique_sdk.MCP.call_tool_async(
                 user_id=self._event.user_id,
                 company_id=self._event.company_id,
                 name=self.name,
@@ -386,9 +389,6 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
                 arguments=arguments,
             )
 
-            self.logger.info(
-                f"Calling MCP tool {self.name} with arguments: {arguments}"
-            )
             self.logger.debug(f"Result: {result}")
 
             return result

{unique_toolkit-1.68.13 → unique_toolkit-1.69.1}/unique_toolkit/app/fast_api_factory.py

@@ -63,6 +63,21 @@ def build_unique_custom_app(
         """Health check endpoint."""
         return JSONResponse(content={"status": "healthy", "service": title})
 
+    from unique_toolkit.monitoring import _MONITORING_AVAILABLE
+
+    if _MONITORING_AVAILABLE:
+        from starlette.responses import Response as StarletteResponse
+
+        from unique_toolkit.monitoring import get_metrics
+
+        @app.get(path="/metrics")
+        async def metrics_endpoint() -> StarletteResponse:
+            """Prometheus metrics endpoint."""
+            return StarletteResponse(
+                content=get_metrics(),
+                media_type="text/plain; version=0.0.4; charset=utf-8",
+            )
+
     @app.post(path=webhook_path)
     async def webhook_handler(
         request: Request, background_tasks: BackgroundTasks

unique_toolkit-1.69.1/unique_toolkit/monitoring/__init__.py

@@ -0,0 +1,39 @@
+try:
+    from prometheus_client import CollectorRegistry  # noqa: F401
+
+    _MONITORING_AVAILABLE = True
+except ImportError:
+    _MONITORING_AVAILABLE = False
+
+
+def _check_monitoring_available() -> None:
+    if not _MONITORING_AVAILABLE:
+        raise ImportError(
+            "prometheus_client is not installed. "
+            "Install it with: pip install unique_toolkit[monitoring]"
+        )
+
+
+if _MONITORING_AVAILABLE:
+    from unique_toolkit.monitoring.middleware import MetricsMiddleware
+    from unique_toolkit.monitoring.registry import (
+        REGISTRY,
+        MetricNamespace,
+        get_metrics,
+        metric_scope,
+        track_execution,
+    )
+
+    __all__ = [
+        "REGISTRY",
+        "MetricNamespace",
+        "MetricsMiddleware",
+        "get_metrics",
+        "metric_scope",
+        "track_execution",
+        "_MONITORING_AVAILABLE",
+    ]
+else:
+
+    def __getattr__(name: str):
+        _check_monitoring_available()

unique_toolkit-1.69.1/unique_toolkit/monitoring/middleware.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import time
+
+from prometheus_client import Counter, Gauge, Histogram
+
+from .registry import REGISTRY
+
+_http_requests_total = Counter(
+    "http_requests_total",
+    "Total HTTP requests",
+    ["method", "path", "status_code"],
+    registry=REGISTRY,
+)
+_http_request_duration_seconds = Histogram(
+    "http_request_duration_seconds",
+    "HTTP request duration",
+    ["method", "path"],
+    registry=REGISTRY,
+)
+_http_requests_in_progress = Gauge(
+    "http_requests_in_progress",
+    "In-flight HTTP requests",
+    ["method"],
+    registry=REGISTRY,
+)
+
+_DEFAULT_EXCLUDED_PATHS = frozenset({"/health", "/metrics", "/"})
+
+
+class MetricsMiddleware:
+    """ASGI middleware that tracks HTTP request metrics.
+
+    Works with any ASGI app (FastAPI, Quart, Starlette).
+
+    Handles:
+    - Non-HTTP scopes (WebSocket, lifespan) — passed through, not tracked
+    - Excluded paths (health, metrics) — passed through, not tracked
+    - App exceptions — duration + error still recorded, in-progress gauge cleaned up
+    - Streaming responses — status captured from first http.response.start message
+
+    Usage::
+
+        # FastAPI
+        app.add_middleware(MetricsMiddleware)
+
+        # Quart
+        app.asgi_app = MetricsMiddleware(app.asgi_app)
+    """
+
+    def __init__(self, app, excluded_paths: set[str] | None = None):
+        self.app = app
+        self._excluded_paths = (
+            excluded_paths if excluded_paths is not None else _DEFAULT_EXCLUDED_PATHS
+        )
+
+    async def __call__(self, scope, receive, send):
+        # Pass through non-HTTP (WebSocket, lifespan)
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        path = scope.get("path", "/")
+
+        # Pass through excluded paths (health checks, metrics endpoint)
+        if path in self._excluded_paths:
+            return await self.app(scope, receive, send)
+
+        method = scope.get("method", "UNKNOWN")
+        status_code = "500"  # Default if app crashes before sending response
+
+        _http_requests_in_progress.labels(method=method).inc()
+        start = time.perf_counter()
+
+        async def send_wrapper(message):
+            nonlocal status_code
+            if message["type"] == "http.response.start":
+                status_code = str(message["status"])
+            await send(message)
+
+        try:
+            await self.app(scope, receive, send_wrapper)
+        finally:
+            duration = time.perf_counter() - start
+            _http_requests_in_progress.labels(method=method).dec()
+            _http_request_duration_seconds.labels(method=method, path=path).observe(
+                duration
+            )
+            _http_requests_total.labels(
+                method=method, path=path, status_code=status_code
+            ).inc()