unique_toolkit 0.7.13__py3-none-any.whl → 0.7.17__py3-none-any.whl

unique_toolkit/_common/validators.py

@@ -1,8 +1,57 @@
-from unique_toolkit.language_model import LanguageModel, LanguageModelName
+from typing import Annotated
 
+from pydantic import BeforeValidator, PlainSerializer
 
-def validate_and_init_language_model(value: LanguageModelName | LanguageModel | str):
-    if isinstance(value, LanguageModel):
-        return value
+from unique_toolkit.language_model import LanguageModelName
+from unique_toolkit.language_model.infos import (
+    LanguageModelInfo,
+    LanguageModelProvider,
+)
 
-    return LanguageModel(value)
+# TODO @klcd: Inform on deprecation of str as input
+LMI = Annotated[
+    LanguageModelInfo,
+    BeforeValidator(
+        lambda v: validate_and_init_language_model_info(v),
+        json_schema_input_type=str | LanguageModelName | LanguageModelInfo,
+    ),
+    PlainSerializer(
+        lambda v: serialize_lmi(v),
+        when_used="json",
+        return_type=str | LanguageModelInfo,
+    ),
+]
+
+
+def serialize_lmi(model: LanguageModelInfo) -> str | LanguageModelInfo:
+    if model.provider == LanguageModelProvider.CUSTOM:
+        return model
+
+    return model.name
+
+
+def validate_and_init_language_model_info(
+    v: str | LanguageModelName | LanguageModelInfo,
+) -> LanguageModelInfo:
+    """Validate and initialize a LanguageModelInfo object.
+
+    Args:
+        v: The input value to validate and initialize.
+
+    Returns:
+        LanguageModelInfo: The validated and initialized LanguageModelInfo object.
+
+    """
+    if isinstance(v, LanguageModelName):
+        return LanguageModelInfo.from_name(v)
+    if isinstance(v, str):
+        if v in [name.value for name in LanguageModelName]:
+            return LanguageModelInfo.from_name(LanguageModelName(v))
+
+        return LanguageModelInfo(
+            name=v,
+            version="custom",
+            provider=LanguageModelProvider.CUSTOM,
+        )
+
+    return v

unique_toolkit/app/schemas.py

@@ -109,31 +109,26 @@ class ChatEventPayload(BaseModel):
 @deprecated("""Use `ChatEventPayload` instead.
             This class will be removed in the next major version.""")
 class EventPayload(ChatEventPayload):
-    user_message: EventUserMessage
-    assistant_message: EventAssistantMessage
-    additional_parameters: Optional[EventAdditionalParameters] = None
-
-
-@deprecated(
-    """Use the more specific `ChatEvent` instead that has the same properties. \
-This class will be removed in the next major version."""
-)
-class Event(BaseModel):
-    model_config = model_config
-
-    id: str
-    event: EventName
-    user_id: str
-    company_id: str
-    payload: EventPayload
-    created_at: Optional[int] = None
-    version: Optional[str] = None
+    pass
+    # user_message: EventUserMessage
+    # assistant_message: EventAssistantMessage
+    # additional_parameters: Optional[EventAdditionalParameters] = None
 
 
 class ChatEvent(BaseEvent):
     model_config = model_config
 
-    event: EventName
     payload: ChatEventPayload
     created_at: Optional[int] = None
     version: Optional[str] = None
+
+
+@deprecated(
+    """Use the more specific `ChatEvent` instead that has the same properties. \
+This class will be removed in the next major version."""
+)
+class Event(ChatEvent):
+    pass
+    # The below should only affect type hints
+    # event: EventName T
+    # payload: EventPayload

unique_toolkit/chat/service.py

@@ -41,6 +41,7 @@ from unique_toolkit.language_model.infos import (
 )
 from unique_toolkit.language_model.schemas import (
     LanguageModelMessages,
+    LanguageModelResponse,
     LanguageModelStreamResponse,
     LanguageModelTool,
 )
@@ -81,7 +82,7 @@ class ChatService:
     @deprecated(
         "The event property is deprecated and will be removed in a future version."
     )
-    def event(self) -> Event | ChatEvent | None:
+    def event(self) -> Event | ChatEvent:
         """
         Get the event object (deprecated).
 
@@ -279,14 +280,30 @@ class ChatService:
         Args:
             debug_info (dict): The new debug information.
         """
