aisberg 0.1.0__py3-none-any.whl

aisberg/modules/chat.py

@@ -0,0 +1,403 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Union, Generator, AsyncGenerator, List, Awaitable
+
+from ..models.chat import (
+    LanguageModelInput,
+    HumanMessage,
+    ChatCompletionResponse,
+    AIMessage,
+    ToolMessage,
+    ChatCompletionChunk,
+)
+from ..models.tools import Tool, tools_to_payload
+from ..exceptions import ToolExecutionError
+import json
+from ..api import async_endpoints, endpoints
+from ..abstract.modules import AsyncModule, SyncModule
+
+
+class AbstractChatModule(ABC):
+    """
+    Abstract base class for chat modules.
+    Handles common logic for processing tool calls in conversational LLM flows.
+    """
+
+    def __init__(self, parent, http_client):
+        """
+        Initialize the AbstractChatModule.
+
+        Args:
+            parent: Parent client instance.
+            http_client: HTTP client for making requests.
+        """
+        self._parent = parent
+        self._client = http_client
+
+    @abstractmethod
+    def complete(
+        self,
+        input: LanguageModelInput,
+        model: str,
+        temperature: float = 0.7,
+        tools: Optional[List[Tool]] = None,
+        auto_execute_tools: bool = False,
+        **kwargs,
+    ) -> Union[ChatCompletionResponse, "Awaitable[ChatCompletionResponse]"]:
+        """
+        Sends a chat message (or list of messages) to the model and returns a structured response.
+        Supports tool usage and optional automatic tool execution.
+
+        Args:
+            input (LanguageModelInput): Message(s) to send to the model
+            model (str): Identifier of the language model to use.
+            temperature (float): Sampling temperature for response creativity (0.0–1.0, default: 0.7).
+            tools (Optional[list]): List of tools that can be invoked by the model.
+            auto_execute_tools (bool): If True, automatically executes tool calls generated by the model (default: False).
+            **kwargs: Additional arguments passed to the API.
+
+        Returns:
+            ChatCompletionResponse: Structured model response.
+        """
+        pass
+
+    @abstractmethod
+    def stream(
+        self,
+        input: LanguageModelInput,
+        model: str,
+        temperature: float = 0.7,
+        tools: Optional[List[Tool]] = None,
+        auto_execute_tools: bool = True,
+        full_chunk: bool = True,
+        **kwargs,
+    ) -> Union[
+        Generator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None, None],
+        AsyncGenerator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None],
+    ]:
+        """
+        Streams a chat completion response from the model. If tools are provided,
+        makes an initial non-streaming call to determine tool invocations,
+        executes them if auto_execute_tools is enabled, and then streams the final response.
+
+        Args:
+            input (LanguageModelInput): Message(s) to send to the model
+            model (str): Identifier of the language model to use.
+            temperature (float): Sampling temperature for response creativity (0.0–1.0, default: 0.7).
+            tools (Optional[list]): List of tools that can be invoked by the model.
+            auto_execute_tools (bool): If True, automatically executes tool calls generated by the model (default: False).
+            full_chunk (bool): If True, yields full response objects; if False, yields only the content (default: True).
+            **kwargs: Additional keyword arguments for the API call.
+
+        Yields:
+            Union[str, ChatCompletionChunk, ChatCompletionResponse]: Response chunks or final responses as they become available.
+        """
+        pass
+
+    @abstractmethod
+    def _handle_tool_calls(
+        self,
+        response: ChatCompletionResponse,
+        input,
+        model,
+        temperature,
+        stream: bool = False,
+        full_chunk: bool = True,
+        **kwargs,
+    ):
+        """
+        Handles execution of tool calls found in a model response. Automatically executes the requested tools and
+        integrates their results into the conversation, then continues the chat either in streaming or blocking mode.
+
+        Args:
+            response (ChatCompletionResponse): The initial model response containing tool calls.
+            input: Original input messages.
+            model: Model identifier.
+            temperature: Sampling temperature.
+            stream: If True, returns a streaming generator.
+            full_chunk: (Streaming only) If True, yields full chunks.
+            **kwargs: Additional arguments for the next API call.
+
+        Returns:
+            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,
+        response: ChatCompletionResponse,
+        input: Union[LanguageModelInput, List[HumanMessage]],
+    ):
+        """
+        Utility: Build the message list including tool calls and their results.
+        """
+        tool_calls = response.choices[0].message.tool_calls
+        messages: list = (
+            input if isinstance(input, list) else [HumanMessage(content=input)]
+        )
+
+        messages.append(
+            AIMessage(
+                content=response.choices[0].message.content or "",
+                tool_calls=tool_calls,
+            )
+        )
+
+        for tool_call in tool_calls or []:
+            tool_name = tool_call.function.name
+            tool_arguments_str = tool_call.function.arguments
+            try:
+                tool_arguments = (
+                    json.loads(tool_arguments_str) if tool_arguments_str else {}
+                )
+                tool_result = self._parent.tools.execute(tool_name, tool_arguments)
+                messages.append(
+                    ToolMessage(
+                        content=json.dumps(tool_result)
+                        if not isinstance(tool_result, str)
+                        else tool_result,
+                        tool_call_id=tool_call.id,
+                    )
+                )
+            except json.JSONDecodeError:
+                messages.append(
+                    ToolMessage(
+                        content=f"Error: Invalid JSON arguments provided for tool {tool_name}.",
+                        tool_call_id=tool_call.id,
+                    )
+                )
+            except ToolExecutionError as e:
+                messages.append(
+                    ToolMessage(
+                        content=f"Error: {str(e)}",
+                        tool_call_id=tool_call.id,
+                    )
+                )
+        return messages
+
+
+class SyncChatModule(AbstractChatModule, SyncModule):
+    """
+    `ChatModule` is a synchronous module that provides a high-level interface for interacting with
+    language models in a conversational format. It supports standard completions, streaming responses,
+    and automatic tool execution within conversations. 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 complete(
+        self,
+        input: LanguageModelInput,
+        model: str,
+        temperature: float = 0.7,
+        tools: Optional[list[Tool]] = None,
+        auto_execute_tools: bool = False,
+        **kwargs,
+    ) -> ChatCompletionResponse:
+        response = endpoints.chat(
+            self._client, input, model, temperature, tools_to_payload(tools), **kwargs
+        )
+
+        if (
+            auto_execute_tools
+            and response.choices
+            and response.choices[0].message.tool_calls
+        ):
+            return self._handle_tool_calls(
+                input=input,
+                response=response,
+                model=model,
+                temperature=temperature,
+                stream=False,
+                **kwargs,
+            )
+
+        return response
+
+    def stream(
+        self,
+        input: LanguageModelInput,
+        model: str,
+        temperature: float = 0.7,
+        tools: Optional[list[Tool]] = None,
+        auto_execute_tools: bool = True,
+        full_chunk: bool = True,
+        **kwargs,
+    ) -> Generator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None, None]:
+        if tools:
+            resp = endpoints.chat(
+                self._client,
+                input,
+                model,
+                temperature,
+                tools_to_payload(tools),
+                **kwargs,
+            )
+            response = ChatCompletionResponse.model_validate(resp)
+
+            if not auto_execute_tools:
+                yield response if full_chunk else response.choices[0].message.content
+                return
+
+            yield from self._handle_tool_calls(
+                response,
+                input,
+                model,
+                temperature,
+                stream=True,
+                full_chunk=full_chunk,
+                **kwargs,
+            )
+        else:
+            for chunk in endpoints.stream(
+                self._client,
+                input,
+                model,
+                temperature,
+                full_chunk=full_chunk,
+                **kwargs,
+            ):
+                yield ChatCompletionChunk.model_validate(chunk) if full_chunk else chunk
+
+    def _handle_tool_calls(
+        self,
+        response: ChatCompletionResponse,
+        input,
+        model,
+        temperature,
+        stream: bool = False,
+        full_chunk: bool = True,
+        **kwargs,
+    ):
+        tool_calls = response.choices[0].message.tool_calls
+
+        if tool_calls is None:
+            return response
+
+        messages = self._build_tool_messages(response, input)
+
+        if stream:
+            return self.stream(
+                messages, model, temperature, full_chunk=full_chunk, **kwargs
+            )
+        else:
+            return self.complete(messages, model, temperature, **kwargs)
+
+
+class AsyncChatModule(AbstractChatModule, AsyncModule):
+    """
+    `ChatModule` is an asynchronous module that provides a high-level interface for interacting with
+    language models in a conversational format. It supports standard completions, streaming responses,
+    and automatic tool execution within conversations. 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 complete(
+        self,
+        input: LanguageModelInput,
+        model: str,
+        temperature: float = 0.7,
+        tools: Optional[list[Tool]] = None,
+        auto_execute_tools: bool = False,
+        **kwargs,
+    ) -> ChatCompletionResponse:
+        resp = await async_endpoints.chat(
+            self._client, input, model, temperature, tools_to_payload(tools), **kwargs
+        )
+        response = ChatCompletionResponse.model_validate(resp)
+
+        if (
+            auto_execute_tools
+            and response.choices
+            and response.choices[0].message.tool_calls
+        ):
+            return await self._handle_tool_calls(
+                input=input,
+                response=response,
+                model=model,
+                temperature=temperature,
+                stream=False,
+                **kwargs,
+            )
+
+        return response
+
+    async def stream(
+        self,
+        input: LanguageModelInput,
+        model: str,
+        temperature: float = 0.7,
+        tools: Optional[list[Tool]] = None,
+        auto_execute_tools: bool = True,
+        full_chunk: bool = True,
+        **kwargs,
+    ) -> AsyncGenerator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None]:
+        if tools:
+            resp = await async_endpoints.chat(
+                self._client,
+                input,
+                model,
+                temperature,
+                tools_to_payload(tools),
+                **kwargs,
+            )
+            response = ChatCompletionResponse.model_validate(resp)
+
+            if not auto_execute_tools:
+                if full_chunk:
+                    yield response
+                else:
+                    yield response.choices[0].message.content
+                return
+
+            async for val in await self._handle_tool_calls(
+                response,
+                input,
+                model,
+                temperature,
+                stream=True,
+                full_chunk=full_chunk,
+                **kwargs,
+            ):
+                yield val
+        else:
+            async for chunk in async_endpoints.chat_stream(
+                self._client,
+                input,
+                model,
+                temperature,
+                full_chunk=full_chunk,
+                **kwargs,
+            ):
+                yield ChatCompletionChunk.model_validate(chunk) if full_chunk else chunk
+
+    async def _handle_tool_calls(
+        self,
+        response: ChatCompletionResponse,
+        input,
+        model,
+        temperature,
+        stream: bool = False,
+        full_chunk: bool = True,
+        **kwargs,
+    ):
+        tool_calls = response.choices[0].message.tool_calls
+
+        if tool_calls is None:
+            return response
+
+        messages = self._build_tool_messages(response, input)
+
+        if stream:
+            return self.stream(
+                messages, model, temperature, full_chunk=full_chunk, **kwargs
+            )
+        else:
+            return await self.complete(messages, model, temperature, **kwargs)

aisberg/modules/collections.py

@@ -0,0 +1,117 @@
+from typing import List
+from abc import ABC
+
+from ..models.collections import GroupCollections, Collection, CollectionDetails
+
+from abc import abstractmethod
+from ..abstract.modules import SyncModule, AsyncModule
+from ..api import endpoints, async_endpoints
+
+
+class AbstractCollectionsModule(ABC):
+    def __init__(self, parent, client):
+        self._parent = parent
+        self._client = client
+
+    @abstractmethod
+    def list(self) -> List[GroupCollections]:
+        """
+        Get a list of available collections. Collections are grouped by your belonging groups.
+
+        Returns:
+            List[GroupCollections]: 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_by_group(self, group_id: str) -> List[Collection]:
+        """
+        Get collections by group ID.
+
+        Args:
+            group_id (str): The ID of the group for which to retrieve collections.
+
+        Returns:
+            List[Collection]: A list of collections for the specified group.
+
+        Raises:
+            ValueError: If no collections are found for the specified group ID.
+            Exception: If there is an error fetching the collections.
+        """
+        pass
+
+    @abstractmethod
+    def details(self, collection_id: str, group_id: str) -> CollectionDetails:
+        """
+        Get details of a specific collection.
+
+        Args:
+            collection_id (str): The ID of the collection to retrieve.
+            group_id (str): The ID of the group to which the collection belongs.
+
+        Returns:
+            CollectionDetails: The details of the specified collection.
+
+        Raises:
+            ValueError: If the specified collection is not found.
+        """
+        pass
+
+    @staticmethod
+    def _get_collections_by_group(
+        collections: List[GroupCollections], group_id: str
+    ) -> List[Collection]:
+        for group in collections:
+            if group.group == group_id:
+                return group.collections
+        raise ValueError("No collections found for group ID")
+
+
+class SyncCollectionsModule(SyncModule, AbstractCollectionsModule):
+    def __init__(self, parent, client):
+        SyncModule.__init__(self, parent, client)
+        AbstractCollectionsModule.__init__(self, parent, client)
+
+    def list(self) -> List[GroupCollections]:
+        return endpoints.collections(self._client)
+
+    def get_by_group(self, group_id: str) -> List[Collection]:
+        collections = self.list()
+        return self._get_collections_by_group(collections, group_id)
+
+    def details(self, collection_id: str, group_id: str) -> CollectionDetails:
+        points = endpoints.collection(self._client, collection_id, group_id)
+        if points is None:
+            raise ValueError("No collection found")
+        return CollectionDetails(
+            name=collection_id,
+            group=group_id,
+            points=points,
+        )
+
+
+class AsyncCollectionsModule(AsyncModule, AbstractCollectionsModule):
+    def __init__(self, parent, client):
+        AsyncModule.__init__(self, parent, client)
+        AbstractCollectionsModule.__init__(self, parent, client)
+
+    async def list(self) -> List[GroupCollections]:
+        return await async_endpoints.collections(self._client)
+
+    async def get_by_group(self, group_id: str) -> List[Collection]:
+        collections = await self.list()
+        return self._get_collections_by_group(collections, group_id)
+
+    async def details(self, collection_id: str, group_id: str) -> CollectionDetails:
+        points = await async_endpoints.collection(self._client, collection_id, group_id)
+        if points is None:
+            raise ValueError("No collection found")
+        return CollectionDetails(
+            name=collection_id,
+            group=group_id,
+            points=points,
+        )

