unique_toolkit 0.5.55__py3-none-any.whl → 0.6.0__py3-none-any.whl

unique_toolkit/embedding/constants.py

@@ -0,0 +1,2 @@
+DOMAIN_NAME = "embedding"
+DEFAULT_TIMEOUT = 600_000

unique_toolkit/embedding/functions.py

@@ -0,0 +1,79 @@
+import logging
+
+import unique_sdk
+
+from unique_toolkit.embedding.constants import DEFAULT_TIMEOUT, DOMAIN_NAME
+from unique_toolkit.embedding.schemas import Embeddings
+
+logger = logging.getLogger(f"toolkit.{DOMAIN_NAME}.{__name__}")
+
+
+def embed_texts(
+    user_id: str,
+    company_id: str,
+    texts: list[str],
+    timeout: int = DEFAULT_TIMEOUT,
+) -> Embeddings:
+    """
+    Embed text.
+
+    Args:
+        user_id (str): The user ID.
+        company_id (str): The company ID.
+        texts (list[str]): The texts to embed.
+        timeout (int): The timeout in milliseconds. Defaults to 600000.
+
+    Returns:
+        Embeddings: The Embedding object.
+
+    Raises:
+        Exception: If an error occurs.
+    """
+
+    try:
+        data = {
+            "user_id": user_id,
+            "company_id": company_id,
+            "texts": texts,
+            "timeout": timeout,
+        }
+        response = unique_sdk.Embeddings.create(**data)
+        return Embeddings(**response)
+    except Exception as e:
+        logger.error(f"Error embedding texts: {e}")
+        raise e
+
+
+async def embed_texts_async(
+    user_id: str,
+    company_id: str,
+    texts: list[str],
+    timeout: int = DEFAULT_TIMEOUT,
+) -> Embeddings:
+    """
+    Embed text asynchronously.
+
+    Args:
+        user_id (str): The user ID.
+        company_id (str): The company ID.
+        texts (list[str]): The texts to embed.
+        timeout (int): The timeout in milliseconds. Defaults to 600000.
+
+    Returns:
+        Embeddings: The Embedding object.
+
+    Raises:
+        Exception: If an error occurs.
+    """
+    try:
+        data = {
+            "user_id": user_id,
+            "company_id": company_id,
+            "texts": texts,
+            "timeout": timeout,
+        }
+        response = await unique_sdk.Embeddings.create_async(**data)
+        return Embeddings(**response)
+    except Exception as e:
+        logger.error(f"Error embedding texts: {e}")
+        raise e

unique_toolkit/embedding/service.py

@@ -1,10 +1,10 @@
-import logging
-from typing import Optional
-
-import unique_sdk
+from typing_extensions import deprecated
 
 from unique_toolkit._common._base_service import BaseService
-from unique_toolkit.app.schemas import Event
+from unique_toolkit._common.validate_required_values import validate_required_values
+from unique_toolkit.app.schemas import BaseEvent, Event
+from unique_toolkit.embedding.constants import DEFAULT_TIMEOUT
+from unique_toolkit.embedding.functions import embed_texts, embed_texts_async
 from unique_toolkit.embedding.schemas import Embeddings
 
 
@@ -13,14 +13,37 @@ class EmbeddingService(BaseService):
     Provides methods to interact with the Embedding service.
 
     Attributes:
-        event (Event): The Event object.
-        logger (Optional[logging.Logger]): The logger object. Defaults t∏o None.
+        company_id (str | None): The company ID.
+        user_id (str | None): The user ID.
     """
 
-    def __init__(self, event: Event, logger: Optional[logging.Logger] = None):
-        super().__init__(event, logger)
+    def __init__(
+        self,
+        event: Event | BaseEvent | None = None,
+        company_id: str | None = None,
+        user_id: str | None = None,
+    ):
+        self._event = event
+        if event:
+            self.company_id = event.company_id
+            self.user_id = event.user_id
+        else:
+            [company_id, user_id] = validate_required_values([company_id, user_id])
+            self.company_id = company_id
+            self.user_id = user_id
+
+    @property
+    @deprecated(
+        "The event property is deprecated and will be removed in a future version."
+    )
+    def event(self) -> Event | BaseEvent | None:
+        """
+        Get the event object (deprecated).
 
-    DEFAULT_TIMEOUT = 600_000
+        Returns:
+            Event | BaseEvent | None: The event object.
+        """
+        return self._event
 
     def embed_texts(
         self,
@@ -32,7 +55,7 @@ class EmbeddingService(BaseService):
 
         Args:
             text (str): The text to embed.
-            timeout (int): The timeout in milliseconds. Defaults to None.
+            timeout (int): The timeout in milliseconds. Defaults to 600000.
 
         Returns:
             Embeddings: The Embedding object.
@@ -40,13 +63,12 @@ class EmbeddingService(BaseService):
         Raises:
             Exception: If an error occurs.
         """
