unique_toolkit 0.5.6__py3-none-any.whl → 0.5.7__py3-none-any.whl

unique_toolkit/content/service.py

@@ -7,7 +7,7 @@ from typing import Optional, cast
 import requests
 import unique_sdk
 
-from unique_toolkit.app.performance.async_wrapper import async_warning, to_async
+from unique_toolkit._common._base_service import BaseService
 from unique_toolkit.chat.state import ChatState
 from unique_toolkit.content.schemas import (
     Content,
@@ -18,21 +18,29 @@ from unique_toolkit.content.schemas import (
 )
 
 
-class ContentService:
-    DEFAULT_SEARCH_LANGUAGE = "english"
+class ContentService(BaseService):
+    """
+    Provides methods for searching, downloading and uploading content in the knowledge base.
+
+    Attributes:
+        state (ChatState): The chat state.
+        logger (Optional[logging.Logger]): The logger. Defaults to None.
+    """
 
     def __init__(self, state: ChatState, logger: Optional[logging.Logger] = None):
-        self.state = state
-        self.logger = logger or logging.getLogger(__name__)
+        super().__init__(state, logger)
+
+    DEFAULT_SEARCH_LANGUAGE = "english"
 
     def search_content_chunks(
         self,
         search_string: str,
         search_type: ContentSearchType,
         limit: int,
-        reranker_config: Optional[RerankerConfig] = None,
         search_language: str = DEFAULT_SEARCH_LANGUAGE,
+        reranker_config: Optional[RerankerConfig] = None,
         scope_ids: Optional[list[str]] = None,
+        chat_only: Optional[bool] = None,
     ) -> list[ContentChunk]:
         """
         Performs a synchronous search for content chunks in the knowledge base.
@@ -41,32 +49,50 @@ class ContentService:
             search_string (str): The search string.
             search_type (ContentSearchType): The type of search to perform.
             limit (int): The maximum number of results to return.
-            reranker_config (Optional[RerankerConfig]): The reranker configuration. Defaults to None.
             search_language (str): The language for the full-text search. Defaults to "english".
+            reranker_config (Optional[RerankerConfig]): The reranker configuration. Defaults to None.
             scope_ids (Optional[list[str]]): The scope IDs. Defaults to None.
+            chat_only (Optional[bool]): Whether to search only in the current chat. Defaults to None.
 
         Returns:
             list[ContentChunk]: The search results.
         """
-        return self._trigger_search_content_chunks(
-            search_string=search_string,
-            search_type=search_type,
-            limit=limit,
-            reranker_config=reranker_config,
-            search_language=search_language,
-            scope_ids=scope_ids,
-        )
+        if not scope_ids:
+            self.logger.warning("No scope IDs provided for search.")
 
-    @to_async
-    @async_warning
-    def async_search_content_chunks(
+        try:
+            searches = unique_sdk.Search.create(
+                user_id=self.state.user_id,
+                company_id=self.state.company_id,
+                chatId=self.state.chat_id,
+                searchString=search_string,
+                searchType=search_type.name,
+                scopeIds=scope_ids,
+                limit=limit,
+                reranker=reranker_config.model_dump() if reranker_config else None,
+                language=search_language,
+                chatOnly=chat_only,
+            )
+        except Exception as e:
+            self.logger.error(f"Error while searching content chunks: {e}")
+            raise e
+
+        def map_to_content_chunks(searches: list[unique_sdk.Search]):
+            return [ContentChunk(**search) for search in searches]
+
+        # TODO change return type in sdk from Search to list[Search]
+        searches = cast(list[unique_sdk.Search], searches)
+        return map_to_content_chunks(searches)
+
+    async def search_content_chunks_async(
         self,
         search_string: str,
         search_type: ContentSearchType,
         limit: int,
-        reranker_config: Optional[RerankerConfig] = None,
         search_language: str = DEFAULT_SEARCH_LANGUAGE,
+        reranker_config: Optional[RerankerConfig] = None,
         scope_ids: Optional[list[str]] = None,
+        chat_only: Optional[bool] = None,
     ):
         """
         Performs an asynchronous search for content chunks in the knowledge base.
@@ -75,36 +101,19 @@ class ContentService:
             search_string (str): The search string.
             search_type (ContentSearchType): The type of search to perform.
             limit (int): The maximum number of results to return.
-            reranker_config (Optional[RerankerConfig]): The reranker configuration. Defaults to None.
             search_language (str): The language for the full-text search. Defaults to "english".
-            scope_ids (Optional[list[str]]): The scope IDs. Defaults to [].
+            reranker_config (Optional[RerankerConfig]): The reranker configuration. Defaults to None.
+            scope_ids (Optional[list[str]]): The scope IDs. Defaults to None.
+            chat_only (Optional[bool]): Whether to search only in the current chat. Defaults to None.
 
         Returns:
             list[ContentChunk]: The search results.
         """
-        return self._trigger_search_content_chunks(
-            search_string=search_string,
-            search_type=search_type,
-            limit=limit,
-            reranker_config=reranker_config,
-            search_language=search_language,
-            scope_ids=scope_ids,
-        )
-
-    def _trigger_search_content_chunks(
-        self,
-        search_string: str,
-        search_type: ContentSearchType,
-        limit: int,
-        reranker_config: Optional[RerankerConfig] = None,
-        search_language: str = DEFAULT_SEARCH_LANGUAGE,
-        scope_ids: Optional[list[str]] = None,
-    ) -> list[ContentChunk]:
         if not scope_ids:
             self.logger.warning("No scope IDs provided for search.")
 
         try:
-            searches = unique_sdk.Search.create(
+            searches = await unique_sdk.Search.create_async(
                 user_id=self.state.user_id,
                 company_id=self.state.company_id,
                 chatId=self.state.chat_id,
@@ -114,7 +123,7 @@ class ContentService:
                 limit=limit,
                 reranker=reranker_config.model_dump() if reranker_config else None,
                 language=search_language,
-                chatOnly=self.state.chat_only,
+                chatOnly=chat_only,
             )
         except Exception as e:
             self.logger.error(f"Error while searching content chunks: {e}")
@@ -141,11 +150,21 @@ class ContentService:
         Returns:
             list[Content]: The search results.
         """
-        return self._trigger_search_contents(where)
+        try:
+            contents = unique_sdk.Content.search(
+                user_id=self.state.user_id,
+                company_id=self.state.company_id,
+                chatId=self.state.chat_id,
+                # TODO add type parameter
+                where=where,  # type: ignore
+            )
+        except Exception as e:
+            self.logger.error(f"Error while searching contents: {e}")
+            raise e
 
-    @to_async
-    @async_warning
-    def async_search_contents(
+        return self._map_contents(contents)
+
+    async def search_contents_async(
         self,
         where: dict,
     ) -> list[Content]:
@@ -158,35 +177,8 @@ class ContentService:
         Returns:
             list[Content]: The search results.
         """
-        return self._trigger_search_contents(where)
-
-    def _trigger_search_contents(
-        self,
-        where: dict,
-    ) -> list[Content]:
-        def map_content_chunk(content_chunk):
-            return ContentChunk(
-                id=content_chunk["id"],
-                text=content_chunk["text"],
-                start_page=content_chunk["startPage"],
-                end_page=content_chunk["endPage"],
-                order=content_chunk["order"],
-            )
-
-        def map_content(content):
-            return Content(
-                id=content["id"],
-                key=content["key"],
-                title=content["title"],
-                url=content["url"],
-                chunks=[map_content_chunk(chunk) for chunk in content["chunks"]],
-            )
-
-        def map_contents(contents):
-            return [map_content(content) for content in contents]
-
         try:
-            contents = unique_sdk.Content.search(
+            contents = await unique_sdk.Content.search_async(
                 user_id=self.state.user_id,
                 company_id=self.state.company_id,
                 chatId=self.state.chat_id,
@@ -197,7 +189,29 @@ class ContentService:
             self.logger.error(f"Error while searching contents: {e}")
             raise e
 
-        return map_contents(contents)
+        return self._map_contents(contents)
+
+    @staticmethod
+    def _map_content_chunk(content_chunk: dict):
+        return ContentChunk(
+            id=content_chunk["id"],
+            text=content_chunk["text"],
+            start_page=content_chunk["startPage"],
+            end_page=content_chunk["endPage"],
+            order=content_chunk["order"],
+        )
+
+    def _map_content(self, content: dict):
+        return Content(
+            id=content["id"],
+            key=content["key"],
+            title=content["title"],
+            url=content["url"],
+            chunks=[self._map_content_chunk(chunk) for chunk in content["chunks"]],
+        )
+
+    def _map_contents(self, contents):
+        return [self._map_content(content) for content in contents]
 
     def upload_content(
         self,

unique_toolkit/embedding/service.py

@@ -4,22 +4,29 @@ from typing import Optional
 import numpy as np
 import unique_sdk
 
-from unique_toolkit.app.performance.async_wrapper import async_warning, to_async
+from unique_toolkit._common._base_service import BaseService
 from unique_toolkit.chat.state import ChatState
 from unique_toolkit.embedding.schemas import Embeddings
 
 
-class EmbeddingService:
+class EmbeddingService(BaseService):
+    """
+    Provides methods to interact with the Embedding service.
+
+    Attributes:
+        state (ChatState): The ChatState object.
+        logger (Optional[logging.Logger]): The logger object. Defaults t∏o None.
+    """
+
     def __init__(self, state: ChatState, logger: Optional[logging.Logger] = None):
-        self.state = state
-        self.logger = logger or logging.getLogger(__name__)
+        super().__init__(state, logger)
 
-    _DEFAULT_TIMEOUT = 600_000
+    DEFAULT_TIMEOUT = 600_000
 
     def embed_texts(
         self,
         texts: list[str],
-        timeout: int = _DEFAULT_TIMEOUT,
+        timeout: int = DEFAULT_TIMEOUT,
     ) -> Embeddings:
         """
         Embed text.
@@ -34,17 +41,18 @@ class EmbeddingService:
         Raises:
             Exception: If an error occurs.
         """
-        return self._trigger_embed_texts(
-            texts=texts,
-            timeout=timeout,
-        )
+        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
 
-    @to_async
-    @async_warning
-    def async_embed_texts(
+    async def embed_texts_async(
         self,
         texts: list[str],
-        timeout: int = _DEFAULT_TIMEOUT,
+        timeout: int = DEFAULT_TIMEOUT,
     ) -> Embeddings:
         """
         Embed text asynchronously.
@@ -59,24 +67,21 @@ class EmbeddingService:
         Raises:
             Exception: If an error occurs.
         """
-        return self._trigger_embed_texts(
-            texts=texts,
-            timeout=timeout,
-        )
+        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 _trigger_embed_texts(self, texts: list[str], timeout: int) -> Embeddings:
-        request = {
+    def _get_request_obj(self, texts: list[str], timeout: int) -> dict:
+        return {
             "user_id": self.state.user_id,
             "company_id": self.state.company_id,
             "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
 
     def get_cosine_similarity(
         self,

unique_toolkit/language_model/service.py

@@ -3,7 +3,7 @@ from typing import Optional, cast
 
 import unique_sdk
 
-from unique_toolkit.app.performance.async_wrapper import async_warning, to_async
+from unique_toolkit._common._base_service import BaseService
 from unique_toolkit.chat.state import ChatState
 from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.language_model.infos import LanguageModelName
@@ -15,20 +15,27 @@ from unique_toolkit.language_model.schemas import (
 )
 
 
-class LanguageModelService:
+class LanguageModelService(BaseService):
+    """
+    Provides methods to interact with the Language Model by generating responses.
+
+    Attributes:
+        state (ChatState): The ChatState object.
+        logger (Optional[logging.Logger]): The logger object. Defaults to None.
+    """
+
     def __init__(self, state: ChatState, logger: Optional[logging.Logger] = None):
-        self.state = state
-        self.logger = logger or logging.getLogger(__name__)
+        super().__init__(state, logger)
 
-    _DEFAULT_COMPLETE_TIMEOUT = 240_000
-    _DEFAULT_COMPLETE_TEMPERATURE = 0.0
+    DEFAULT_COMPLETE_TIMEOUT = 240_000
+    DEFAULT_COMPLETE_TEMPERATURE = 0.0
 
     def complete(
         self,
         messages: LanguageModelMessages,
         model_name: LanguageModelName,
-        temperature: float = _DEFAULT_COMPLETE_TEMPERATURE,
-        timeout: int = _DEFAULT_COMPLETE_TIMEOUT,
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: Optional[list[LanguageModelTool]] = None,
     ):
         """
@@ -44,22 +51,32 @@ class LanguageModelService:
         Returns:
             LanguageModelResponse: The LanguageModelResponse object.
         """
-        return self._trigger_complete(
-            messages=messages,
-            model_name=model_name,
-            temperature=temperature,
-            timeout=timeout,
-            tools=tools,
-        )
+        options = self._add_tools_to_options({}, tools)
+        messages = messages.model_dump(exclude_none=True)
+        try:
+            response = unique_sdk.ChatCompletion.create(
+                company_id=self.state.company_id,
+                # TODO change or extend types in unique_sdk
+                model=model_name.name,  # type: ignore
+                messages=cast(
+                    list[unique_sdk.Integrated.ChatCompletionRequestMessage],
+                    messages,
+                ),
+                timeout=timeout,
+                temperature=temperature,
+                options=options,  # type: ignore
+            )
+            return LanguageModelResponse(**response)
+        except Exception as e:
+            self.logger.error(f"Error completing: {e}")
+            raise e
 
-    @to_async
-    @async_warning
-    def async_complete(
+    async def complete_async(
         self,
         messages: LanguageModelMessages,
         model_name: LanguageModelName,
-        temperature: float = _DEFAULT_COMPLETE_TEMPERATURE,
-        timeout: int = _DEFAULT_COMPLETE_TIMEOUT,
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: Optional[list[LanguageModelTool]] = None,
     ):
         """
@@ -75,26 +92,10 @@ class LanguageModelService:
         Returns:
             str: The completed message content.
         """
-        return self._trigger_complete(
-            messages=messages,
-            model_name=model_name,
-            temperature=temperature,
-            timeout=timeout,
-            tools=tools,
-        )
-
-    def _trigger_complete(
-        self,
-        messages: LanguageModelMessages,
-        model_name: LanguageModelName,
-        temperature: float,
-        timeout: int,
-        tools: Optional[list[LanguageModelTool]] = None,
-    ) -> LanguageModelResponse:
         options = self._add_tools_to_options({}, tools)
         messages = messages.model_dump(exclude_none=True, exclude={"tool_calls"})
         try:
-            response = unique_sdk.ChatCompletion.create(
+            response = await unique_sdk.ChatCompletion.create_async(
                 company_id=self.state.company_id,
                 # TODO change or extend types in unique_sdk
                 model=model_name.name,  # type: ignore
@@ -106,20 +107,19 @@ class LanguageModelService:
                 temperature=temperature,
                 options=options,  # type: ignore
             )
+            return LanguageModelResponse(**response)
         except Exception as e:
             self.logger.error(f"Error completing: {e}")
             raise e
 
-        return LanguageModelResponse(**response)
-
     def stream_complete(
         self,
         messages: LanguageModelMessages,
         model_name: LanguageModelName,
         content_chunks: list[ContentChunk] = [],
         debug_info: dict = {},
-        temperature: float = _DEFAULT_COMPLETE_TEMPERATURE,
-        timeout: int = _DEFAULT_COMPLETE_TIMEOUT,
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: Optional[list[LanguageModelTool]] = None,
         start_text: Optional[str] = None,
     ):
@@ -139,27 +139,44 @@ class LanguageModelService:
         Returns:
             The LanguageModelStreamResponse object once the stream has finished.
         """
-        return self._trigger_stream_complete(
-            messages=messages,
-            content_chunks=content_chunks,
-            model_name=model_name,
-            debug_info=debug_info,
-            timeout=timeout,
-            temperature=temperature,
-            tools=tools,
-            start_text=start_text,
-        )
+        options = self._add_tools_to_options({}, tools)
+        search_context = self._to_search_context(content_chunks)
+        messages = messages.model_dump(exclude_none=True)
 
-    @to_async
-    @async_warning
-    def async_stream_complete(
+        try:
+            response = unique_sdk.Integrated.chat_stream_completion(
+                user_id=self.state.user_id,
+                company_id=self.state.company_id,
+                assistantMessageId=self.state.assistant_message_id,  # type: ignore
+                userMessageId=self.state.user_message_id,  # type: ignore
+                messages=cast(
+                    list[unique_sdk.Integrated.ChatCompletionRequestMessage],
+                    messages,
+                ),
+                chatId=self.state.chat_id,
+                searchContext=search_context,
+                # TODO change or extend types in unique_sdk
+                model=model_name.name,  # type: ignore
+                timeout=timeout,
+                temperature=temperature,
+                assistantId=self.state.assistant_id,
+                debugInfo=debug_info,
+                options=options,  # type: ignore
+                startText=start_text,
+            )
+            return LanguageModelStreamResponse(**response)
+        except Exception as e:
+            self.logger.error(f"Error streaming completion: {e}")
+            raise e
+
+    async def stream_complete_async(
         self,
         messages: LanguageModelMessages,
         model_name: LanguageModelName,
         content_chunks: list[ContentChunk] = [],
         debug_info: dict = {},
-        temperature: float = _DEFAULT_COMPLETE_TEMPERATURE,
-        timeout: int = _DEFAULT_COMPLETE_TIMEOUT,
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: Optional[list[LanguageModelTool]] = None,
         start_text: Optional[str] = None,
     ):
@@ -179,48 +196,13 @@ class LanguageModelService:
         Returns:
             The LanguageModelStreamResponse object once the stream has finished.
         """
-        return self._trigger_stream_complete(
-            messages=messages,
-            content_chunks=content_chunks,
-            model_name=model_name,
-            debug_info=debug_info,
-            timeout=timeout,
-            temperature=temperature,
-            tools=tools,
-            start_text=start_text,
-        )
 
-    def _trigger_stream_complete(
-        self,
-        messages: LanguageModelMessages,
-        model_name: LanguageModelName,
-        content_chunks: list[ContentChunk],
-        debug_info: dict,
-        timeout: int,
-        temperature: float,
-        tools: Optional[list[LanguageModelTool]] = None,
-        start_text: Optional[str] = None,
-    ) -> LanguageModelStreamResponse:
         options = self._add_tools_to_options({}, tools)
-        search_context = [
-            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 content_chunks
-        ]
-
-        messages = messages.model_dump(exclude_none=True, exclude={"tool_calls"})
+        search_context = self._to_search_context(content_chunks)
+        messages = messages.model_dump(exclude_none=True, exclude=["tool_calls"])
 
         try:
-            response = unique_sdk.Integrated.chat_stream_completion(
+            response = await unique_sdk.Integrated.chat_stream_completion_async(
                 user_id=self.state.user_id,
                 company_id=self.state.company_id,
                 assistantMessageId=self.state.assistant_message_id,  # type: ignore
@@ -240,11 +222,29 @@ class LanguageModelService:
                 options=options,  # type: ignore
                 startText=start_text,
             )
+            return LanguageModelStreamResponse(**response)
         except Exception as e:
             self.logger.error(f"Error streaming completion: {e}")
             raise e
 
-        return LanguageModelStreamResponse(**response)
+    @staticmethod
+    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
+        ]
 
     @staticmethod
     def _add_tools_to_options(

{unique_toolkit-0.5.6.dist-info → unique_toolkit-0.5.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 0.5.6
+Version: 0.5.7
 Summary: 
 License: MIT
 Author: Martin Fadler
@@ -15,10 +15,9 @@ Requires-Dist: pydantic (>=2.8.2,<3.0.0)
 Requires-Dist: pyhumps (>=3.8.0,<4.0.0)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: regex (>=2024.5.15,<2025.0.0)
-Requires-Dist: requests (>=2.32.3,<3.0.0)
 Requires-Dist: tiktoken (>=0.7.0,<0.8.0)
 Requires-Dist: typing-extensions (>=4.9.0,<5.0.0)
-Requires-Dist: unique-sdk (>=0.9.2,<0.10.0)
+Requires-Dist: unique-sdk (>=0.9.4,<0.10.0)
 Description-Content-Type: text/markdown
 
 # Unique Toolkit
@@ -102,6 +101,12 @@ 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.7] - 2024-07-31
+- Replace mocked async service calls with async calls in `unique_sdk` 
+- Change async methods name from `async_*` to `*_async`
+- Remove `chat_only` and `scope_ids` attributes from `ChatState` class
+- Replace `AsyncExecutor` by simpler utility function `run_async_tasks_parallel`
+
 ## [0.5.6] - 2024-07-30
 - Bug fix: `ContentService.search_content_chunks` and it's `async` equivalent now accept `None` as a valid parameter value for `scope_ids`.