aisberg 0.1.0__py3-none-any.whl

aisberg/models/chat.py

@@ -0,0 +1,143 @@
+from typing import Union, Sequence, Tuple, Dict, Any, List, Optional, Literal
+from pydantic import BaseModel, Field
+from enum import Enum
+
+
+class RoleEnum(str, Enum):
+    user = "user"
+    assistant = "assistant"
+    system = "system"
+    tool = "tool"
+
+
+class ToolCall(BaseModel):
+    id: str
+    type: str = "function"
+    function: "FunctionCall"
+
+
+class FunctionCall(BaseModel):
+    name: str
+    arguments: str  # JSON string des arguments
+
+
+class BaseMessage(BaseModel):
+    role: RoleEnum
+    content: Union[str, None]
+
+
+class HumanMessage(BaseMessage):
+    role: RoleEnum = Field(default=RoleEnum.user, frozen=True)
+
+
+class AIMessage(BaseMessage):
+    role: RoleEnum = Field(default=RoleEnum.assistant, frozen=True)
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+class SystemMessage(BaseMessage):
+    role: RoleEnum = Field(default=RoleEnum.system, frozen=True)
+
+
+class ToolMessage(BaseMessage):
+    role: RoleEnum = Field(default=RoleEnum.tool, frozen=True)
+    tool_call_id: str  # ID du tool_call auquel ce message répond
+
+
+# Types de représentations acceptées comme "message"
+MessageLikeRepresentation = Union[
+    BaseMessage,  # BaseMessage ou ses sous-classes
+    str,  # Simple string message
+    Tuple[str, str],  # (role, content)
+    Dict[str, Any],  # dict style {"role": "user", "content": "..."}
+]
+
+# Entrée possible pour un LLM
+LanguageModelInput = Union[str, Sequence[MessageLikeRepresentation]]
+
+
+class ChatMessageContent(BaseModel):
+    role: str
+    content: Optional[str] = ""
+    refusal: Optional[str] = None
+    annotations: Optional[Dict[str, Any]] = None
+    function_call: Optional[FunctionCall] = None
+    tool_calls: Optional[List[ToolCall]] = None
+    reasoning_content: Optional[str] = None
+
+
+class Delta(BaseModel):
+    content: Optional[str] = ""
+    function_call: Optional[dict] = None
+    tool_calls: Optional[list] = None
+
+
+class ChatChoice(BaseModel):
+    index: int
+    message: ChatMessageContent
+    finish_reason: Optional[str]
+
+
+class ChoiceChunk(BaseModel):
+    index: int
+    delta: Delta
+    finish_reason: Optional[str]
+
+
+class ChatUsage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str
+    created: int
+    model: str
+    choices: List[ChatChoice]
+    usage: Optional[ChatUsage] = None
+
+
+class ChatCompletionChunk(BaseModel):
+    id: str
+    object: Literal["chat.completion.chunk"]
+    created: int
+    model: str
+    choices: List[ChoiceChunk]
+    usage: Optional[ChatUsage] = None
+
+
+def parse_message(msg: MessageLikeRepresentation) -> BaseMessage:
+    if isinstance(msg, BaseMessage):
+        return msg
+    elif isinstance(msg, HumanMessage):
+        return msg
+    elif isinstance(msg, AIMessage):
+        return msg
+    elif isinstance(msg, SystemMessage):
+        return msg
+    elif isinstance(msg, ToolMessage):
+        return msg
+    elif isinstance(msg, str):
+        return HumanMessage(content=msg)
+    elif isinstance(msg, tuple) and len(msg) == 2:
+        return BaseMessage(role=RoleEnum(msg[0]), content=msg[1])
+    elif isinstance(msg, dict):
+        return BaseMessage(**msg)
+    else:
+        raise ValueError(f"Invalid message representation: {msg}")
+
+
+def format_messages(input: LanguageModelInput) -> list[dict]:
+    """
+    Normalise et formate les messages pour l’API LLM (OpenAI-like).
+    """
+    if isinstance(input, str):
+        messages = [HumanMessage(content=input)]
+    else:
+        messages = [parse_message(msg) for msg in input]
+    return [
+        msg.model_dump(include={"role", "content", "tool_calls", "tool_call_id"})
+        for msg in messages
+    ]

aisberg/models/collections.py

@@ -0,0 +1,36 @@
+from typing import List, Optional, Union
+from pydantic import BaseModel
+
+
+class Collection(BaseModel):
+    name: str
+
+
+class GroupCollections(BaseModel):
+    group: str
+    collections: List[Collection]
+
+
+# Modèle plus structuré pour payload
+class Payload(BaseModel):
+    method: Optional[str] = None
+    norm: Optional[Union[str, bool]] = None
+    filetype: Optional[str] = None
+    filename: Optional[str] = None
+    dense_encoder: Optional[str] = None
+    Category: Optional[str] = None
+    text: Optional[str] = None
+    timestamp: Optional[str] = None
+    collection_name: Optional[str] = None
+    sparse_encoder: Optional[str] = None
+
+
+class PointDetails(BaseModel):
+    id: str
+    payload: Payload
+
+
+class CollectionDetails(BaseModel):
+    name: str
+    group: str
+    points: List[PointDetails]