-        request = self._get_request_obj(texts=texts, timeout=timeout)
-        try:
-            response = unique_sdk.Embeddings.create(**request)
-            return Embeddings(**response)
-        except Exception as e:
-            self.logger.error(f"Error embedding texts: {e}")
-            raise e
+        return embed_texts(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            texts=texts,
+            timeout=timeout,
+        )
 
     async def embed_texts_async(
         self,
@@ -58,7 +80,7 @@ class EmbeddingService(BaseService):
 
         Args:
             text (str): The text to embed.
-            timeout (int): The timeout in milliseconds. Defaults to None.
+            timeout (int): The timeout in milliseconds. Defaults to 600000.
 
         Returns:
             Embeddings: The Embedding object.
@@ -66,18 +88,9 @@ class EmbeddingService(BaseService):
         Raises:
             Exception: If an error occurs.
         """
-        request = self._get_request_obj(texts=texts, timeout=timeout)
-        try:
-            response = await unique_sdk.Embeddings.create_async(**request)
-            return Embeddings(**response)
-        except Exception as e:
-            self.logger.error(f"Error embedding texts: {e}")
-            raise e
-
-    def _get_request_obj(self, texts: list[str], timeout: int) -> dict:
-        return {
-            "user_id": self.event.user_id,
-            "company_id": self.event.company_id,
-            "texts": texts,
-            "timeout": timeout,
-        }
+        return await embed_texts_async(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            texts=texts,
+            timeout=timeout,
+        )

unique_toolkit/evaluators/__init__.py

@@ -0,0 +1 @@
+from .constants import DOMAIN_NAME as DOMAIN_NAME

unique_toolkit/evaluators/constants.py

@@ -0,0 +1 @@
+DOMAIN_NAME = "evaluators"

unique_toolkit/evaluators/context_relevancy/constants.py

@@ -4,6 +4,7 @@ from unique_toolkit.evaluators.context_relevancy.prompts import (
     CONTEXT_RELEVANCY_METRIC_USER_MSG,
 )
 from unique_toolkit.evaluators.schemas import (
+    EvaluationMetricInputFieldName,
     EvaluationMetricName,
 )
 from unique_toolkit.language_model.infos import LanguageModel
@@ -14,9 +15,8 @@ USER_MSG_KEY = "userPrompt"
 
 # Required input fields for context relevancy evaluation
 context_relevancy_required_input_fields = [
-    "input_text",
-    "output_text",
-    "context_texts",
+    EvaluationMetricInputFieldName.INPUT_TEXT,
+    EvaluationMetricInputFieldName.CONTEXT_TEXTS,
 ]
 
 

unique_toolkit/evaluators/context_relevancy/utils.py

@@ -79,7 +79,9 @@ async def check_context_relevancy_async(
     try:
         msgs = _get_msgs(input, config)
         result = await LanguageModelService.complete_async_util(
-            company_id=company_id, messages=msgs, model_name=model_name
+            company_id=company_id,
+            messages=msgs,
+            model_name=model_name,
         )
         result_content = result.choices[0].message.content
         if not result_content:
@@ -89,7 +91,8 @@ async def check_context_relevancy_async(
                 user_message=error_message,
             )
         return parse_eval_metric_result(
-            result_content, EvaluationMetricName.CONTEXT_RELEVANCY
+            result_content,  # type: ignore
+            EvaluationMetricName.CONTEXT_RELEVANCY,
         )
     except Exception as e:
         error_message = "Error occurred during context relevancy metric analysis"

unique_toolkit/evaluators/hallucination/utils.py

@@ -91,7 +91,8 @@ async def check_hallucination_async(
                 user_message=error_message,
             )
         return parse_eval_metric_result(
-            result_content, EvaluationMetricName.HALLUCINATION
+            result_content,  # type: ignore
+            EvaluationMetricName.HALLUCINATION,
         )
     except Exception as e:
         error_message = "Error occurred during hallucination metric analysis"

unique_toolkit/language_model/__init__.py

@@ -1,3 +1,4 @@
+from .constants import DOMAIN_NAME as DOMAIN_NAME
 from .infos import LanguageModel as LanguageModel
 from .infos import LanguageModelName as LanguageModelName
 from .prompt import (

unique_toolkit/language_model/constants.py

@@ -0,0 +1,4 @@
+DOMAIN_NAME = "language_model"
+DEFAULT_COMPLETE_TIMEOUT = 240_000
+DEFAULT_COMPLETE_TEMPERATURE = 0.0
+DEFAULT_SEED = 42

unique_toolkit/language_model/functions.py

@@ -0,0 +1,362 @@
+import logging
+from typing import Optional, Type, cast
+
+import unique_sdk
+from pydantic import BaseModel
+
+from unique_toolkit.content.schemas import ContentChunk
+from unique_toolkit.evaluators import DOMAIN_NAME
+from unique_toolkit.language_model.constants import (
+    DEFAULT_COMPLETE_TEMPERATURE,
+    DEFAULT_COMPLETE_TIMEOUT,
+)
+from unique_toolkit.language_model.infos import LanguageModelName
+from unique_toolkit.language_model.schemas import (
+    LanguageModelMessages,
+    LanguageModelResponse,
+    LanguageModelStreamResponse,
+    LanguageModelTool,
+)
+
+logger = logging.getLogger(f"toolkit.{DOMAIN_NAME}.{__name__}")
+
+
+def complete(
+    company_id: str,
+    messages: LanguageModelMessages,
+    model_name: LanguageModelName | str,
+    temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+    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,
+) -> LanguageModelResponse:
+    """
+    Calls the completion endpoint synchronously without streaming the response.
+
+    Args:
+        company_id (str): The company ID associated with the request.
+        messages (LanguageModelMessages): The messages to complete.
+        model_name (LanguageModelName | str): The model name to use for the completion.
+        temperature (float): The temperature setting for the completion. Defaults to 0.
+        timeout (int): The timeout value in milliseconds. Defaults to 240_000.
+        tools (Optional[list[LanguageModelTool]]): Optional list of tools to include.
+        other_options (Optional[dict]): Additional options to use. Defaults to None.
+
+    Returns:
+        LanguageModelResponse: The response object containing the completed result.
+    """
+    options, model, messages_dict, _ = _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:
+        response = unique_sdk.ChatCompletion.create(
+            company_id=company_id,
+            model=model,
+            messages=cast(
+                list[unique_sdk.Integrated.ChatCompletionRequestMessage],
+                messages_dict,
+            ),
+            timeout=timeout,
+            options=options,  # type: ignore
+        )
+        return LanguageModelResponse(**response)
+    except Exception as e:
+        logger.error(f"Error completing: {e}")
+        raise e
+
+
+async def complete_async(
+    company_id: str,
+    messages: LanguageModelMessages,
+    model_name: LanguageModelName | str,
+    temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+    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,
+) -> LanguageModelResponse:
+    """
+    Calls the completion endpoint asynchronously without streaming the response.
+
+    This method sends a request to the completion endpoint using the provided messages, model name,
+    temperature, timeout, and optional tools. It returns a `LanguageModelResponse` object containing
+    the completed result.
+
+    Args:
+        company_id (str): The company ID associated with the request.
+        messages (LanguageModelMessages): The messages to complete.
+        model_name (LanguageModelName | str): The model name to use for the completion.
+        temperature (float): The temperature setting for the completion. Defaults to 0.
+        timeout (int): The timeout value in milliseconds for the request. Defaults to 240_000.
+        tools (Optional[list[LanguageModelTool]]): Optional list of tools to include in the request.
+        other_options (Optional[dict]): The other options to use. Defaults to None.
+
+    Returns:
+        LanguageModelResponse: The response object containing the completed result.
+
+    Raises:
+        Exception: If an error occurs during the request, an exception is raised and logged.
+    """
+    options, model, messages_dict, _ = _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:
+        response = await unique_sdk.ChatCompletion.create_async(
+            company_id=company_id,
+            model=model,
+            messages=cast(
+                list[unique_sdk.Integrated.ChatCompletionRequestMessage],
+                messages_dict,
+            ),
+            timeout=timeout,
+            options=options,  # type: ignore
+        )
+        return LanguageModelResponse(**response)
+    except Exception as e:
+        logger.error(f"Error completing: {e}")  # type: ignore
+        raise e
+
+
+def stream_complete_to_chat(
+    company_id: str,
+    user_id: str,
+    assistant_message_id: str,
+    user_message_id: str,
+    chat_id: str,
+    assistant_id: str,
+    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,
+) -> LanguageModelStreamResponse:
+    """
+    Streams a completion synchronously.
+
+    Args:
+        company_id (str): The company ID associated with the request.
+        user_id (str): The user ID for the request.
+        assistant_message_id (str): The assistant message ID.
+        user_message_id (str): The user message ID.
+        chat_id (str): The chat ID.
+        assistant_id (str): The assistant ID.
+        messages (LanguageModelMessages): The messages to complete.
+        model_name (LanguageModelName | str): The model name.
+        content_chunks (list[ContentChunk]): Content chunks for context.
+        debug_info (dict): Debug information.
+        temperature (float): Temperature setting.
+        timeout (int): Timeout in milliseconds.
+        tools (Optional[list[LanguageModelTool]]): Optional tools.
+        start_text (Optional[str]): Starting text.
+        other_options (Optional[dict]): Additional options.
+
+    Returns:
+        LanguageModelStreamResponse: The streaming response object.
+    """
+    options, model, messages_dict, search_context = _prepare_completion_params_util(
+        messages=messages,
+        model_name=model_name,
+        temperature=temperature,
+        tools=tools,
+        other_options=other_options,
+        content_chunks=content_chunks,
+    )
+
+    try:
+        response = unique_sdk.Integrated.chat_stream_completion(
+            user_id=user_id,
+            company_id=company_id,
+            assistantMessageId=assistant_message_id,
+            userMessageId=user_message_id,
+            messages=cast(
+                list[unique_sdk.Integrated.ChatCompletionRequestMessage],
+                messages_dict,
+            ),
+            chatId=chat_id,
+            searchContext=search_context,
+            model=model,
+            timeout=timeout,
+            assistantId=assistant_id,
+            debugInfo=debug_info,
+            options=options,  # type: ignore
+            startText=start_text,
+        )
+        return LanguageModelStreamResponse(**response)
+    except Exception as e:
+        logger.error(f"Error streaming completion: {e}")
+        raise e
+
+
+async def stream_complete_to_chat_async(
+    company_id: str,
+    user_id: str,
+    assistant_message_id: str,
+    user_message_id: str,
+    chat_id: str,
+    assistant_id: str,
+    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,
+) -> LanguageModelStreamResponse:
+    """
+    Streams a completion asynchronously.
+
+    Args: [same as stream_complete]
+
+    Returns:
+        LanguageModelStreamResponse: The streaming response object.
+    """
+    options, model, messages_dict, search_context = _prepare_completion_params_util(
+        messages=messages,
+        model_name=model_name,
+        temperature=temperature,
+        tools=tools,
+        other_options=other_options,
+        content_chunks=content_chunks,
+    )
+
+    try:
+        response = await unique_sdk.Integrated.chat_stream_completion_async(
+            user_id=user_id,
+            company_id=company_id,
+            assistantMessageId=assistant_message_id,
+            userMessageId=user_message_id,
+            messages=cast(
+                list[unique_sdk.Integrated.ChatCompletionRequestMessage],
+                messages_dict,
+            ),
+            chatId=chat_id,
+            searchContext=search_context,
+            model=model,
+            timeout=timeout,
+            assistantId=assistant_id,
+            debugInfo=debug_info,
+            options=options,  # type: ignore
+            startText=start_text,
+        )
+        return LanguageModelStreamResponse(**response)
+    except Exception as e:
+        logger.error(f"Error streaming completion: {e}")
+        raise e
+
+
+def _add_tools_to_options(
+    options: dict,
+    tools: Optional[list[LanguageModelTool]],
+) -> dict:
+    if tools:
+        options["tools"] = [
+            {
+                "type": "function",
+                "function": tool.model_dump(exclude_none=True),
+            }
+            for tool in tools
+        ]
+    return options
+
+
+def _to_search_context(chunks: list[ContentChunk]) -> dict | None:
+    if not chunks:
+        return None
+    return [
+        unique_sdk.Integrated.SearchResult(
+            id=chunk.id,
+            chunkId=chunk.chunk_id,
+            key=chunk.key,
+            title=chunk.title,
+            url=chunk.url,
+            startPage=chunk.start_page,
+            endPage=chunk.end_page,
+            order=chunk.order,
+            object=chunk.object,
+        )  # type: ignore
+        for chunk in chunks
+    ]
+
+
+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(
+    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.
+
+    Returns:
+        tuple containing:
+        - options (dict): Combined options including tools and temperature
+        - model (str): Resolved model name
+        - messages_dict (dict): Processed messages
+        - search_context (Optional[dict]): Processed content chunks if provided
+    """
+
+    options = _add_tools_to_options({}, tools)
+    if structured_output_model:
+        options = _add_response_format_to_options(
+            options, structured_output_model, structured_output_enforce_schema
+        )
+    options["temperature"] = temperature
+    if other_options:
+        options.update(other_options)
+
+    model = model_name.name if isinstance(model_name, LanguageModelName) else model_name
+
+    # Different methods need different message dump parameters
+    messages_dict = messages.model_dump(
+        exclude_none=True,
+        by_alias=content_chunks is not None,  # Use by_alias for streaming methods
+    )
+
+    search_context = (
+        _to_search_context(content_chunks) if content_chunks is not None else None
+    )
+
+    return options, model, messages_dict, search_context