aisberg 0.1.0__py3-none-any.whl

aisberg/api/endpoints.py

@@ -0,0 +1,328 @@
+from io import BytesIO
+
+import httpx
+from ..models.chat import (
+    LanguageModelInput,
+    format_messages,
+    ChatCompletionResponse,
+    ChatCompletionChunk,
+)
+from typing import Optional, Generator, Union, List, Any, Tuple
+
+from ..models.collections import GroupCollections, PointDetails
+from ..models.embeddings import (
+    EncodingFormat,
+    EncodingResponse,
+    ChunksDataList,
+    RerankerResponse,
+)
+from ..models.models import Model
+from ..models.token import TokenInfo
+from ..models.workflows import WorkflowDetails, Workflow
+from ..utils import parse_chat_line, WorkflowLineParser
+from ..requests.sync_requests import req, req_stream
+from ..models.requests import AnyDict, AnyList
+
+
+def models(client: httpx.Client) -> List[Model]:
+    """
+    Get the list of available models.
+    """
+    resp = req(client, "GET", "/v1/models", AnyDict).data
+    if not resp or not isinstance(resp, list):
+        raise ValueError("Invalid response format for models")
+    return [Model.model_validate(item) for item in resp]
+
+
+def workflows(client: httpx.Client) -> List[Workflow]:
+    """
+    Get the list of available workflows.
+    """
+    resp = req(client, "GET", "/workflow/light", AnyList)
+    return [Workflow.model_validate(item) for item in resp.root]
+
+
+def workflow(client: httpx.Client, workflow_id: str) -> WorkflowDetails:
+    """
+    Get details of a specific workflow.
+    """
+    try:
+        resp = req(client, "GET", f"/workflow/details/{workflow_id}", WorkflowDetails)
+        return resp
+    except httpx.HTTPStatusError as e:
+        if e.response.status_code == 404:
+            raise ValueError(f"Workflow with ID {workflow_id} not found")
+        raise e
+
+
+def collections(client: httpx.Client) -> List[GroupCollections]:
+    """
+    Get the list of available collections.
+    """
+    resp = req(client, "GET", "/collections", AnyList)
+    return [GroupCollections.model_validate(item) for item in resp.root]
+
+
+def collection(
+    client: httpx.Client, collection_id: str, group_id: str
+) -> List[PointDetails]:
+    """
+    Get details of a specific collection.
+    """
+    try:
+        resp = req(client, "GET", f"/collections/{collection_id}/{group_id}", AnyList)
+        return [PointDetails.model_validate(item) for item in resp.root]
+    except httpx.HTTPStatusError as e:
+        if e.response.status_code == 404:
+            raise ValueError(
+                f"Collection with ID {collection_id} not found in group {group_id}"
+            )
+        raise e
+
+
+def me(client: httpx.Client) -> TokenInfo:
+    """
+    Get the details of the current user.
+    """
+    return req(client, "GET", "/users/me", TokenInfo)
+
+
+def chat(
+    client: httpx.Client,
+    input: LanguageModelInput,
+    model: str = None,
+    temperature: float = 0.7,
+    tools: Optional[list] = None,
+    group: Optional[str] = None,
+    **kwargs,
+) -> ChatCompletionResponse:
+    """
+    Send a chat message and get a response from an LLM endpoint.
+    """
+    if model is None:
+        raise ValueError("Model must be specified")
+
+    formatted_messages = format_messages(input)
+
+    payload = {
+        "model": model,
+        "messages": formatted_messages,
+        "temperature": temperature,
+        "stream": False,
+        **kwargs,
+    }
+
+    if group is not None:
+        payload["group"] = group
+
+    if tools is not None:
+        payload["tools"] = tools
+
+    return req(
+        client,
+        "POST",
+        "/v1/chat/completions",
+        ChatCompletionResponse,
+        json=payload,
+    )
+
+
+def stream(
+    client: httpx.Client,
+    input: LanguageModelInput,
+    model: str,
+    temperature: float = 0.7,
+    full_chunk: bool = True,
+    group: Optional[str] = None,
+    **kwargs,
+) -> Generator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None, None]:
+    """
+    Stream de complétions OpenAI.
+    - Si `full_chunk` est True (défaut) : chaque yield est le JSON complet du chunk.
+    - Sinon : on garde la compat ascendante → on ne yield que le delta.content + marquages.
+    """
+    formatted_messages = format_messages(input)
+
+    payload = {
+        "model": model,
+        "messages": formatted_messages,
+        "temperature": temperature,
+        "stream": True,
+        **kwargs,
+    }
+
+    if group is not None:
+        payload["group"] = group
+
+    for chunk in req_stream(
+        client,
+        "POST",
+        "/v1/chat/completions",
+        parse_line=lambda line: parse_chat_line(line, full_chunk=full_chunk),
+        json=payload,
+    ):
+        data = ChatCompletionChunk.model_validate(chunk)
+
+        if data is None:
+            continue
+
+        if full_chunk:
+            yield data
+        else:
+            yield data.choices[0].delta.content if data.choices else ""
+
+
+def embeddings(
+    client: httpx.Client,
+    input: str,
+    model: str,
+    encoding_format: EncodingFormat,
+    normalize: bool,
+    group: Optional[str] = None,
+    **kwargs,
+) -> EncodingResponse:
+    """
+    Get embeddings for a given input using the specified model.
+    """
+    payload = {
+        "model": model,
+        "input": input,
+        "encoding_format": encoding_format,
+        "normalize": normalize,
+        **kwargs,
+    }
+
+    if group is not None:
+        payload["group"] = group
+
+    return req(
+        client,
+        "POST",
+        "/v1/embeddings",
+        EncodingResponse,
+        json=payload,
+    )
+
+
+def retrieve(
+    client: httpx.Client,
+    query: str,
+    collections_names: List[str],
+    limit: int,
+    score_threshold: float,
+    filters: list,
+    beta: float,
+    group: Optional[str] = None,
+    **kwargs,
+) -> ChunksDataList:
+    """
+    Retrieve the most relevant documents based on the given query from specified collections.
+    """
+    data = {
+        "query": query,
+        "collections_names": collections_names,
+        "limit": limit,
+        "score": score_threshold,
+        "filters": filters,
+        "beta": beta,
+        **kwargs,
+    }
+
+    if group is not None:
+        data["group"] = group
+
+    return req(
+        client,
+        "POST",
+        "/collections/run/search",
+        ChunksDataList,
+        json=data,
+    )
+
+
+def rerank(
+    client: httpx.Client,
+    query: str,
+    documents: List[str],
+    model: str,
+    top_n: int,
+    return_documents: bool,
+    group: Optional[str] = None,
+    **kwargs,
+) -> RerankerResponse:
+    """
+    Rerank a list of documents based on their relevance to a given query using the specified model.
+    """
+    payload = {
+        "query": query,
+        "documents": documents,
+        "model": model,
+        "top_n": top_n,
+        "return_documents": return_documents,
+        **kwargs,
+    }
+
+    if group is not None:
+        payload["group"] = group
+
+    return req(
+        client,
+        "POST",
+        "/v1/rerank",
+        RerankerResponse,
+        json=payload,
+    )
+
+
+def run_workflow(
+    client: httpx.Client,
+    workflow_id: str,
+    data: dict,
+) -> Any:
+    """
+    Run a specific workflow with the provided data.
+    """
+    try:
+        parser = WorkflowLineParser()
+        for chunk in req_stream(
+            client,
+            "POST",
+            f"/workflow/run/{workflow_id}",
+            parse_line=parser,
+            json=data,
+        ):
+            yield chunk
+    except httpx.HTTPStatusError as e:
+        if e.response.status_code == 404:
+            raise ValueError(f"Workflow with ID {workflow_id} not found")
+        raise e
+
+
+def parse_document(
+    client: httpx.Client,
+    file: Tuple[bytes, str],
+    source: str,
+    group: Optional[str] = None,
+) -> str:
+    """
+    Parse a document using the specified model.
+    """
+    payload = {
+        "source": source,
+    }
+
+    if group is not None:
+        payload["group"] = group
+
+    files = {"file": (file[1], BytesIO(file[0]), "application/octet-stream")}
+
+    response = req(
+        client,
+        "POST",
+        "/document-parser/parsing/parse",
+        AnyDict,
+        files=files,
+        data=payload,
+    )
+    print(response)
+    return response

