aisberg 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

aisberg/api/async_endpoints.py

@@ -1,4 +1,4 @@
-from io import BytesIO
+import json
 
 import httpx
 from ..models.chat import (
@@ -7,9 +7,10 @@ from ..models.chat import (
     ChatCompletionResponse,
     ChatCompletionChunk,
 )
-from typing import Optional, AsyncGenerator, Union, List, Any, Tuple
+from typing import Optional, AsyncGenerator, Union, List, Any
 
-from ..models.collections import GroupCollections, PointDetails
+from ..models.collections import GroupCollections, PointDetails, ChunkingDictInput
+from ..models.documents import DocumentParserResponse
 from ..models.embeddings import (
     EncodingFormat,
     EncodingResponse,
@@ -21,7 +22,7 @@ from ..models.token import TokenInfo
 from ..models.workflows import WorkflowDetails, Workflow
 from ..utils import parse_chat_line, WorkflowLineParser
 from ..requests.async_requests import areq, areq_stream
-from ..models.requests import AnyDict, AnyList
+from ..models.requests import AnyDict, AnyList, HttpxFileField
 
 
 async def models(client: httpx.AsyncClient) -> List[Model]:
@@ -67,7 +68,7 @@ async def collections(client: httpx.AsyncClient) -> List[GroupCollections]:
 
 
 async def collection(
-    client: httpx.AsyncClient, collection_id: str, group_id: str
+    client: httpx.AsyncClient, collection_id: str, group_id: Optional[str] = None
 ) -> List[PointDetails]:
     """
     Get details of a specific collection.
@@ -85,6 +86,128 @@ async def collection(
         raise e
 
 
+async def create_collection(
+    client: httpx.AsyncClient,
+    name: str,
+    model: str,
+    group: Optional[str] = None,
+):
+    """
+    Create a new collection with the specified name and optional group.
+    """
+    payload = {"collection_name": name, "embedding_model": model}
+    if group is not None:
+        payload["group"] = group
+
+    return await areq(
+        client,
+        "POST",
+        "/collections",
+        AnyDict,
+        json=payload,
+    )
+
+
+async def delete_collection(
+    client: httpx.AsyncClient,
+    name: str,
+    group: Optional[str] = None,
+):
+    """
+    Delete a collection with the specified name and optional group.
+    """
+    payload = {"collections": [name]}
+    if group is not None:
+        payload["group"] = group
+
+    return await areq(
+        client,
+        "DELETE",
+        "/collections",
+        AnyDict,
+        json=payload,
+    )
+
+
+async def insert_points_in_collection(
+    client: httpx.AsyncClient,
+    name: str,
+    files: HttpxFileField,
+    normalize: bool,
+    chunking_dict: Optional[ChunkingDictInput] = None,
+    group: Optional[str] = None,
+):
+    """
+    Insert points into a collection with the specified name.
+    """
+    payload = {
+        "chunking_dict": json.dumps({"method": "custom", "params": {}}),
+        "normalize": normalize,
+    }
+    if group is not None:
+        payload["group"] = group
+
+    if chunking_dict is not None:
+        payload["chunking_dict"] = chunking_dict.model_dump_json()
+
+    return await areq(
+        client,
+        "POST",
+        f"/collections/{name}",
+        AnyDict,
+        data=payload,
+        files=files,
+    )
+
+
+async def delete_points_in_collection(
+    client: httpx.AsyncClient,
+    points_ids: List[str],
+    name: str,
+    group: Optional[str] = None,
+):
+    """
+    Delete points into a collection with the specified name.
+    """
+    payload = {
+        "points": points_ids,
+        "collection": name,
+    }
+    if group is not None:
+        payload["group"] = group
+
+    return await areq(
+        client,
+        "DELETE",
+        "/collections/chunks",
+        AnyDict,
+        json=payload,
+    )
+
+
+async def delete_all_points_in_collection(
+    client: httpx.AsyncClient,
+    name: str,
+    group: Optional[str] = None,
+):
+    """
+    Delete All points into a collection with the specified name.
+    """
+    payload = {
+        "collection": name,
+    }
+    if group is not None:
+        payload["group"] = group
+
+    return await areq(
+        client,
+        "DELETE",
+        "/collections/all/chunks",
+        AnyDict,
+        json=payload,
+    )
+
+
 async def me(client: httpx.AsyncClient) -> TokenInfo:
     """
     Get the details of the current user.
@@ -303,31 +426,26 @@ async def run_workflow(
         raise e
 
 
-async def parse_document(
+async def parse_documents(
     client: httpx.AsyncClient,
-    file: Tuple[bytes, str],
-    source: str,
+    files: HttpxFileField,
     group: Optional[str] = None,
-) -> str:
+    **kwargs,
+) -> DocumentParserResponse:
     """
-    Parse a document using the specified model.
+    Parse a single or multiple documents using the document parser endpoint.
+    Returns the ID of the parsed document to be downloaded later from the S3 bucket.
     """
-    payload = {
-        "source": source,
-    }
-
+    payload = {**kwargs}
     if group is not None:
         payload["group"] = group
 
-    files = {"file": (file[1], BytesIO(file[0]), "application/octet-stream")}
-
-    response = areq(
+    response = await areq(
         client,
         "POST",
         "/document-parser/parsing/parse",
-        AnyDict,
+        DocumentParserResponse,
         files=files,
-        json=payload,
+        data=payload,
     )
-    print(response)
     return response

aisberg/api/endpoints.py

@@ -1,4 +1,4 @@
-from io import BytesIO
+import json
 
 import httpx
 from ..models.chat import (
@@ -7,9 +7,10 @@ from ..models.chat import (
     ChatCompletionResponse,
     ChatCompletionChunk,
 )
-from typing import Optional, Generator, Union, List, Any, Tuple
+from typing import Optional, Generator, Union, List, Any
 
-from ..models.collections import GroupCollections, PointDetails
+from ..models.collections import GroupCollections, PointDetails, ChunkingDictInput
+from ..models.documents import DocumentParserResponse
 from ..models.embeddings import (
     EncodingFormat,
     EncodingResponse,
@@ -21,7 +22,7 @@ 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
+from ..models.requests import AnyDict, AnyList, HttpxFileField
 
 
 def models(client: httpx.Client) -> List[Model]:
@@ -64,7 +65,7 @@ def collections(client: httpx.Client) -> List[GroupCollections]:
 
 
 def collection(
-    client: httpx.Client, collection_id: str, group_id: str
+    client: httpx.Client, collection_id: str, group_id: Optional[str] = None
 ) -> List[PointDetails]:
     """
     Get details of a specific collection.
@@ -80,6 +81,128 @@ def collection(
         raise e
 
 
+def create_collection(
+    client: httpx.Client,
+    name: str,
+    model: str,
+    group: Optional[str] = None,
+):
+    """
+    Create a new collection with the specified name and optional group.
+    """
+    payload = {"collection_name": name, "embedding_model": model}
+    if group is not None:
+        payload["group"] = group
+
+    return req(
+        client,
+        "POST",
+        "/collections",
+        AnyDict,
+        json=payload,
+    )
+
+
+def delete_collection(
+    client: httpx.Client,
+    name: str,
+    group: Optional[str] = None,
+):
+    """
+    Delete a collection with the specified name and optional group.
+    """
+    payload = {"collections": [name]}
+    if group is not None:
+        payload["group"] = group
+
+    return req(
+        client,
+        "DELETE",
+        "/collections",
+        AnyDict,
+        json=payload,
+    )
+
+
+def insert_points_in_collection(
+    client: httpx.Client,
+    name: str,
+    files: HttpxFileField,
+    normalize: bool,
+    chunking_dict: Optional[ChunkingDictInput] = None,
+    group: Optional[str] = None,
+):
+    """
+    Insert points into a collection with the specified name.
+    """
+    payload = {
+        "chunking_dict": json.dumps({"method": "custom", "params": {}}),
+        "normalize": normalize,
+    }
+    if group is not None:
+        payload["group"] = group
+
+    if chunking_dict is not None:
+        payload["chunking_dict"] = chunking_dict.model_dump_json()
+
+    return req(
+        client,
+        "POST",
+        f"/collections/{name}",
+        AnyDict,
+        data=payload,
+        files=files,
+    )
+
+
+def delete_points_in_collection(
+    client: httpx.Client,
+    points_ids: List[str],
+    name: str,
+    group: Optional[str] = None,
+):
+    """
+    Delete points into a collection with the specified name.
+    """
+    payload = {
+        "points": points_ids,
+        "collection": name,
+    }
+    if group is not None:
+        payload["group"] = group
+
+    return req(
+        client,
+        "DELETE",
+        "/collections/chunks",
+        AnyDict,
+        json=payload,
+    )
+
+
+def delete_all_points_in_collection(
+    client: httpx.Client,
+    name: str,
+    group: Optional[str] = None,
+):
+    """
+    Delete All points into a collection with the specified name.
+    """
+    payload = {
+        "collection": name,
+    }
+    if group is not None:
+        payload["group"] = group
+
+    return req(
+        client,
+        "DELETE",
+        "/collections/all/chunks",
+        AnyDict,
+        json=payload,
+    )
+
+
 def me(client: httpx.Client) -> TokenInfo:
     """
     Get the details of the current user.
@@ -298,31 +421,26 @@ def run_workflow(
         raise e
 
 
-def parse_document(
+def parse_documents(
     client: httpx.Client,
-    file: Tuple[bytes, str],
-    source: str,
+    files: HttpxFileField,
     group: Optional[str] = None,
-) -> str:
+    **kwargs,
+) -> DocumentParserResponse:
     """
-    Parse a document using the specified model.
+    Parse a single or multiple documents using the document parser endpoint.
+    Returns the ID of the parsed document to be downloaded later from the S3 bucket.
     """
-    payload = {
-        "source": source,
-    }
-
+    payload = {**kwargs}
     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,
+        DocumentParserResponse,
         files=files,
         data=payload,
     )
-    print(response)
     return response

aisberg/async_client.py

@@ -10,6 +10,8 @@ from .modules import (
     AsyncModelsModule,
     AsyncWorkflowsModule,
     ToolsModule,
+    AsyncDocumentsModule,
+    SyncS3Module,
 )
 
 
@@ -33,6 +35,12 @@ class AisbergAsyncClient:
         self.me = AsyncMeModule(self, self._client)
         self.collections = AsyncCollectionsModule(self, self._client)
         self.embeddings = AsyncEmbeddingsModule(self, self._client)
+        self.documents = AsyncDocumentsModule(self, self._client)
+        self._s3 = SyncS3Module(
+            settings.s3_access_key_id,
+            settings.s3_secret_access_key,
+            settings.s3_endpoint,
+        )
 
     async def initialize(self):
         """

aisberg/client.py

@@ -9,6 +9,8 @@ from .modules import (
     SyncModelsModule,
     SyncWorkflowsModule,
     ToolsModule,
+    SyncDocumentsModule,
+    SyncS3Module,
 )
 
 
@@ -39,6 +41,12 @@ class AisbergClient:
         self.me = SyncMeModule(self, self._client)
         self.collections = SyncCollectionsModule(self, self._client)
         self.embeddings = SyncEmbeddingsModule(self, self._client)
+        self.documents = SyncDocumentsModule(self, self._client)
+        self._s3 = SyncS3Module(
+            settings.s3_access_key_id,
+            settings.s3_secret_access_key,
+            settings.s3_endpoint,
+        )
 
         # Validate API key
         self._validate_api_key()

aisberg/config.py

@@ -5,10 +5,16 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
 
 class Settings(BaseSettings):
     # Variables attendues
+    # -- API --
     aisberg_api_key: Union[str, None] = None
     aisberg_base_url: Union[str, None] = None
     timeout: int = 30
 
+    # -- S3 --
+    s3_access_key_id: Union[str, None] = None
+    s3_secret_access_key: Union[str, None] = None
+    s3_endpoint: Union[str, None] = None
+
     # Pour indiquer le fichier .env
     model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
 

aisberg/models/collections.py

@@ -11,6 +11,20 @@ class GroupCollections(BaseModel):
     collections: List[Collection]
 
 
+class CollectionDataset(BaseModel):
+    chunks: List[str]
+    metadata: Optional[dict] = []
+
+
+class CollectionCreateResponse(BaseModel):
+    message: Optional[str] = None
+
+
+class ChunkingDictInput(BaseModel):
+    method: Optional[str] = "custom"
+    params: Optional[dict] = {}
+
+
 # Modèle plus structuré pour payload
 class Payload(BaseModel):
     method: Optional[str] = None
@@ -32,5 +46,5 @@ class PointDetails(BaseModel):
 
 class CollectionDetails(BaseModel):
     name: str
-    group: str
+    group: Optional[str] = None
     points: List[PointDetails]

aisberg/models/documents.py

@@ -0,0 +1,46 @@
+from pydantic import BaseModel
+from typing import Optional, List, Tuple, Union
+from io import BytesIO
+
+
+class DocumentParserResponse(BaseModel):
+    """
+    Response model for document parsing.
+    """
+
+    message: Optional[str] = None
+    parsedFiles: Optional[List[str]] = None
+    bucketName: Optional[str] = None
+
+
+class FileObject(BaseModel):
+    """
+    Represents a file object with its name and content.
+    """
+
+    name: str
+    buffer: bytes
+
+
+class DocumentParserDocOutput(BaseModel):
+    type: str
+    data: Union[str, dict, list]
+
+
+class ParsedDocument(BaseModel):
+    """
+    Represents a parsed document with its content and metadata.
+    """
+
+    content: DocumentParserDocOutput
+    metadata: Optional[dict] = None
+
+
+DocumentParserFileInput = Union[
+    str,
+    bytes,
+    BytesIO,
+    Tuple[bytes, str],
+    "FileObject",
+    List[Union[str, bytes, BytesIO, Tuple[bytes, str], "FileObject"]],
+]

aisberg/models/requests.py

@@ -1,4 +1,5 @@
-from typing import List, Any
+from typing import List, Any, Tuple
+from io import BytesIO
 
 from pydantic import BaseModel, RootModel, ConfigDict
 
@@ -9,3 +10,6 @@ class AnyDict(BaseModel):
 
 class AnyList(RootModel[List[Any]]):
     pass
+
+
+HttpxFileField = List[Tuple[str, Tuple[str, BytesIO, str]]]

aisberg/modules/__init__.py

@@ -5,6 +5,8 @@ from .me import AsyncMeModule, SyncMeModule
 from .models import AsyncModelsModule, SyncModelsModule
 from .workflows import AsyncWorkflowsModule, SyncWorkflowsModule
 from .tools import ToolsModule
+from .documents import AsyncDocumentsModule, SyncDocumentsModule
+from .s3 import SyncS3Module
 
 __all__ = [
     "AsyncChatModule",
@@ -20,4 +22,7 @@ __all__ = [
     "AsyncWorkflowsModule",
     "SyncWorkflowsModule",
     "ToolsModule",
+    "AsyncDocumentsModule",
+    "SyncDocumentsModule",
+    "SyncS3Module",
 ]

aisberg/modules/chat.py

@@ -14,6 +14,9 @@ from ..exceptions import ToolExecutionError
 import json
 from ..api import async_endpoints, endpoints
 from ..abstract.modules import AsyncModule, SyncModule
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 class AbstractChatModule(ABC):
@@ -58,7 +61,7 @@ class AbstractChatModule(ABC):
         Returns:
             ChatCompletionResponse: Structured model response.
         """
-        pass
+        ...
 
     @abstractmethod
     def stream(
@@ -91,7 +94,7 @@ class AbstractChatModule(ABC):
         Yields:
             Union[str, ChatCompletionChunk, ChatCompletionResponse]: Response chunks or final responses as they become available.
         """
-        pass
+        ...
 
     @abstractmethod
     def _handle_tool_calls(
@@ -121,7 +124,7 @@ class AbstractChatModule(ABC):
             ChatCompletionResponse or AsyncGenerator: If stream is True, returns a generator yielding tool results and final response.
             Otherwise, returns a complete ChatCompletionResponse with tool results integrated.
         """
-        pass
+        ...
 
     def _build_tool_messages(
         self,
@@ -160,6 +163,10 @@ class AbstractChatModule(ABC):
                     )
                 )
             except json.JSONDecodeError:
+                logger.error(
+                    f"Invalid JSON arguments provided for tool {tool_name}."
+                    f" Arguments: {tool_arguments_str}"
+                )
                 messages.append(
                     ToolMessage(
                         content=f"Error: Invalid JSON arguments provided for tool {tool_name}.",
@@ -167,6 +174,7 @@ class AbstractChatModule(ABC):
                     )
                 )
             except ToolExecutionError as e:
+                logger.error(f"Error executing tool {tool_name}: {str(e)}")
                 messages.append(
                     ToolMessage(
                         content=f"Error: {str(e)}",