unique_toolkit 0.5.49__py3-none-any.whl → 0.5.50__py3-none-any.whl

unique_toolkit/language_model/service.py

@@ -1,7 +1,8 @@
 import logging
-from typing import Optional, cast
+from typing import Optional, Type, cast
 
 import unique_sdk
+from pydantic import BaseModel
 
 from unique_toolkit._common._base_service import BaseService
 from unique_toolkit.app.schemas import Event
@@ -37,6 +38,8 @@ class LanguageModelService(BaseService):
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: Optional[list[LanguageModelTool]] = None,
+        structured_output_model: Optional[Type[BaseModel]] = None,
+        structured_output_enforce_schema: bool = False,
         other_options: Optional[dict] = None,
     ):
         """
@@ -53,12 +56,14 @@ class LanguageModelService(BaseService):
         Returns:
             LanguageModelResponse: The LanguageModelResponse object.
         """
-        options, model, messages_dict, _ = self.prepare_completion_params_util(
+        options, model, messages_dict, _ = self._prepare_completion_params_util(
             messages=messages,
             model_name=model_name,
             temperature=temperature,
             tools=tools,
             other_options=other_options,
+            structured_output_model=structured_output_model,
+            structured_output_enforce_schema=structured_output_enforce_schema,
         )
 
         try:
@@ -77,9 +82,8 @@ class LanguageModelService(BaseService):
             self.logger.error(f"Error completing: {e}")
             raise e
 
-    @classmethod
     async def complete_async_util(
-        cls,
+        self,
         company_id: str,
         messages: LanguageModelMessages,
         model_name: LanguageModelName | str,
@@ -87,6 +91,8 @@ class LanguageModelService(BaseService):
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: Optional[list[LanguageModelTool]] = None,
         other_options: Optional[dict] = None,
+        structured_output_model: Optional[Type[BaseModel]] = None,
+        structured_output_enforce_schema: bool = False,
         logger: Optional[logging.Logger] = logging.getLogger(__name__),
     ) -> LanguageModelResponse:
         """
@@ -112,12 +118,14 @@ class LanguageModelService(BaseService):
         Raises:
             Exception: If an error occurs during the request, an exception is raised and logged.
         """
-        options, model, messages_dict, _ = cls.prepare_completion_params_util(
+        options, model, messages_dict, _ = self._prepare_completion_params_util(
             messages=messages,
             model_name=model_name,
             temperature=temperature,
             tools=tools,
             other_options=other_options,
+            structured_output_model=structured_output_model,
+            structured_output_enforce_schema=structured_output_enforce_schema,
         )
 
         try:
@@ -205,7 +213,7 @@ class LanguageModelService(BaseService):
             The LanguageModelStreamResponse object once the stream has finished.
         """
         options, model, messages_dict, search_context = (
-            self.prepare_completion_params_util(
+            self._prepare_completion_params_util(
                 messages=messages,
                 model_name=model_name,
                 temperature=temperature,
@@ -268,7 +276,7 @@ class LanguageModelService(BaseService):
             The LanguageModelStreamResponse object once the stream has finished.
         """
         options, model, messages_dict, search_context = (
-            self.prepare_completion_params_util(
+            self._prepare_completion_params_util(
                 messages=messages,
                 model_name=model_name,
                 temperature=temperature,
@@ -335,15 +343,32 @@ class LanguageModelService(BaseService):
             ]
         return options
 
-    @classmethod
-    def prepare_completion_params_util(
-        cls,
+    @staticmethod
+    def _add_response_format_to_options(
+        options: dict,
+        structured_output_model: Type[BaseModel],
+        structured_output_enforce_schema: bool = False,
+    ) -> dict:
+        options["responseFormat"] = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": structured_output_model.__name__,
+                "strict": structured_output_enforce_schema,
+                "schema": structured_output_model.model_json_schema(),
+            },
+        }
+        return options
+
+    def _prepare_completion_params_util(
+        self,
         messages: LanguageModelMessages,
         model_name: LanguageModelName | str,
         temperature: float,
         tools: Optional[list[LanguageModelTool]] = None,
         other_options: Optional[dict] = None,
         content_chunks: Optional[list[ContentChunk]] = None,
+        structured_output_model: Optional[Type[BaseModel]] = None,
+        structured_output_enforce_schema: bool = False,
     ) -> tuple[dict, str, dict, Optional[dict]]:
         """
         Prepares common parameters for completion requests.
@@ -356,8 +381,15 @@ class LanguageModelService(BaseService):
             - search_context (Optional[dict]): Processed content chunks if provided
         """
 
-        options = cls._add_tools_to_options({}, tools)
+        options = self._add_tools_to_options({}, tools)
+
+        if structured_output_model:
+            options = self._add_response_format_to_options(
+                options, structured_output_model, structured_output_enforce_schema
+            )
+
         options["temperature"] = temperature
+
         if other_options:
             options.update(other_options)
 

{unique_toolkit-0.5.49.dist-info → unique_toolkit-0.5.50.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 0.5.49
+Version: 0.5.50
 Summary: 
 License: Proprietary
 Author: Martin Fadler
@@ -100,6 +100,10 @@ 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).
 
+## [0.5.50] - 2025-01-30
+- Add the possibility to define completion output structure through a pydantic class
+
+
 ## [0.5.49] - 2025-01-24
 - Add `parsed` and `refusal` to `LanguageModelAssistantMessage` to support structured output
 

{unique_toolkit-0.5.49.dist-info → unique_toolkit-0.5.50.dist-info}/RECORD

@@ -40,11 +40,11 @@ unique_toolkit/language_model/builder.py,sha256=nsRqWO_2dgFehK5CgtqR5aqXgYUU0QL6
 unique_toolkit/language_model/infos.py,sha256=NgoV05ausVWMqrYqgH6i3s7tYG7mejupROIF_bwEGZo,13050
 unique_toolkit/language_model/prompt.py,sha256=JSawaLjQg3VR-E2fK8engFyJnNdk21zaO8pPIodzN4Q,3991
 unique_toolkit/language_model/schemas.py,sha256=87511yupvea-U6sfKWfelETevNMVPevhj7mEqX5FszU,7461
-unique_toolkit/language_model/service.py,sha256=brNCPRA0XxgqHi2rI5i2lyFCkUiw4MNMe1VaR3UgWmY,15500
+unique_toolkit/language_model/service.py,sha256=jBHFeGtPbaOeVBxg4XBwzCLjpkIDDAx_9eW7X_fOibk,16900
 unique_toolkit/language_model/utils.py,sha256=bPQ4l6_YO71w-zaIPanUUmtbXC1_hCvLK0tAFc3VCRc,1902
 unique_toolkit/short_term_memory/schemas.py,sha256=OhfcXyF6ACdwIXW45sKzjtZX_gkcJs8FEZXcgQTNenw,1406
 unique_toolkit/short_term_memory/service.py,sha256=Jd9P72-VvJy7hnqNrjmrmB5BHmsKuOpTiT0Jr-dBbsQ,1682
-unique_toolkit-0.5.49.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
-unique_toolkit-0.5.49.dist-info/METADATA,sha256=y2wgnOvbhd2PaXa5wshMxjmQ-hNrkaKhqalKAtr5K4E,15819
-unique_toolkit-0.5.49.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-unique_toolkit-0.5.49.dist-info/RECORD,,
+unique_toolkit-0.5.50.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
+unique_toolkit-0.5.50.dist-info/METADATA,sha256=ybCuYkNYzgWuZhRbEf_vSlVG1zMXWLaAEH8Es53iy20,15931
+unique_toolkit-0.5.50.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+unique_toolkit-0.5.50.dist-info/RECORD,,