aisberg 0.1.0__py3-none-any.whl

aisberg/modules/embeddings.py

@@ -0,0 +1,309 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Union, List, Literal
+
+from ..models.embeddings import (
+    EncodingResponse,
+    ChunksDataList,
+    RerankerResponse,
+    ChunkData,
+)
+from ..models.collections import Collection
+from ..api import async_endpoints, endpoints
+from ..abstract.modules import AsyncModule, SyncModule
+
+
+class AbstractEmbeddingsModule(ABC):
+    """
+    Abstract base class for embeddings modules.
+    Handles common logic for embedding operations across synchronous and asynchronous modules.
+    """
+
+    def __init__(self, parent, http_client):
+        """
+        Initialize the AbstractEmbeddingsModule.
+
+        Args:
+            parent: Parent client instance.
+            http_client: HTTP client for making requests.
+        """
+        self._parent = parent
+        self._client = http_client
+
+    @abstractmethod
+    def encode(
+        self,
+        input: str,
+        model: str,
+        encoding_format: Optional[Literal["float", "base64"]] = "float",
+        normalize: Optional[bool] = False,
+        **kwargs,
+    ) -> EncodingResponse:
+        """Encode a list of texts into embeddings.
+
+        Args:
+            input (str): The text or texts to encode. Can be a single string or a list of strings.
+            model (str): The model to use for encoding. Defaults to "text-embedding-3-small".
+            encoding_format (str): The format of the encoding. Defaults to "float". Can be "float" or "base64".
+            normalize (bool): Whether to normalize the embeddings. Defaults to False.
+            **kwargs: Additional parameters for the encoding.
+
+        Returns:
+            EncodingResponse: The response containing the encoded embeddings.
+        """
+        pass
+
+    @abstractmethod
+    def retrieve(
+        self,
+        query: str,
+        collections_names: List[Union[str, Collection]],
+        limit: int = 10,
+        score_threshold: float = 0.0,
+        filters: List = None,
+        beta: float = 0.7,
+    ) -> ChunksDataList:
+        """Retrieve similar texts based on a query.
+
+        Args:
+            query (str): The query text to retrieve similar texts for.
+            collections_names (List[str]): A list of collection names to search in.
+            limit (int): The maximum number of results to return. Defaults to 10.
+            score_threshold (float): The minimum score threshold for results. Defaults to 0.0.
+            filters (list): A list of filters to apply to the retrieval.
+            beta (float): Dense/Sparse trade-off parameter. Defaults to 0.7. 0 means full sparse, 1 means full dense.
+
+        Returns:
+            List[ChunkData]: A list of ChunkData objects containing the retrieved texts and their metadata.
+        """
+        pass
+
+    @abstractmethod
+    def rerank(
+        self,
+        query: str,
+        documents: Union[ChunksDataList, List[Union[str, ChunkData]]],
+        model: str,
+        top_n: int = 10,
+        return_documents: bool = True,
+        threshold: Optional[float] = None,
+    ) -> RerankerResponse:
+        """Rerank texts based on a query.
+
+        Args:
+            query (str): The query text to rerank the documents against.
+            documents (Union[ChunksDataList, List[Union[str, ChunkData]]]): A list of documents to rerank. Can be a ChunksDataList or a list of strings or ChunkData objects.
+            model (str): The model to use for reranking. Defaults to "text-embedding-3-small".
+            top_n (int): The number of top results to return. Defaults to 10.
+            return_documents (bool): Whether to return the original documents in the response. Defaults to True.
+            threshold (Optional[float]): A threshold for filtering results. If provided, only results with a score above this threshold will be returned. Defaults to None.
+
+        Returns:
+            RerankerResponse: The response containing the reranked documents and their scores.
+
+        Raises:
+            ValueError: If the documents list is empty or contains invalid document types.
+            Exception: If the documents list is not of the expected type.
+        """
+        pass
+
+    @staticmethod
+    def _format_collections_names(
+        collections_names: List[Union[str, Collection]],
+    ) -> List[str]:
+        """Format the input collections names into a list of strings."""
+        coll_names = []
+
+        for coll in collections_names:
+            if isinstance(coll, Collection):
+                coll_names.append(coll.name)
+            elif isinstance(coll, str):
+                coll_names.append(coll)
+            else:
+                raise ValueError(
+                    f"Invalid collection type: {type(coll)}. Expected str or Collection."
+                )
+
+        return coll_names
+
+    @staticmethod
+    def _format_chunks_data_list(
+        documents: Union[ChunksDataList, List[Union[str, ChunkData]]],
+    ) -> List[str]:
+        """Format the input documents into a ChunksDataList."""
+        chunks = []
+
+        if isinstance(documents, ChunksDataList):
+            chunks = documents.texts()
+        elif isinstance(documents, list):
+            if len(documents) == 0:
+                raise ValueError("Documents list is empty.")
+
+            for doc in documents:
+                if isinstance(doc, ChunkData):
+                    chunks.append(doc.text)
+                elif isinstance(doc, str):
+                    chunks.append(doc)
+                else:
+                    raise ValueError(
+                        f"Invalid document type: {type(doc)}. Expected str or ChunkData."
+                    )
+        else:
+            raise Exception(
+                f"Documents list is not of the expected type: {type(documents)}. Expected ChunksDataList or list of str or ChunkData."
+            )
+
+        return chunks
+
+
+class SyncEmbeddingsModule(AbstractEmbeddingsModule, SyncModule):
+    """
+    `SyncEmbeddingsModule` is a synchronous module that provides a high-level interface for interacting with
+    embeddings tools. The module abstracts all communication with the backend API,
+    providing both blocking and generator-based usage.
+    """
+
+    def __init__(self, parent, http_client):
+        super().__init__(parent, http_client)
+        SyncModule.__init__(self, parent, http_client)
+
+    def encode(
+        self,
+        input: str,
+        model: str,
+        encoding_format: Optional[Literal["float", "base64"]] = "float",
+        normalize: Optional[bool] = False,
+        **kwargs,
+    ) -> EncodingResponse:
+        resp = endpoints.embeddings(
+            self._client,
+            input=input,
+            model=model,
+            encoding_format=encoding_format,
+            normalize=normalize,
+            **kwargs,
+        )
+        return EncodingResponse.model_validate(resp)
+
+    def retrieve(
+        self,
+        query: str,
+        collections_names: List[Union[str, Collection]],
+        limit: int = 10,
+        score_threshold: float = 0.0,
+        filters: List = None,
+        beta: float = 0.7,
+    ) -> ChunksDataList:
+        if filters is None:
+            filters = []
+
+        resp = endpoints.retrieve(
+            self._client,
+            query=query,
+            collections_names=self._format_collections_names(collections_names),
+            limit=limit,
+            score_threshold=score_threshold,
+            filters=filters,
+            beta=beta,
+        )
+        return ChunksDataList.model_validate(resp)
+
+    def rerank(
+        self,
+        query: str,
+        documents: Union[ChunksDataList, List[Union[str, ChunkData]]],
+        model: str,
+        top_n: int = 10,
+        return_documents: bool = True,
+        threshold: Optional[float] = None,
+    ) -> RerankerResponse:
+        resp = endpoints.rerank(
+            self._client,
+            query,
+            self._format_chunks_data_list(documents),
+            model,
+            top_n,
+            return_documents,
+        )
+        resp = RerankerResponse.model_validate(resp)
+
+        if threshold is not None:
+            resp = resp.filter_by_relevance_score(threshold)
+
+        return resp
+
+
+class AsyncEmbeddingsModule(AbstractEmbeddingsModule, AsyncModule):
+    """
+    `AsyncEmbeddingsModule` is an asynchronous module that provides a high-level interface for interacting with
+    embeddings tools. The module abstracts all communication with the backend API,
+    providing both blocking and generator-based usage.
+    """
+
+    def __init__(self, parent, http_client):
+        super().__init__(parent, http_client)
+        AsyncModule.__init__(self, parent, http_client)
+
+    async def encode(
+        self,
+        input: str,
+        model: str,
+        encoding_format: Optional[Literal["float", "base64"]] = "float",
+        normalize: Optional[bool] = False,
+        **kwargs,
+    ) -> EncodingResponse:
+        resp = await async_endpoints.embeddings(
+            self._client,
+            input=input,
+            model=model,
+            encoding_format=encoding_format,
+            normalize=normalize,
+            **kwargs,
+        )
+        return EncodingResponse.model_validate(resp)
+
+    async def retrieve(
+        self,
+        query: str,
+        collections_names: List[Union[str, Collection]],
+        limit: int = 10,
+        score_threshold: float = 0.0,
+        filters: List = None,
+        beta: float = 0.7,
+    ) -> ChunksDataList:
+        if filters is None:
+            filters = []
+
+        resp = await async_endpoints.retrieve(
+            self._client,
+            query=query,
+            collections_names=self._format_collections_names(collections_names),
+            limit=limit,
+            score_threshold=score_threshold,
+            filters=filters,
+            beta=beta,
+        )
+        return ChunksDataList.model_validate(resp)
+
+    async def rerank(
+        self,
+        query: str,
+        documents: Union[ChunksDataList, List[Union[str, ChunkData]]],
+        model: str,
+        top_n: int = 10,
+        return_documents: bool = True,
+        threshold: Optional[float] = None,
+    ) -> RerankerResponse:
+        resp = await async_endpoints.rerank(
+            self._client,
+            query,
+            self._format_chunks_data_list(documents),
+            model,
+            top_n,
+            return_documents,
+        )
+        resp = RerankerResponse.model_validate(resp)
+
+        if threshold is not None:
+            resp = resp.filter_by_relevance_score(threshold)
+
+        return resp

