unique_toolkit 2026.36.0.dev1__py3-none-any.whl → 2026.36.0.dev2__py3-none-any.whl

unique_toolkit/language_model/functions.py

@@ -46,12 +46,15 @@ logger = logging.getLogger(f"toolkit.language_model.{__name__}")
 def _completion_headers(
     chat_id: str | None,
     assistant_id: str | None,
+    assistant_message_id: str | None = None,
 ) -> dict[str, str] | None:
     headers: dict[str, str] = {}
     if chat_id:
         headers["x-chat-id"] = chat_id
     if assistant_id:
         headers["x-assistant-id"] = assistant_id
+    if assistant_message_id:
+        headers["x-assistant-message-id"] = assistant_message_id
     return headers or None
 
 
@@ -68,6 +71,7 @@ def complete(
     user_id: str | None = None,
     chat_id: str | None = None,
     assistant_id: str | None = None,
+    assistant_message_id: str | None = None,
 ) -> LanguageModelResponse:
     """Call the completion endpoint synchronously without streaming the response.
 
@@ -103,7 +107,7 @@ def complete(
         )
 
     try:
-        headers = _completion_headers(chat_id, assistant_id)
+        headers = _completion_headers(chat_id, assistant_id, assistant_message_id)
         if headers is None:
             response = unique_sdk.ChatCompletion.create(
                 company_id=company_id,
@@ -142,6 +146,7 @@ async def complete_async(
     structured_output_enforce_schema: bool = False,
     chat_id: str | None = None,
     assistant_id: str | None = None,
+    assistant_message_id: str | None = None,
 ) -> LanguageModelResponse:
     """Call the completion endpoint asynchronously without streaming the response.
 
@@ -186,7 +191,7 @@ async def complete_async(
     )
 
     try:
-        headers = _completion_headers(chat_id, assistant_id)
+        headers = _completion_headers(chat_id, assistant_id, assistant_message_id)
         if headers is None:
             response = await unique_sdk.ChatCompletion.create_async(
                 company_id=company_id,
@@ -502,6 +507,7 @@ def complete_with_references(
     other_options: dict[str, Any] | None = None,
     chat_id: str | None = None,
     assistant_id: str | None = None,
+    assistant_message_id: str | None = None,
 ) -> LanguageModelStreamResponse:
     # Use toolkit language model functions for chat completion
 
@@ -516,6 +522,7 @@ def complete_with_references(
         other_options=other_options,
         chat_id=chat_id,
         assistant_id=assistant_id,
+        assistant_message_id=assistant_message_id,
     )
 
     return _create_language_model_stream_response_with_references(
@@ -539,6 +546,7 @@ async def complete_with_references_async(
     other_options: dict[str, Any] | None = None,
     chat_id: str | None = None,
     assistant_id: str | None = None,
+    assistant_message_id: str | None = None,
 ) -> LanguageModelStreamResponse:
     # Use toolkit language model functions for chat completion
     response = await complete_async(
@@ -552,6 +560,7 @@ async def complete_with_references_async(
         other_options=other_options,
         chat_id=chat_id,
         assistant_id=assistant_id,
+        assistant_message_id=assistant_message_id,
     )
 
     return _create_language_model_stream_response_with_references(

unique_toolkit/language_model/service.py

@@ -5,7 +5,11 @@ from pydantic import BaseModel
 from typing_extensions import deprecated
 
 from unique_toolkit._common.validate_required_values import validate_required_values
-from unique_toolkit.app.schemas import AssistantWebhookEvent, BaseEvent
+from unique_toolkit.app.schemas import (
+    AssistantWebhookEvent,
+    BaseEvent,
+    ChatEventPayload,
+)
 from unique_toolkit.app.unique_settings import UniqueSettings
 from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.language_model.constants import (
@@ -67,10 +71,13 @@ class LanguageModelService:
             self._event = event
             self._chat_id = None
             self._assistant_id = None
+            self._assistant_message_id = None
             if isinstance(event, AssistantWebhookEvent):
                 payload = event.payload
                 self._chat_id = payload.chat_id
                 self._assistant_id = payload.assistant_id
+                if isinstance(payload, ChatEventPayload):
+                    self._assistant_message_id = payload.assistant_message.id
             return
 
         [company_id, user_id] = validate_required_values([company_id, user_id])
@@ -79,6 +86,7 @@ class LanguageModelService:
         self._user_id: str = user_id
         self._chat_id: str | None = None
         self._assistant_id: str | None = None
+        self._assistant_message_id: str | None = None
 
     @classmethod
     def from_event(cls, event: BaseEvent[Any]):
@@ -89,6 +97,8 @@ class LanguageModelService:
         if isinstance(event, AssistantWebhookEvent):
             service._chat_id = event.payload.chat_id
             service._assistant_id = event.payload.assistant_id
+            if isinstance(event.payload, ChatEventPayload):
+                service._assistant_message_id = event.payload.assistant_message.id
         return service
 
     @classmethod
@@ -251,6 +261,7 @@ class LanguageModelService:
             structured_output_enforce_schema=structured_output_enforce_schema,
             chat_id=self._chat_id,
             assistant_id=self._assistant_id,
+            assistant_message_id=self._assistant_message_id,
         )
 
     async def complete_async(
@@ -281,6 +292,7 @@ class LanguageModelService:
             structured_output_enforce_schema=structured_output_enforce_schema,
             chat_id=self._chat_id,
             assistant_id=self._assistant_id,
+            assistant_message_id=self._assistant_message_id,
         )
 
     @classmethod
@@ -344,6 +356,7 @@ class LanguageModelService:
             start_text=start_text,
             chat_id=self._chat_id,
             assistant_id=self._assistant_id,
+            assistant_message_id=self._assistant_message_id,
         )
 
     async def complete_with_references_async(
@@ -371,4 +384,5 @@ class LanguageModelService:
             start_text=start_text,
             chat_id=self._chat_id,
             assistant_id=self._assistant_id,
+            assistant_message_id=self._assistant_message_id,
         )

{unique_toolkit-2026.36.0.dev1.dist-info → unique_toolkit-2026.36.0.dev2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 2026.36.0.dev1
+Version: 2026.36.0.dev2
 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-2026.36.0.dev1.dist-info → unique_toolkit-2026.36.0.dev2.dist-info}/RECORD

@@ -391,14 +391,14 @@ unique_toolkit/language_model/_responses_api_utils.py,sha256=w2ZyleOQcWLpOWGJNNb
 unique_toolkit/language_model/builder.py,sha256=KwBzsD6-fvQnB4p2I7M-5Tk6j4tmQ973MYwU1ulXFb4,4526
 unique_toolkit/language_model/constants.py,sha256=B-topqW0r83dkC_25DeQfnPk3n53qzIHUCBS7YJ0-1U,119
 unique_toolkit/language_model/default_language_model.py,sha256=mcyuFs_V6tp52C8xJSHiEHO4kFGj6xed1gvozvCRD7k,1181
-unique_toolkit/language_model/functions.py,sha256=1J-7p9fvVbU1agrVZu-_UGiQHDSBEX0_eZ4SidntdFM,26226
+unique_toolkit/language_model/functions.py,sha256=W6JC4LUh-LZ9waDgph--xVmKeQ7Ls6NRAqDx9DwQdeY,26691
 unique_toolkit/language_model/infos.py,sha256=zs0UFhds4UkSRr3zxyTJ-jJ9u2O2x9VgZx3a9d4zbYU,146221
 unique_toolkit/language_model/invocation_stats.py,sha256=XAM_qckI-BoMBSgDpVcxFD1cFSdi1-GHZToHz_jLynM,2491
 unique_toolkit/language_model/model_costs.py,sha256=WM1u2fWAKpPB_creS-ew8YAzc1V8oiaCbvhpSJaO8B0,3789
 unique_toolkit/language_model/prompt.py,sha256=JSawaLjQg3VR-E2fK8engFyJnNdk21zaO8pPIodzN4Q,3991
 unique_toolkit/language_model/reference.py,sha256=T3wr2jP_uRg5wKVIbKqLbxJlcgxfbZQGq3XdjbQ2kSo,6984
 unique_toolkit/language_model/schemas.py,sha256=IUmp_JMHx0czadJlmsLL9nmc-nCPmVRVLp2T5DSYsqY,35402
-unique_toolkit/language_model/service.py,sha256=9Q1xQCGjGVHZVMjkXaAPKAc7fLY2lGY6sFgdasoRQZE,12616
+unique_toolkit/language_model/service.py,sha256=lgpzWhqxLBOeFacCafQi4e7BqAxJbOlr8GnCYcKeTdM,13274
 unique_toolkit/language_model/settings.py,sha256=7szohExI-qmTfUf2Zp5AB5TzrsZy1UekAGyXx_UaPMc,1524
 unique_toolkit/language_model/stream_transform.py,sha256=iNNGpLY_7G6Bpp9Y4KtbIkQm3WB4yZt75BXcqx0DVkI,2775
 unique_toolkit/language_model/utils.py,sha256=bPQ4l6_YO71w-zaIPanUUmtbXC1_hCvLK0tAFc3VCRc,1902
@@ -420,6 +420,6 @@ unique_toolkit/short_term_memory/service.py,sha256=2JMBEBn6GoyK-a0xvdcWpPft-3C74
 unique_toolkit/smart_rules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unique_toolkit/smart_rules/compile.py,sha256=vRe0aMJWuVIK-VsqomgxQImzke7MDmwgSxxJfYCa7zw,1386
 unique_toolkit/test_utilities/events.py,sha256=FLrhMh0A_mQutAUjoD0FdMcVQ4Gq31r35OVGzRjvLro,6760
-unique_toolkit-2026.36.0.dev1.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
-unique_toolkit-2026.36.0.dev1.dist-info/METADATA,sha256=TRfcAs3vlbrpma_tG96J6ZZq6eAu9L9NWGIrq2Y2eVY,7821
-unique_toolkit-2026.36.0.dev1.dist-info/RECORD,,
+unique_toolkit-2026.36.0.dev2.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
+unique_toolkit-2026.36.0.dev2.dist-info/METADATA,sha256=Caxi8wvrymP225m-SxkPfSMbMQ4agwTI-81eGleVGXw,7821
+unique_toolkit-2026.36.0.dev2.dist-info/RECORD,,