aisberg/modules/document.py

@@ -0,0 +1,117 @@
+from typing import List
+from abc import ABC
+
+from ..models.document import GroupDocument, Collection, CollectionDetails
+
+from abc import abstractmethod
+from ..abstract.modules import SyncModule, AsyncModule
+from ..api import endpoints, async_endpoints
+
+
+class AbstractDocumentModule(ABC):
+    def __init__(self, parent, client):
+        self._parent = parent
+        self._client = client
+
+    @abstractmethod
+    def list(self) -> List[GroupDocument]:
+        """
+        Get a list of available collections. Document are grouped by your belonging groups.
+
+        Returns:
+            List[GroupDocument]: 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_by_group(self, group_id: str) -> List[Collection]:
+        """
+        Get collections by group ID.
+
+        Args:
+            group_id (str): The ID of the group for which to retrieve collections.
+
+        Returns:
+            List[Collection]: A list of collections for the specified group.
+
+        Raises:
+            ValueError: If no collections are found for the specified group ID.
+            Exception: If there is an error fetching the collections.
+        """
+        pass
+
+    @abstractmethod
+    def details(self, collection_id: str, group_id: str) -> CollectionDetails:
+        """
+        Get details of a specific collection.
+
+        Args:
+            collection_id (str): The ID of the collection to retrieve.
+            group_id (str): The ID of the group to which the collection belongs.
+
+        Returns:
+            CollectionDetails: The details of the specified collection.
+
+        Raises:
+            ValueError: If the specified collection is not found.
+        """
+        pass
+
+    @staticmethod
+    def _get_collections_by_group(
+        collections: List[GroupDocument], group_id: str
+    ) -> List[Collection]:
+        for group in collections:
+            if group.group == group_id:
+                return group.collections
+        raise ValueError("No collections found for group ID")
+
+
+class SyncDocumentModule(SyncModule, AbstractDocumentModule):
+    def __init__(self, parent, client):
+        SyncModule.__init__(self, parent, client)
+        AbstractDocumentModule.__init__(self, parent, client)
+
+    def list(self) -> List[GroupDocument]:
+        return endpoints.collections(self._client)
+
+    def get_by_group(self, group_id: str) -> List[Collection]:
+        collections = self.list()
+        return self._get_collections_by_group(collections, group_id)
+
+    def details(self, collection_id: str, group_id: str) -> CollectionDetails:
+        points = endpoints.collection(self._client, collection_id, group_id)
+        if points is None:
+            raise ValueError("No collection found")
+        return CollectionDetails(
+            name=collection_id,
+            group=group_id,
+            points=points,
+        )
+
+
+class AsyncDocumentModule(AsyncModule, AbstractDocumentModule):
+    def __init__(self, parent, client):
+        AsyncModule.__init__(self, parent, client)
+        AbstractDocumentModule.__init__(self, parent, client)
+
+    async def list(self) -> List[GroupDocument]:
+        return await async_endpoints.collections(self._client)
+
+    async def get_by_group(self, group_id: str) -> List[Collection]:
+        collections = await self.list()
+        return self._get_collections_by_group(collections, group_id)
+
+    async def details(self, collection_id: str, group_id: str) -> CollectionDetails:
+        points = await async_endpoints.collection(self._client, collection_id, group_id)
+        if points is None:
+            raise ValueError("No collection found")
+        return CollectionDetails(
+            name=collection_id,
+            group=group_id,
+            points=points,
+        )