aisberg/models/embeddings.py

@@ -0,0 +1,92 @@
+from typing import List, Optional, Dict, Any, Literal, Union
+from pydantic import BaseModel, RootModel, ConfigDict
+
+
+class EmbeddingItem(BaseModel):
+    embedding: List[float]
+    index: int
+    object: str
+
+
+class UsageDetails(BaseModel):
+    prompt_tokens: int
+    total_tokens: int
+    completion_tokens: int
+    prompt_tokens_details: Optional[dict] = None
+
+
+class EncodingResponse(BaseModel):
+    id: str
+    object: str
+    model: str
+    created: int
+    data: List[EmbeddingItem]
+    usage: UsageDetails
+
+
+class ChunkData(BaseModel):
+    id: str
+    idx: int
+    filename: str
+    text: str
+    collection_name: str
+    filetype: str
+    dense_encoder: str
+    sparse_encoder: str
+    timestamp: str
+    norm: Optional[Union[bool, str]] = None
+    rrf_score: Optional[float] = None
+    chunk_method: Optional[str] = None
+
+    # Pour permettre l'ajout de clés dynamiques pour les métadonnées
+    model_config = ConfigDict(extra="allow")
+
+    def metadata(self) -> Dict[str, Any]:
+        """Retourne les métadonnées du ChunkData."""
+        fields = set(self.model_fields.keys())
+        return {k: v for k, v in self.model_dump().items() if k not in fields}
+
+
+class ChunksDataList(RootModel[List[ChunkData]]):
+    def metadata(self) -> List[Dict[str, Any]]:
+        """Retourne une liste des métadonnées de chaque ChunkData."""
+        return [chunk.metadata() for chunk in self.root]
+
+    def texts(self) -> List[str]:
+        """Retourne une liste des textes de chaque ChunkData."""
+        return [chunk.text for chunk in self.root]
+
+    def __iter__(self):
+        return iter(self.root)
+
+    def __len__(self):
+        return len(self.root)
+
+
+class RerankerDocument(BaseModel):
+    text: str
+
+
+class RerankerResult(BaseModel):
+    index: int
+    document: RerankerDocument
+    relevance_score: float
+
+
+class RerankerResponse(BaseModel):
+    id: str
+    model: str
+    usage: Dict[str, int]
+    results: List[RerankerResult]
+
+    def filter_by_relevance_score(self, threshold: float) -> "RerankerResponse":
+        """Retourne une nouvelle instance filtrée de RerankerResponse."""
+        filtered_results = [
+            result for result in self.results if result.relevance_score >= threshold
+        ]
+        return RerankerResponse(
+            id=self.id, model=self.model, usage=self.usage, results=filtered_results
+        )
+
+
+EncodingFormat = Optional[Literal["float", "base64"]]

aisberg/models/models.py

@@ -0,0 +1,39 @@
+from typing import Optional, List, Any
+from pydantic import BaseModel
+
+
+class ModelPermission(BaseModel):
+    id: Optional[str] = None
+    object: Optional[str] = None
+    created: Optional[int] = None
+    allow_create_engine: Optional[bool] = None
+    allow_sampling: Optional[bool] = None
+    allow_logprobs: Optional[bool] = None
+    allow_search_indices: Optional[bool] = None
+    allow_view: Optional[bool] = None
+    allow_fine_tuning: Optional[bool] = None
+    organization: Optional[str] = None
+    group: Optional[Any] = None
+    is_blocking: Optional[bool] = None
+
+
+class ModelMeta(BaseModel):
+    vocab_type: Optional[int] = None
+    n_vocab: Optional[int] = None
+    n_ctx_train: Optional[int] = None
+    n_embd: Optional[int] = None
+    n_params: Optional[int] = None
+    size: Optional[int] = None
+
+
+class Model(BaseModel):
+    id: Optional[str] = None
+    created: Optional[int] = None
+    object: Optional[str] = None
+    owned_by: Optional[str] = None
+    root: Optional[str] = None
+    parent: Optional[Any] = None
+    max_model_len: Optional[int] = None
+    permission: Optional[List[ModelPermission]] = None
+    language: Optional[List[str]] = None
+    meta: Optional[ModelMeta] = None

aisberg/models/requests.py