aisberg/modules/me.py

@@ -0,0 +1,77 @@
+from typing import List
+from abc import ABC
+from ..models.token import TokenInfo
+
+from abc import abstractmethod
+from ..abstract.modules import SyncModule, AsyncModule
+from ..api import endpoints, async_endpoints
+
+
+class AbstractMeModule(ABC):
+    def __init__(self, parent, client):
+        self._parent = parent
+        self._client = client
+
+    @abstractmethod
+    def _fetch_info(self) -> TokenInfo:
+        """
+        Get information about the current API token.
+
+        Returns:
+            TokenInfo: Information about the API token.
+
+        Raises:
+            Exception: If there is an error fetching the token information.
+        """
+        pass
+
+    def info(self) -> TokenInfo:
+        """
+        Get information about the current API token.
+
+        Returns:
+            TokenInfo: Information about the API token.
+
+        Raises:
+            Exception: If there is an error fetching the token information.
+        """
+        return self._fetch_info()
+
+    def groups(self) -> List[str]:
+        """
+        Get a list of groups the current user belongs to.
+
+        Returns:
+            list[str]: A list of group IDs.
+
+        Raises:
+            Exception: If there is an error fetching the groups.
+        """
+        resp = self.info()
+        return getattr(resp, "groups", []) or []
+
+
+class SyncMeModule(SyncModule, AbstractMeModule):
+    def __init__(self, parent, client):
+        SyncModule.__init__(self, parent, client)
+        AbstractMeModule.__init__(self, parent, client)
+
+    def _fetch_info(self) -> TokenInfo:
+        return endpoints.me(self._client)
+
+
+class AsyncMeModule(AsyncModule, AbstractMeModule):
+    def __init__(self, parent, client):
+        AsyncModule.__init__(self, parent, client)
+        AbstractMeModule.__init__(self, parent, client)
+
+    async def _fetch_info(self) -> TokenInfo:
+        resp = await async_endpoints.me(self._client)
+        return TokenInfo.model_validate(resp)
+
+    async def info(self) -> TokenInfo:
+        return await self._fetch_info()
+
+    async def groups(self) -> List[str]:
+        resp = await self.info()
+        return getattr(resp, "groups", []) or []