-
+        [
+            company_id,
+            user_id,
+            assistant_message_id,
+            user_message_id,
+            chat_id,
+            user_message_text,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_message_id,
+                self._user_message_id,
+                self._chat_id,
+                self._user_message_text,
+            ]
+        )
         return await modify_message_async(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            assistant_message_id=self._assistant_message_id,
-            chat_id=self._chat_id,
-            user_message_id=self._user_message_id,
-            user_message_text=self._user_message_text,
+            user_id=user_id,
+            company_id=company_id,
+            assistant_message_id=assistant_message_id,
+            chat_id=chat_id,
+            user_message_id=user_message_id,
+            user_message_text=user_message_text,
             assistant=False,
             debug_info=debug_info,
         )
@@ -298,14 +315,31 @@ class ChatService:
         Args:
             debug_info (dict): The new debug information.
         """
+        [
+            company_id,
+            user_id,
+            assistant_message_id,
+            user_message_id,
+            chat_id,
+            user_message_text,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_message_id,
+                self._user_message_id,
+                self._chat_id,
+                self._user_message_text,
+            ]
+        )
 
         return modify_message(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            assistant_message_id=self._assistant_message_id,
-            chat_id=self._chat_id,
-            user_message_id=self._user_message_id,
-            user_message_text=self._user_message_text,
+            user_id=user_id,
+            company_id=company_id,
+            assistant_message_id=assistant_message_id,
+            chat_id=chat_id,
+            user_message_id=user_message_id,
+            user_message_text=user_message_text,
             assistant=False,
             debug_info=debug_info,
         )
@@ -334,13 +368,31 @@ class ChatService:
         Raises:
             Exception: If the modification fails.
         """