aisberg/async_client.py

@@ -0,0 +1,107 @@
+from .config import settings
+import httpx
+from typing import Dict, Callable
+
+from .modules import (
+    AsyncChatModule,
+    AsyncCollectionsModule,
+    AsyncEmbeddingsModule,
+    AsyncMeModule,
+    AsyncModelsModule,
+    AsyncWorkflowsModule,
+    ToolsModule,
+)
+
+
+class AisbergAsyncClient:
+    def __init__(self, base_url=None, api_key=None, timeout=None):
+        self.api_key = api_key or settings.aisberg_api_key
+        self.base_url = base_url or settings.aisberg_base_url
+        self.timeout = timeout or settings.timeout
+        self.tool_registry: Dict[str, Callable] = {}
+        self._client = httpx.AsyncClient(
+            base_url=self.base_url,
+            timeout=self.timeout,
+            headers={"Authorization": f"Bearer {self.api_key}"},
+        )
+
+        # Modules
+        self.tools = ToolsModule(self)
+        self.chat = AsyncChatModule(self, self._client)
+        self.models = AsyncModelsModule(self, self._client)
+        self.workflows = AsyncWorkflowsModule(self, self._client)
+        self.me = AsyncMeModule(self, self._client)
+        self.collections = AsyncCollectionsModule(self, self._client)
+        self.embeddings = AsyncEmbeddingsModule(self, self._client)
+
+    async def initialize(self):
+        """
+        Initialise le client asynchrone.
+        Cette méthode est appelée pour s'assurer que le client est prêt à être utilisé.
+        """
+        await self._validate_api_key()
+        return self
+
+    async def _validate_api_key(self):
+        """
+        Valide la clé API en effectuant une requête à l'API.
+        """
+        try:
+            await self.me.info()
+        except httpx.ConnectTimeout as e:
+            raise ConnectionError(
+                f"Le host {self.base_url} n'est pas accessible. Vérifiez votre connexion réseau ou l'URL de l'API."
+            ) from e
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 401:
+                raise ValueError(
+                    "Clé API invalide ou expirée. Veuillez vérifier votre clé API."
+                ) from e
+            elif e.response.status_code == 403:
+                raise PermissionError(
+                    "Accès interdit. Vérifiez vos permissions pour utiliser l'API."
+                ) from e
+            else:
+                raise ValueError(
+                    f"Erreur lors de la validation de la clé API: {e.response.text}"
+                ) from e
+        except Exception:
+            raise ValueError(
+                f"Clé API invalide/expirée ou le host {self.base_url} n'est pas accessible."
+            )
+
+    async def close(self):
+        """
+        Ferme le client.
+        """
+        await self._client.aclose()
+
+    async def __aenter__(self):
+        """
+        Enter the context manager, returning the client instance.
+
+        Example:
+            async with AisbergAsyncClient() as client:
+                # Use the client here
+                pass
+
+        Returns:
+            AisbergAsyncClient: L'instance du client Aisberg.
+        """
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        """
+        Exit the context manager, closing the client.
+        Example:
+            async with AisbergAsyncClient() as client:
+                # Use the client here
+                pass
+        """
+        await self.close()
+
+    def __repr__(self):
+        return f"<AisbergAsyncClient base_url={self.base_url}>"
+
+    def __str__(self):
+        return f"AisbergAsyncClient(base_url={self.base_url})"