aisberg/modules/models.py

@@ -0,0 +1,108 @@
+from typing import List
+from abc import ABC
+
+from ..models.models import Model
+
+from abc import abstractmethod
+from ..abstract.modules import SyncModule, AsyncModule
+from ..api import endpoints, async_endpoints
+
+
+class AbstractModelsModule(ABC):
+    def __init__(self, parent, client):
+        self._parent = parent
+        self._client = client
+
+    @abstractmethod
+    def list(self) -> List[Model]:
+        """
+        Get a list of available collections. Models are grouped by your belonging groups.
+
+        Returns:
+            List[GroupModels]: A list of available collections.
+
+        Raises:
+            ValueError: If no collections are found.
+            Exception: If there is an error fetching the collections.
+        """
+        pass
+
+    @abstractmethod
+    def get(self, model_id: str) -> Model:
+        """Get details of a specific model.
+
+        Args:
+            model_id (str): The ID of the model to retrieve.
+
+        Returns:
+            Model: The details of the specified model.
+
+        Raises:
+            ValueError: If the specified model is not found.
+        """
+        pass
+
+    @abstractmethod
+    def is_available(self, model_id: str) -> bool:
+        """Check if a specific model is available.
+
+        Args:
+            model_id (str): The ID of the model to check.
+
+        Returns:
+            bool: True if the model is available, False otherwise.
+        """
+        pass
+
+    @staticmethod
+    def _get_model_by_id(models: List[Model], model_id: str) -> Model | None:
+        for model in models:
+            if model.id == model_id:
+                return model
+        return None
+
+
+class SyncModelsModule(SyncModule, AbstractModelsModule):
+    def __init__(self, parent, client):
+        SyncModule.__init__(self, parent, client)
+        AbstractModelsModule.__init__(self, parent, client)
+
+    def list(self) -> List[Model]:
+        return endpoints.models(self._client)
+
+    def get(self, model_id: str) -> Model:
+        models = self.list()
+        model = self._get_model_by_id(models, model_id)
+        if model is None:
+            raise ValueError("No model found")
+        return model
+
+    def is_available(self, model_id: str) -> bool:
+        try:
+            self.get(model_id)
+            return True
+        except ValueError:
+            return False
+
+
+class AsyncModelsModule(AsyncModule, AbstractModelsModule):
+    def __init__(self, parent, client):
+        AsyncModule.__init__(self, parent, client)
+        AbstractModelsModule.__init__(self, parent, client)
+
+    async def list(self) -> List[Model]:
+        return await async_endpoints.models(self._client)
+
+    async def get(self, model_id: str) -> Model:
+        models = await self.list()
+        model = self._get_model_by_id(models, model_id)
+        if model is None:
+            raise ValueError("No model found")
+        return model
+
+    async def is_available(self, model_id: str) -> bool:
+        try:
+            await self.get(model_id)
+            return True
+        except ValueError:
+            return False

