unique_toolkit 0.0.2__py3-none-any.whl → 0.5.1__py3-none-any.whl

unique_toolkit/content/service.py

@@ -0,0 +1,356 @@
+import logging
+import os
+import tempfile
+from pathlib import Path
+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.chat.state import ChatState
+from unique_toolkit.content.schemas import (
+    Content,
+    ContentChunk,
+    ContentSearchType,
+    ContentUploadInput,
+)
+
+
+class ContentService:
+    def __init__(self, state: ChatState, logger: Optional[logging.Logger] = None):
+        self.state = state
+        self.logger = logger or logging.getLogger(__name__)
+
+    def search_content_chunks(
+        self,
+        search_string: str,
+        search_type: ContentSearchType,
+        limit: int,
+        scope_ids: Optional[list[str]] = None,
+    ) -> list[ContentChunk]:
+        """
+        Performs a synchronous search for content chunks in the knowledge base.
+
+        Args:
+            search_string (str): The search string.
+            search_type (ContentSearchType): The type of search to perform.
+            limit (int): The maximum number of results to return.
+            scope_ids (Optional[list[str]]): The scope IDs. Defaults to None.
+
+        Returns:
+            list[ContentChunk]: The search results.
+        """
+        return self._trigger_search_content_chunks(
+            search_string,
+            search_type,
+            limit,
+            scope_ids,
+        )
+
+    @to_async
+    @async_warning
+    def async_search_content_chunks(
+        self,
+        search_string: str,
+        search_type: ContentSearchType,
+        limit: int,
+        scope_ids: Optional[list[str]],
+    ):
+        """
+        Performs an asynchronous search for content chunks in the knowledge base.
+
+        Args:
+            search_string (str): The search string.
+            search_type (ContentSearchType): The type of search to perform.
+            limit (int): The maximum number of results to return.
+            scope_ids (Optional[list[str]]): The scope IDs. Defaults to [].
+
+        Returns:
+            list[ContentChunk]: The search results.
+        """
+        return self._trigger_search_content_chunks(
+            search_string,
+            search_type,
+            limit,
+            scope_ids,
+        )
+
+    def _trigger_search_content_chunks(
+        self,
+        search_string: str,
+        search_type: ContentSearchType,
+        limit: int,
+        scope_ids: Optional[list[str]],
+    ) -> list[ContentChunk]:
+        scope_ids = scope_ids or self.state.scope_ids or []
+
+        if not scope_ids:
+            self.logger.warning("No scope IDs provided for search.")
+
+        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,
+                chatOnly=self.state.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)
+
+    def search_contents(
+        self,
+        where: dict,
+    ) -> list[Content]:
+        """
+        Performs a search in the knowledge base by filter (and not a smilarity search)
+        This function loads complete content of the files from the knowledge base in contrast to search_content_chunks.
+
+        Args:
+            where (dict): The search criteria.
+
+        Returns:
+            list[Content]: The search results.
+        """
+        return self._trigger_search_contents(where)
+
+    @to_async
+    @async_warning
+    def async_search_contents(
+        self,
+        where: dict,
+    ) -> list[Content]:
+        """
+        Performs an asynchronous search for content files in the knowledge base by filter.
+
+        Args:
+            where (dict): The search criteria.
+
+        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(
+                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
+
+        return map_contents(contents)
+
+    def upload_content(
+        self,
+        path_to_content: str,
+        content_name: str,
+        mime_type: str,
+        scope_id: Optional[str] = None,
+        chat_id: Optional[str] = None,
+    ):
+        """
+        Uploads content to the knowledge base.
+
+        Args:
+            path_to_content (str): The path to the content to upload.
+            content_name (str): The name of the content.
+            mime_type (str): The MIME type of the content.
+            scope_id (Optional[str]): The scope ID. Defaults to None.
+            chat_id (Optional[str]): The chat ID. Defaults to None.
+
+        Returns:
+            Content: The uploaded content.
+        """
+
+        byte_size = os.path.getsize(path_to_content)
+        created_content = self._trigger_upsert_content(
+            input=ContentUploadInput(
+                key=content_name, title=content_name, mime_type=mime_type
+            ),
+            scope_id=scope_id,
+            chat_id=chat_id,
+        )
+
+        write_url = created_content.write_url
+
+        if not write_url:
+            error_msg = "Write url for uploaded content is missing"
+            self.logger.error(error_msg)
+            raise ValueError(error_msg)
+
+        # upload to azure blob storage SAS url uploadUrl the pdf file translatedFile make sure it is treated as a application/pdf
+        with open(path_to_content, "rb") as file:
+            requests.put(
+                url=write_url,
+                data=file,
+                headers={
+                    "X-Ms-Blob-Content-Type": mime_type,
+                    "X-Ms-Blob-Type": "BlockBlob",
+                },
+            )
+
+        read_url = created_content.read_url
+
+        if not read_url:
+            error_msg = "Read url for uploaded content is missing"
+            self.logger.error(error_msg)
+            raise ValueError(error_msg)
+
+        if chat_id:
+            self._trigger_upsert_content(
+                input=ContentUploadInput(
+                    key=content_name,
+                    title=content_name,
+                    mime_type=mime_type,
+                    byte_size=byte_size,
+                ),
+                content_url=read_url,
+                chat_id=chat_id,
+            )
+        else:
+            self._trigger_upsert_content(
+                input=ContentUploadInput(
+                    key=content_name,
+                    title=content_name,
+                    mime_type=mime_type,
+                    byte_size=byte_size,
+                ),
+                content_url=read_url,
+                scope_id=scope_id,
+            )
+
+        return created_content
+
+    def _trigger_upsert_content(
+        self,
+        input: ContentUploadInput,
+        scope_id: Optional[str] = None,
+        chat_id: Optional[str] = None,
+        content_url: Optional[str] = None,
+    ):
+        if not chat_id and not scope_id:
+            raise ValueError("chat_id or scope_id must be provided")
+
+        try:
+            if input.byte_size:
+                input_json = {
+                    "key": input.key,
+                    "title": input.title,
+                    "mimeType": input.mime_type,
+                    "byteSize": input.byte_size,
+                }
+            else:
+                input_json = {
+                    "key": input.key,
+                    "title": input.title,
+                    "mimeType": input.mime_type,
+                }
+            content = unique_sdk.Content.upsert(
+                user_id=self.state.user_id,
+                company_id=self.state.company_id,
+                input=input_json,  # type: ignore
+                fileUrl=content_url,
+                scopeId=scope_id,
+                chatId=chat_id,
+                sourceOwnerType=None,  # type: ignore
+                storeInternally=False,
+            )
+            return Content(**content)
+        except Exception as e:
+            self.logger.error(f"Error while uploading content: {e}")
+            raise e
+
+    def download_content(
+        self,
+        content_id: str,
+        content_name: str,
+        chat_id: Optional[str] = None,
+    ) -> Path:
+        """
+        Downloads content to temporary directory
+
+        Args:
+            content_id (str): The id of the uploaded content.
+            content_name (str): The name of the uploaded content.
+            chat_id (Optional[str]): The chat_id, defaults to None.
+
+        Returns:
+            content_path: The path to the downloaded content in the temporary directory.
+
+        Raises:
+            Exception: If the download fails.
+        """
+
+        print("download chat id", chat_id)
+
+        url = f"{unique_sdk.api_base}/content/{content_id}/file"
+        if chat_id:
+            url = f"{url}?chatId={chat_id}"
+
+        # Create a random directory inside /tmp
+        random_dir = tempfile.mkdtemp(dir="/tmp")
+
+        # Create the full file path
+        content_path = Path(random_dir) / content_name
+
+        # Download the file and save it to the random directory
+        headers = {
+            "x-api-version": unique_sdk.api_version,
+            "x-app-id": unique_sdk.app_id,
+            "x-user-id": self.state.user_id,
+            "x-company-id": self.state.company_id,
+            "Authorization": "Bearer %s" % (unique_sdk.api_key,),
+        }
+
+        response = requests.get(url, headers=headers)
+        if response.status_code == 200:
+            with open(content_path, "wb") as file:
+                file.write(response.content)
+        else:
+            error_msg = f"Error downloading file: Status code {response.status_code}"
+            self.logger.error(error_msg)
+            raise Exception(error_msg)
+
+        return content_path

unique_toolkit/content/utils.py

@@ -0,0 +1,188 @@
+import re
+
+import tiktoken
+
+from unique_toolkit.content.schemas import (
+    ContentChunk,
+)
+
+
+def _map_content_id_to_chunks(content_chunks: list[ContentChunk]):
+    doc_id_to_chunks: dict[str, list[ContentChunk]] = {}
+    for chunk in content_chunks:
+        source_chunks = doc_id_to_chunks.get(chunk.id)
+        if not source_chunks:
+            doc_id_to_chunks[chunk.id] = [chunk]
+        else:
+            source_chunks.append(chunk)
+    return doc_id_to_chunks
+
+
+def sort_content_chunks(content_chunks: list[ContentChunk]):
+    """
+    Sorts the content chunks based on their 'order' in the original content.
+    This function sorts the search results based on their 'order' in ascending order.
+    It also performs text modifications by replacing the string within the tags <|/content|>
+    with 'text part {order}' and removing any <|info|> tags (Which is useful in referencing the chunk).
+    Parameters:
+    - content_chunks (list): A list of ContentChunkt objects.
+    Returns:
+    - list: A list of ContentChunk objects sorted according to their order.
+    """
+    doc_id_to_chunks = _map_content_id_to_chunks(content_chunks)
+    sorted_chunks: list[ContentChunk] = []
+    for chunks in doc_id_to_chunks.values():
+        chunks.sort(key=lambda x: x.order)
+        for i, s in enumerate(chunks):
+            s.text = re.sub(
+                r"<\|/content\|>", f" text part {s.order}<|/content|>", s.text
+            )
+            s.text = re.sub(r"<\|info\|>(.*?)<\|\/info\|>", "", s.text)
+            pages_postfix = _generate_pages_postfix([s])
+            s.key = s.key + pages_postfix if s.key else s.key
+            s.title = s.title + pages_postfix if s.title else s.title
+        sorted_chunks.extend(chunks)
+    return sorted_chunks
+
+
+def merge_content_chunks(content_chunks: list[ContentChunk]):
+    """
+    Merges multiple search results based on their 'id', removing redundant content and info markers.
+
+    This function groups search results by their 'id' and then concatenates their texts,
+    cleaning up any content or info markers in subsequent chunks beyond the first one.
+
+    Parameters:
+    - content_chunks (list): A list of objects, each representing a search result with 'id' and 'text' keys.
+
+    Returns:
+    - list: A list of objects with merged texts for each unique 'id'.
+    """
+
+    doc_id_to_chunks = _map_content_id_to_chunks(content_chunks)
+    merged_chunks: list[ContentChunk] = []
+    for chunks in doc_id_to_chunks.values():
+        chunks.sort(key=lambda x: x.order)
+        for i, s in enumerate(chunks):
+            ## skip first element
+            if i > 0:
+                ## replace the string within the tags <|content|>...<|/content|> and <|info|> and <|/info|>
+                s.text = re.sub(r"<\|content\|>(.*?)<\|\/content\|>", "", s.text)
+                s.text = re.sub(r"<\|info\|>(.*?)<\|\/info\|>", "", s.text)
+
+        pages_postfix = _generate_pages_postfix(chunks)
+        chunks[0].text = "\n".join(str(s.text) for s in chunks)
+        chunks[0].key = (
+            chunks[0].key + pages_postfix if chunks[0].key else chunks[0].key
+        )
+        chunks[0].title = (
+            chunks[0].title + pages_postfix if chunks[0].title else chunks[0].title
+        )
+        chunks[0].end_page = chunks[-1].end_page
+        merged_chunks.append(chunks[0])
+
+    return merged_chunks
+
+
+def _generate_pages_postfix(chunks: list[ContentChunk]) -> str:
+    """
+    Generates a postfix string of page numbers from a list of source objects.
+    Each source object contains startPage and endPage numbers. The function
+    compiles a list of all unique page numbers greater than 0 from all chunks,
+    and then returns them as a string prefixed with " : " if there are any.
+
+    Parameters:
+    - chunks (list): A list of objects with 'startPage' and 'endPage' keys.
+
+    Returns:
+    - string: A string of page numbers separated by commas, prefixed with " : ".
+    """
+
+    def gen_all_numbers_in_between(start, end) -> list[int]:
+        """
+        Generates a list of all numbers between start and end, inclusive.
+        If start or end is -1, it behaves as follows:
+        - If both start and end are -1, it returns an empty list.
+        - If only end is -1, it returns a list containing only the start.
+        - If start is -1, it returns an empty list.
+
+        Parameters:
+        - start (int): The starting page number.
+        - end (int): The ending page number.
+
+        Returns:
+        - list: A list of numbers from start to end, inclusive.
+        """
+        if start == -1 and end == -1:
+            return []
+        if end == -1:
+            return [start]
+        if start == -1:
+            return []
+        return list(range(start, end + 1))
+
+    page_numbers_array = [
+        gen_all_numbers_in_between(s.start_page, s.end_page) for s in chunks
+    ]
+    page_numbers = [number for sublist in page_numbers_array for number in sublist]
+    page_numbers = [p for p in page_numbers if p > 0]
+    page_numbers = sorted(set(page_numbers))
+    pages_postfix = (
+        " : " + ",".join(str(p) for p in page_numbers) if page_numbers else ""
+    )
+    return pages_postfix
+
+
+def pick_content_chunks_for_token_window(
+    content_chunks: list[ContentChunk],
+    token_limit: int,
+    encoding_model="cl100k_base",
+):
+    """
+    Selects and returns a list of search results that fit within a specified token limit.
+
+    This function iterates over a list of search results, each with a 'text' field, and
+    encodes the text using a predefined encoding scheme. It accumulates search results
+    until the token limit is reached or exceeded.
+
+    Parameters:
+    - content_chunks (list): A list of dictionaries, each containing a 'text' key with string value.
+    - token_limit (int): The maximum number of tokens to include in the output.
+
+    Returns:
+    - list: A list of dictionaries representing the search results that fit within the token limit.
+    """
+    picked_chunks: list[ContentChunk] = []
+    token_count = 0
+
+    encoding = tiktoken.get_encoding(encoding_model)
+
+    for chunk in content_chunks:
+        try:
+            searchtoken_count = len(encoding.encode(chunk.text))
+        except Exception:
+            searchtoken_count = 0
+        if token_count + searchtoken_count > token_limit:
+            break
+
+        picked_chunks.append(chunk)
+        token_count += searchtoken_count
+
+    return picked_chunks
+
+
+def count_tokens(text, encoding_model="cl100k_base") -> int:
+    """
+    Counts the number of tokens in the provided text.
+
+    This function encodes the input text using a predefined encoding scheme
+    and returns the number of tokens in the encoded text.
+
+    Parameters:
+    - text (str): The text to count tokens for.
+
+    Returns:
+    - int: The number of tokens in the text.
+    """
+    encoding = tiktoken.get_encoding(encoding_model)
+    return len(encoding.encode(text))

unique_toolkit/embedding/schemas.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class Embeddings(BaseModel):
+    embeddings: list[list[float]]

unique_toolkit/embedding/service.py

@@ -0,0 +1,89 @@
+import logging
+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.chat.state import ChatState
+from unique_toolkit.embedding.schemas import Embeddings
+
+
+class EmbeddingService:
+    def __init__(self, state: ChatState, logger: Optional[logging.Logger] = None):
+        self.state = state
+        self.logger = logger or logging.getLogger(__name__)
+
+    _DEFAULT_TIMEOUT = 600_000
+
+    def embed_texts(
+        self,
+        texts: list[str],
+        timeout: int = _DEFAULT_TIMEOUT,
+    ) -> Embeddings:
+        """
+        Embed text.
+
+        Args:
+            text (str): The text to embed.
+            timeout (int): The timeout in milliseconds. Defaults to None.
+
+        Returns:
+            Embeddings: The Embedding object.
+
+        Raises:
+            Exception: If an error occurs.
+        """
+        return self._trigger_embed_texts(
+            texts=texts,
+            timeout=timeout,
+        )
+
+    @to_async
+    @async_warning
+    def async_embed_texts(
+        self,
+        texts: list[str],
+        timeout: int = _DEFAULT_TIMEOUT,
+    ) -> Embeddings:
+        """
+        Embed text asynchronously.
+
+        Args:
+            text (str): The text to embed.
+            timeout (int): The timeout in milliseconds. Defaults to None.
+
+        Returns:
+            Embeddings: The Embedding object.
+
+        Raises:
+            Exception: If an error occurs.
+        """
+        return self._trigger_embed_texts(
+            texts=texts,
+            timeout=timeout,
+        )
+
+    def _trigger_embed_texts(self, texts: list[str], timeout: int) -> Embeddings:
+        request = {
+            "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,
+        embedding_1: list[float],
+        embedding_2: list[float],
+    ) -> float:
+        """Get cosine similarity."""
+        return np.dot(embedding_1, embedding_2) / (
+            np.linalg.norm(embedding_1) * np.linalg.norm(embedding_2)
+        )