aisberg/client.py

@@ -0,0 +1,108 @@
+from .config import settings
+import httpx
+from typing import Dict, Callable
+from .modules import (
+    SyncChatModule,
+    SyncCollectionsModule,
+    SyncEmbeddingsModule,
+    SyncMeModule,
+    SyncModelsModule,
+    SyncWorkflowsModule,
+    ToolsModule,
+)
+
+
+class AisbergClient:
+    def __init__(self, base_url=None, api_key=None, timeout=None):
+        self.api_key = api_key or settings.aisberg_api_key
+        self.base_url = base_url or settings.aisberg_base_url
+        self.timeout = timeout or settings.timeout
+
+        if not self.base_url or not self.api_key:
+            raise ValueError(
+                "L'URL de base et la clé API doivent être définies. "
+                "Utilisez les variables d'environnement AISBERG_API_KEY et AISBERG_BASE_URL ou passez-les lors de l'initialisation du client."
+            )
+
+        self.tool_registry: Dict[str, Callable] = {}
+        self._client = httpx.Client(
+            base_url=self.base_url,
+            timeout=self.timeout,
+            headers={"Authorization": f"Bearer {self.api_key}"},
+        )
+
+        # Modules
+        self.tools = ToolsModule(self)
+        self.chat = SyncChatModule(self, self._client)
+        self.models = SyncModelsModule(self, self._client)
+        self.workflows = SyncWorkflowsModule(self, self._client)
+        self.me = SyncMeModule(self, self._client)
+        self.collections = SyncCollectionsModule(self, self._client)
+        self.embeddings = SyncEmbeddingsModule(self, self._client)
+
+        # Validate API key
+        self._validate_api_key()
+
+    def _validate_api_key(self):
+        """
+        Valide la clé API en effectuant une requête à l'API.
+        """
+        try:
+            self.me.info()
+        except httpx.ConnectTimeout as e:
+            raise ConnectionError(
+                f"Le host {self.base_url} n'est pas accessible. Vérifiez votre connexion réseau ou l'URL de l'API."
+            ) from e
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 401:
+                raise ValueError(
+                    "Clé API invalide ou expirée. Veuillez vérifier votre clé API."
+                ) from e
+            elif e.response.status_code == 403:
+                raise PermissionError(
+                    "Accès interdit. Vérifiez vos permissions pour utiliser l'API."
+                ) from e
+            else:
+                raise ValueError(
+                    f"Erreur lors de la validation de la clé API: {e.response.text}"
+                ) from e
+        except Exception as e:
+            raise ValueError(
+                f"Clé API invalide/expirée ou le host {self.base_url} n'est pas accessible. Erreur: {str(e)}"
+            )
+
+    def close(self):
+        """
+        Ferme le client.
+        """
+        self._client.close()
+
+    def __enter__(self):
+        """
+        Enter the context manager, returning the client instance.
+
+        Example:
+            with AisbergClient() as client:
+                # Use the client here
+                pass
+
+        Returns:
+            AisbergClient: L'instance du client Aisberg.
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Exit the context manager, closing the client.
+        Example:
+            with AisbergClient() as client:
+                # Use the client here
+                pass
+        """
+        self.close()
+
+    def __repr__(self):
+        return f"<AisbergClient base_url={self.base_url}>"
+
+    def __str__(self):
+        return f"AisbergClient(base_url={self.base_url})"

aisberg/config.py

@@ -0,0 +1,17 @@
+from typing import Union
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Settings(BaseSettings):
+    # Variables attendues
+    aisberg_api_key: Union[str, None] = None
+    aisberg_base_url: Union[str, None] = None
+    timeout: int = 30
+
+    # Pour indiquer le fichier .env
+    model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
+
+
+# Singleton partagé dans tout le SDK
+settings = Settings()

aisberg/exceptions.py

@@ -0,0 +1,22 @@
+class APIError(Exception):
+    pass
+
+
+class AuthError(APIError):
+    pass
+
+
+class ToolExecutionError(Exception):
+    """Exception levée lors de l'exécution d'un tool"""
+
+    pass
+
+
+class UnspecifiedClassArgumentError(Exception):
+    """Exception levée lorsqu'un argument requis n'est pas spécifié"""
+
+    def __init__(self, argument_name: str):
+        super().__init__(
+            f"L'argument '{argument_name}' est requis mais n'a pas été spécifié."
+        )
+        self.argument_name = argument_name