aisberg/modules/tools.py

@@ -0,0 +1,78 @@
+from ..abstract.modules import BaseModule
+from ..exceptions import ToolExecutionError
+from typing import Callable, Dict, Any
+
+
+class ToolsModule(BaseModule):
+    """
+    Tools module for the aisberg application.
+    This module provides methods to register and execute tools.
+    """
+
+    def __init__(self, parent):
+        """
+        Initialize the ToolsModule.
+
+        Args:
+            parent (AisbergClient): Parent client instance.
+        """
+        super().__init__(parent)
+
+    def register(self, name: str, func: Callable) -> None:
+        """
+        Enregistre une fonction comme tool disponible.
+
+        Args:
+            name (str): Nom du tool (doit correspondre au nom dans la définition du tool)
+            func (Callable): Fonction à exécuter quand le tool est appelé
+        """
+        self._parent.tool_registry[name] = func
+
+    def execute(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """
+        Exécute un tool avec les arguments fournis.
+
+        Args:
+            tool_name (str): Nom du tool à exécuter
+            arguments (Dict[str, Any]): Arguments à passer au tool
+
+        Returns:
+            Any: Résultat de l'exécution du tool
+
+        Raises:
+            ToolExecutionError: Si le tool n'existe pas ou si l'exécution échoue
+        """
+        if tool_name not in self._parent.tool_registry:
+            raise ToolExecutionError(f"Tool '{tool_name}' not registered")
+
+        try:
+            return self._parent.tool_registry[tool_name](**arguments)
+        except Exception as e:
+            raise ToolExecutionError(f"Error executing tool '{tool_name}': {str(e)}")
+
+    def list(self) -> Dict[str, Callable]:
+        """
+        Liste tous les tools enregistrés.
+
+        Returns:
+            Dict[str, Callable]: Dictionnaire des tools enregistrés avec leur nom et fonction
+        """
+        return self._parent.tool_registry
+
+    def clear(self) -> None:
+        """
+        Efface tous les tools enregistrés.
+        """
+        self._parent.tool_registry.clear()
+
+    def remove(self, tool_name: str) -> None:
+        """
+        Supprime un tool enregistré.
+
+        Args:
+            tool_name (str): Nom du tool à supprimer
+        """
+        if tool_name in self._parent.tool_registry:
+            del self._parent.tool_registry[tool_name]
+        else:
+            raise ToolExecutionError(f"Tool '{tool_name}' not registered")