@@ -0,0 +1,11 @@
+from typing import List, Any
+
+from pydantic import BaseModel, RootModel, ConfigDict
+
+
+class AnyDict(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+
+class AnyList(RootModel[List[Any]]):
+    pass

aisberg/models/token.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+from typing import List
+
+
+class TokenInfo(BaseModel):
+    message: str
+    id: str
+    username: str
+    email: str
+    groups: List[str]
+    roles: List[str]

aisberg/models/tools.py

@@ -0,0 +1,73 @@
+from typing import Dict, List, Iterable
+from pydantic import BaseModel, Field
+
+
+class Property(BaseModel):
+    type: str = Field(..., examples=["string"])
+    description: str
+
+
+class Parameters(BaseModel):
+    """Schéma JSON Schema standard OpenAI pour les paramètres."""
+
+    type: str = Field(default="object", frozen=True)
+    properties: Dict[str, Property]
+    required: List[str]
+
+
+class Function(BaseModel):
+    name: str
+    description: str
+    parameters: Parameters
+
+
+class Tool(BaseModel):
+    type: str = Field(default="function", frozen=True)
+    function: Function
+
+
+def make_tool(
+    *,
+    name: str,
+    description: str,
+    params: Dict[str, str],  # {"param": "description", ...}
+) -> Tool:
+    """Crée un objet `Tool` à partir des paramètres fournis.
+
+    Args:
+        name: Nom de la fonction.
+        description: Description de la fonction.
+        params: Dictionnaire des paramètres avec leur description.
+
+    Returns:
+        Un objet `Tool` prêt à être utilisé.
+    """
+    properties = {
+        key: Property(type="string", description=desc)  # <-- adapt typage si besoin
+        for key, desc in params.items()
+    }
+    return Tool(
+        function=Function(
+            name=name,
+            description=description,
+            parameters=Parameters(
+                properties=properties,
+                required=list(params.keys()),
+            ),
+        )
+    )
+
+
+def tools_to_payload(tools: Iterable[Tool]) -> List[Dict]:
+    """
+    Convertit des objets `Tool` (Pydantic) en payload brut..
+
+    Args:
+        tools: Un itérable d'instances `Tool`.
+    Returns:
+        Une liste de dicts sérialisables en JSON.
+    """
+    if not isinstance(tools, Iterable) or not tools:
+        return []
+
+    return [tool.model_dump(mode="json", exclude_none=True) for tool in tools]

aisberg/models/workflows.py

@@ -0,0 +1,66 @@
+from datetime import datetime
+from typing import Optional, List, Dict, Any
+from pydantic import BaseModel, ConfigDict
+
+
+class Workflow(BaseModel):
+    name: str
+    id: str
+    group: str
+    createdAt: datetime
+    updatedAt: datetime
+    description: Optional[str] = None
+    isActive: Optional[bool] = True
+
+
+class Input(BaseModel):
+    input_id: str
+    input_type: str
+    input_name: str
+    value: Dict[str, Any]
+    linkable: bool
+
+
+class Output(BaseModel):
+    output_id: str
+    output_type: str
+    output_name: str
+    linkable: bool
+
+
+class Position(BaseModel):
+    x: float
+    y: float
+    _id: str
+
+
+class Node(BaseModel):
+    inputs: List[Input]
+    outputs: List[Output]
+    node_id: str
+    name: str
+    source_nodes_id: Optional[List[str]]
+    type: str
+    target_nodes: List[Dict[str, Any]]
+    position: Position
+    _id: str
+
+
+class WorkflowDetails(BaseModel):
+    id: str
+    group: str
+    name: str
+    createdAt: datetime
+    updatedAt: datetime
+    starting_node_id: str
+    nodes: List[Node]
+
+
+class WorkflowRunChunk(BaseModel):
+    response: Optional[str] = None
+    model_config = ConfigDict(extra="allow")
+
+
+class WorkflowRunResult(BaseModel):
+    response: Optional[str] = None
+    model_config = ConfigDict(extra="allow")

aisberg/modules/__init__.py

@@ -0,0 +1,23 @@
+from .chat import AsyncChatModule, SyncChatModule
+from .collections import AsyncCollectionsModule, SyncCollectionsModule
+from .embeddings import AsyncEmbeddingsModule, SyncEmbeddingsModule
+from .me import AsyncMeModule, SyncMeModule
+from .models import AsyncModelsModule, SyncModelsModule
+from .workflows import AsyncWorkflowsModule, SyncWorkflowsModule
+from .tools import ToolsModule
+
+__all__ = [
+    "AsyncChatModule",
+    "SyncChatModule",
+    "AsyncCollectionsModule",
+    "SyncCollectionsModule",
+    "AsyncEmbeddingsModule",
+    "SyncEmbeddingsModule",
+    "AsyncMeModule",
+    "SyncMeModule",
+    "AsyncModelsModule",
+    "SyncModelsModule",
+    "AsyncWorkflowsModule",
+    "SyncWorkflowsModule",
+    "ToolsModule",
+]