+        [
+            company_id,
+            user_id,
+            assistant_message_id,
+            user_message_id,
+            chat_id,
+            user_message_text,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_message_id,
+                self._user_message_id,
+                self._chat_id,
+                self._user_message_text,
+            ]
+        )
+
         return modify_message(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            assistant_message_id=self._assistant_message_id,
-            chat_id=self._chat_id,
-            user_message_id=self._user_message_id,
-            user_message_text=self._user_message_text,
+            user_id=user_id,
+            company_id=company_id,
+            assistant_message_id=assistant_message_id,
+            chat_id=chat_id,
+            user_message_id=user_message_id,
+            user_message_text=user_message_text,
             assistant=False,
             content=content,
             references=references,
@@ -373,13 +425,32 @@ class ChatService:
         Raises:
             Exception: If the modification fails.
         """
+
+        [
+            company_id,
+            user_id,
+            assistant_message_id,
+            user_message_id,
+            chat_id,
+            user_message_text,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_message_id,
+                self._user_message_id,
+                self._chat_id,
+                self._user_message_text,
+            ]
+        )
+
         return await modify_message_async(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            assistant_message_id=self._assistant_message_id,
-            chat_id=self._chat_id,
-            user_message_id=self._user_message_id,
-            user_message_text=self._user_message_text,
+            user_id=user_id,
+            company_id=company_id,
+            assistant_message_id=assistant_message_id,
+            chat_id=chat_id,
+            user_message_id=user_message_id,
+            user_message_text=user_message_text,
             assistant=False,
             content=content,
             references=references,
@@ -414,13 +485,31 @@ class ChatService:
         Raises:
             Exception: If the modification fails.
         """
+        [
+            company_id,
+            user_id,
+            assistant_message_id,
+            user_message_id,
+            chat_id,
+            user_message_text,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_message_id,
+                self._user_message_id,
+                self._chat_id,
+                self._user_message_text,
+            ]
+        )
+
         return modify_message(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            assistant_message_id=self._assistant_message_id,
-            chat_id=self._chat_id,
-            user_message_id=self._user_message_id,
-            user_message_text=self._user_message_text,
+            user_id=user_id,
+            company_id=company_id,
+            assistant_message_id=assistant_message_id,
+            chat_id=chat_id,
+            user_message_id=user_message_id,
+            user_message_text=user_message_text,
             assistant=True,
             content=content,
             original_content=original_content,
@@ -456,14 +545,30 @@ class ChatService:
         Raises:
             Exception: If the modification fails.
         """
-
+        [
+            company_id,
+            user_id,
+            assistant_message_id,
+            user_message_id,
+            chat_id,
+            user_message_text,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_message_id,
+                self._user_message_id,
+                self._chat_id,
+                self._user_message_text,
+            ]
+        )
         return await modify_message_async(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            assistant_message_id=self._assistant_message_id,
-            chat_id=self._chat_id,
-            user_message_id=self._user_message_id,
-            user_message_text=self._user_message_text,
+            user_id=user_id,
+            company_id=company_id,
+            assistant_message_id=assistant_message_id,
+            chat_id=chat_id,
+            user_message_id=user_message_id,
+            user_message_text=user_message_text,
             assistant=True,
             content=content,
             original_content=original_content,
@@ -595,11 +700,25 @@ class ChatService:
         Raises:
             Exception: If the creation fails.
         """
+        [
+            company_id,
+            user_id,
+            assistant_id,
+            chat_id,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_id,
+                self._chat_id,
+            ]
+        )
+
         chat_message = create_message(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            chat_id=self._chat_id,
-            assistant_id=self._assistant_id,
+            user_id=user_id,
+            company_id=company_id,
+            chat_id=chat_id,
+            assistant_id=assistant_id,
             role=ChatMessageRole.ASSISTANT,
             content=content,
             original_content=original_content,
@@ -635,12 +754,24 @@ class ChatService:
         Raises:
             Exception: If the creation fails.
         """
-
+        [
+            company_id,
+            user_id,
+            assistant_id,
+            chat_id,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_id,
+                self._chat_id,
+            ]
+        )
         chat_message = await create_message_async(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            chat_id=self._chat_id,
-            assistant_id=self._assistant_id,
+            user_id=user_id,
+            company_id=company_id,
+            chat_id=chat_id,
+            assistant_id=assistant_id,
             role=ChatMessageRole.ASSISTANT,
             content=content,
             original_content=original_content,
@@ -676,11 +807,24 @@ class ChatService:
         Raises:
             Exception: If the creation fails.
         """
+        [
+            company_id,
+            user_id,
+            assistant_id,
+            chat_id,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_id,
+                self._chat_id,
+            ]
+        )
         chat_message = create_message(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            chat_id=self._chat_id,
-            assistant_id=self._assistant_id,
+            user_id=user_id,
+            company_id=company_id,
+            chat_id=chat_id,
+            assistant_id=assistant_id,
             role=ChatMessageRole.USER,
             content=content,
             original_content=original_content,
@@ -716,12 +860,24 @@ class ChatService:
         Raises:
             Exception: If the creation fails.
         """
-
+        [
+            company_id,
+            user_id,
+            assistant_id,
+            chat_id,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+                self._assistant_id,
+                self._chat_id,
+            ]
+        )
         chat_message = await create_message_async(
-            user_id=self._user_id,
-            company_id=self._company_id,
-            chat_id=self._chat_id,
-            assistant_id=self._assistant_id,
+            user_id=user_id,
+            company_id=company_id,
+            chat_id=chat_id,
+            assistant_id=assistant_id,
             role=ChatMessageRole.USER,
             content=content,
             original_content=original_content,
@@ -761,9 +917,19 @@ class ChatService:
         Raises:
             Exception: If the creation fails
         """
+        [
+            company_id,
+            user_id,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+            ]
+        )
+
         return create_message_assessment(
-            user_id=self._user_id,
-            company_id=self._company_id,
+            user_id=user_id,
+            company_id=company_id,
             assistant_message_id=assistant_message_id,
             status=status,
             type=type,
@@ -801,9 +967,19 @@ class ChatService:
         Raises:
             Exception: If the creation fails
         """
+        [
+            company_id,
+            user_id,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+            ]
+        )
+
         return await create_message_assessment_async(
-            user_id=self._user_id,
-            company_id=self._company_id,
+            user_id=user_id,
+            company_id=company_id,
             assistant_message_id=assistant_message_id,
             status=status,
             type=type,
@@ -839,9 +1015,19 @@ class ChatService:
         Raises:
             Exception: If the modification fails
         """
+        [
+            company_id,
+            user_id,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+            ]
+        )
+
         return modify_message_assessment(
-            user_id=self._user_id,
-            company_id=self._company_id,
+            user_id=user_id,
+            company_id=company_id,
             assistant_message_id=assistant_message_id,
             status=status,
             type=type,
@@ -876,9 +1062,19 @@ class ChatService:
         Raises:
             Exception: If the modification fails
         """
+        [
+            company_id,
+            user_id,
+        ] = validate_required_values(
+            [
+                self._company_id,
+                self._user_id,
+            ]
+        )
+
         return await modify_message_assessment_async(
-            user_id=self._user_id,
-            company_id=self._company_id,
+            user_id=user_id,
+            company_id=company_id,
             assistant_message_id=assistant_message_id,
             status=status,
             type=type,
@@ -938,6 +1134,32 @@ class ChatService:
             other_options=other_options,
         )
 
+    def complete(
+        self,
+        messages: LanguageModelMessages,
+        model_name: LanguageModelName | str,
+        content_chunks: list[ContentChunk] = [],
+        debug_info: dict = {},
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
+        tools: Optional[list[LanguageModelTool]] = None,
+        start_text: Optional[str] = None,
+        other_options: Optional[dict] = None,
+    ) -> LanguageModelResponse:
+        response = self.stream_complete(
+            messages=messages,
+            model_name=model_name,
+            content_chunks=content_chunks,
+            debug_info=debug_info,
+            temperature=temperature,
+            timeout=timeout,
+            tools=tools,
+            start_text=start_text,
+            other_options=other_options,
+        )
+
+        return LanguageModelResponse.from_stream_response(response)
+
     async def stream_complete_async(
         self,
         messages: LanguageModelMessages,
@@ -989,3 +1211,29 @@ class ChatService:
             start_text=start_text,
             other_options=other_options,
         )
+
+    async def complete_async(
+        self,
+        messages: LanguageModelMessages,
+        model_name: LanguageModelName | str,
+        content_chunks: list[ContentChunk] = [],
+        debug_info: dict = {},
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
+        tools: Optional[list[LanguageModelTool]] = None,
+        start_text: Optional[str] = None,
+        other_options: Optional[dict] = None,
+    ) -> LanguageModelResponse:
+        response = self.stream_complete_async(
+            messages=messages,
+            model_name=model_name,
+            content_chunks=content_chunks,
+            debug_info=debug_info,
+            temperature=temperature,
+            timeout=timeout,
+            tools=tools,
+            start_text=start_text,
+            other_options=other_options,
+        )
+
+        return LanguageModelResponse.from_stream_response(await response)

unique_toolkit/content/functions.py

@@ -213,6 +213,7 @@ def upload_content_from_bytes(
     scope_id: str | None = None,
     chat_id: str | None = None,
     skip_ingestion: bool = False,
+    ingestion_config: unique_sdk.Content.IngestionConfig | None = None,
 ):
     """
     Uploads content to the knowledge base.
@@ -241,6 +242,7 @@ def upload_content_from_bytes(
             scope_id=scope_id,
             chat_id=chat_id,
             skip_ingestion=skip_ingestion,
+            ingestion_config=ingestion_config,
         )
     except Exception as e:
         logger.error(f"Error while uploading content: {e}")
@@ -256,6 +258,7 @@ def upload_content(
     scope_id: str | None = None,
     chat_id: str | None = None,
     skip_ingestion: bool = False,
+    ingestion_config: unique_sdk.Content.IngestionConfig | None = None,
 ):
     """
     Uploads content to the knowledge base.
@@ -284,6 +287,7 @@ def upload_content(
             scope_id=scope_id,
             chat_id=chat_id,
             skip_ingestion=skip_ingestion,
+            ingestion_config=ingestion_config,
         )
     except Exception as e:
         logger.error(f"Error while uploading content: {e}")
@@ -299,6 +303,7 @@ def _trigger_upload_content(
     scope_id: str | None = None,
     chat_id: str | None = None,
     skip_ingestion: bool = False,
+    ingestion_config: unique_sdk.Content.IngestionConfig | None = None,
 ):
     """
     Uploads content to the knowledge base.
@@ -368,16 +373,20 @@ def _trigger_upload_content(
         logger.error(error_msg)
         raise ValueError(error_msg)
 
+    if ingestion_config is None:
+        ingestion_config = {}
+
+    if skip_ingestion:
+        ingestion_config["uniqueIngestionMode"] = "SKIP_INGESTION"
+
     input_dict = {
         "key": content_name,
         "title": content_name,
         "mimeType": mime_type,
         "byteSize": byte_size,
+        "ingestionConfig": ingestion_config,
     }
 
-    if skip_ingestion:
-        input_dict["ingestionConfig"] = {"uniqueIngestionMode": "SKIP_INGESTION"}
-
     if chat_id:
         _upsert_content(
             user_id=user_id,