dais-sdk 0.6.0__py3-none-any.whl

dais_sdk/__init__.py

@@ -0,0 +1,320 @@
+import asyncio
+import queue
+from typing import cast
+from collections.abc import AsyncGenerator, Generator
+from litellm import CustomStreamWrapper, completion, acompletion
+from litellm.utils import ProviderConfigManager
+from litellm.types.utils import (
+    LlmProviders,
+    ModelResponse as LiteLlmModelResponse,
+    ModelResponseStream as LiteLlmModelResponseStream
+)
+from .debug import enable_debugging
+from .param_parser import ParamParser
+from .stream import AssistantMessageCollector
+from .tool.execute import execute_tool_sync, execute_tool
+from .tool.toolset import (
+    Toolset,
+    python_tool,
+    PythonToolset,
+    McpToolset,
+    LocalMcpToolset,
+    RemoteMcpToolset,
+)
+from .tool.utils import find_tool_by_name
+from .mcp_client import (
+    McpClient,
+    McpTool,
+    McpToolResult,
+    LocalMcpClient,
+    RemoteMcpClient,
+    LocalServerParams,
+    RemoteServerParams,
+    OAuthParams,
+)
+from .types import (
+    GenerateTextResponse,
+    StreamTextResponseSync, StreamTextResponseAsync,
+    FullMessageQueueSync, FullMessageQueueAsync,
+)
+from .types.request_params import LlmRequestParams
+from .types.tool import ToolFn, ToolDef, RawToolDef, ToolLike
+from .types.exceptions import (
+    AuthenticationError,
+    PermissionDeniedError,
+    RateLimitError,
+    ContextWindowExceededError,
+    BadRequestError,
+    InvalidRequestError,
+    InternalServerError,
+    ServiceUnavailableError,
+    ContentPolicyViolationError,
+    APIError,
+    Timeout,
+)
+from .types.message import (
+    ChatMessage, UserMessage, SystemMessage, AssistantMessage, ToolMessage,
+    MessageChunk, TextChunk, UsageChunk, ReasoningChunk, AudioChunk, ImageChunk, ToolCallChunk,
+    openai_chunk_normalizer
+)
+from .logger import logger, enable_logging
+
+class LLM:
+    """
+    The `generate_text` and `stream_text` API will return ToolMessage in the returned sequence
+    only if `params.execute_tools` is True.
+
+    - - -
+
+    Possible exceptions raises for `generate_text` and `stream_text`:
+        - AuthenticationError
+        - PermissionDeniedError
+        - RateLimitError
+        - ContextWindowExceededError
+        - BadRequestError
+        - InvalidRequestError
+        - InternalServerError
+        - ServiceUnavailableError
+        - ContentPolicyViolationError
+        - APIError
+        - Timeout
+    """
+
+    def __init__(self,
+                 provider: LlmProviders,
+                 base_url: str,
+                 api_key: str):
+        self.provider = provider
+        self.base_url = base_url
+        self.api_key = api_key
+        self._param_parser = ParamParser(self.provider, self.base_url, self.api_key)
+
+    @staticmethod
+    async def execute_tool_call(
+            params: LlmRequestParams,
+            incomplete_tool_message: ToolMessage
+            ) -> tuple[str | None, str | None]:
+        """
+        Receive incomplete tool messages, execute the tool calls and
+        return the result and error tuple.
+        """
+        name, arguments = incomplete_tool_message.name, incomplete_tool_message.arguments
+        tool_def = params.find_tool(incomplete_tool_message.name)
+        if tool_def is None:
+            raise LlmRequestParams.ToolDoesNotExistError(name)
+
+        result, error = None, None
+        try:
+            result = await execute_tool(tool_def, arguments)
+        except Exception as e:
+            error = f"{type(e).__name__}: {str(e)}"
+        return result, error
+
+    @staticmethod
+    def execute_tool_call_sync(
+            params: LlmRequestParams,
+            incomplete_tool_message: ToolMessage
+            ) -> tuple[str | None, str | None]:
+        """
+        Synchronous version of `execute_tool_call`.
+        """
+        name, arguments = incomplete_tool_message.name, incomplete_tool_message.arguments
+        tool_def = params.find_tool(incomplete_tool_message.name)
+        if tool_def is None:
+            raise LlmRequestParams.ToolDoesNotExistError(name)
+
+        result, error = None, None
+        try:
+            result = execute_tool_sync(tool_def, arguments)
+        except Exception as e:
+            error = f"{type(e).__name__}: {str(e)}"
+        return result, error
+
+    def _resolve_tool_calls_sync(self, params: LlmRequestParams, assistant_message: AssistantMessage) -> Generator[ToolMessage]:
+        if not params.execute_tools: return
+        if (incomplete_tool_messages
+            := assistant_message.get_incomplete_tool_messages()) is None:
+            return
+        for incomplete_tool_message in incomplete_tool_messages:
+            try:
+                result, error = LLM.execute_tool_call_sync(params, incomplete_tool_message)
+            except LlmRequestParams.ToolDoesNotExistError as e:
+                logger.warning(f"{e.message} Skipping this tool call.")
+                continue
+            yield ToolMessage(
+                tool_call_id=incomplete_tool_message.tool_call_id,
+                name=incomplete_tool_message.name,
+                arguments=incomplete_tool_message.arguments,
+                result=result,
+                error=error)
+
+    async def _resolve_tool_calls(self, params: LlmRequestParams, assistant_message: AssistantMessage) -> AsyncGenerator[ToolMessage]:
+        if not params.execute_tools: return
+        if (incomplete_tool_messages :=
+            assistant_message.get_incomplete_tool_messages()) is None:
+            return
+        for incomplete_tool_message in incomplete_tool_messages:
+            try:
+                result, error = await LLM.execute_tool_call(params, incomplete_tool_message)
+            except LlmRequestParams.ToolDoesNotExistError as e:
+                logger.warning(f"{e.message} Skipping this tool call.")
+                continue
+            yield ToolMessage(
+                tool_call_id=incomplete_tool_message.tool_call_id,
+                name=incomplete_tool_message.name,
+                arguments=incomplete_tool_message.arguments,
+                result=result,
+                error=error)
+
+    def list_models(self) -> list[str]:
+        provider_config = ProviderConfigManager.get_provider_model_info(
+            model=None,
+            provider=self.provider,
+        )
+
+        if provider_config is None:
+            raise ValueError(f"The '{self.provider}' provider is not supported to list models.")
+
+        try:
+            models = provider_config.get_models(
+                api_key=self.api_key,
+                api_base=self.base_url
+            )
+        except Exception as e:
+            raise e
+        return models
+
+    def generate_text_sync(self, params: LlmRequestParams) -> GenerateTextResponse:
+        response = completion(**self._param_parser.parse_nonstream(params))
+        response = cast(LiteLlmModelResponse, response)
+        assistant_message = AssistantMessage.from_litellm_message(response)
+        result: GenerateTextResponse = [assistant_message]
+        for tool_message in self._resolve_tool_calls_sync(params, assistant_message):
+            result.append(tool_message)
+        return result
+
+    async def generate_text(self, params: LlmRequestParams) -> GenerateTextResponse:
+        response = await acompletion(**self._param_parser.parse_nonstream(params))
+        response = cast(LiteLlmModelResponse, response)
+        assistant_message = AssistantMessage.from_litellm_message(response)
+        result: GenerateTextResponse = [assistant_message]
+        async for tool_message in self._resolve_tool_calls(params, assistant_message):
+            result.append(tool_message)
+        return result
+
+    def stream_text_sync(self, params: LlmRequestParams) -> StreamTextResponseSync:
+        """
+        Returns:
+            - stream: Generator yielding `MessageChunk` objects
+            - full_message_queue: Queue containing complete `AssistantMessage`, `ToolMessage` (or `None` when done)
+        """
+        def stream(response: CustomStreamWrapper) -> Generator[MessageChunk]:
+            nonlocal message_collector
+            for chunk in response:
+                chunk = cast(LiteLlmModelResponseStream, chunk)
+                yield from openai_chunk_normalizer(chunk)
+                message_collector.collect(chunk)
+
+            message = message_collector.get_message()
+            full_message_queue.put(message)
+
+            for tool_message in self._resolve_tool_calls_sync(params, message):
+                full_message_queue.put(tool_message)
+            full_message_queue.put(None)
+
+        response = completion(**self._param_parser.parse_stream(params))
+        message_collector = AssistantMessageCollector()
+        returned_stream = stream(cast(CustomStreamWrapper, response))
+        full_message_queue = FullMessageQueueSync()
+        return returned_stream, full_message_queue
+
+    async def stream_text(self, params: LlmRequestParams) -> StreamTextResponseAsync:
+        """
+        Returns:
+            - stream: Generator yielding `MessageChunk` objects
+            - full_message_queue: Queue containing complete `AssistantMessage`, `ToolMessage` (or `None` when done)
+        """
+        async def stream(response: CustomStreamWrapper) -> AsyncGenerator[MessageChunk]:
+            nonlocal message_collector
+            async for chunk in response:
+                chunk = cast(LiteLlmModelResponseStream, chunk)
+                for normalized_chunk in openai_chunk_normalizer(chunk):
+                    yield normalized_chunk
+                message_collector.collect(chunk)
+
+            message = message_collector.get_message()
+            await full_message_queue.put(message)
+            async for tool_message in self._resolve_tool_calls(params, message):
+                await full_message_queue.put(tool_message)
+            await full_message_queue.put(None)
+
+        response = await acompletion(**self._param_parser.parse_stream(params))
+        message_collector = AssistantMessageCollector()
+        returned_stream = stream(cast(CustomStreamWrapper, response))
+        full_message_queue = FullMessageQueueAsync()
+        return returned_stream, full_message_queue
+
+__all__ = [
+    "LLM",
+    "LlmProviders",
+    "LlmRequestParams",
+
+    "Toolset",
+    "python_tool",
+    "PythonToolset",
+    "McpToolset",
+    "LocalMcpToolset",
+    "RemoteMcpToolset",
+
+    "McpClient",
+    "McpTool",
+    "McpToolResult",
+    "LocalMcpClient",
+    "RemoteMcpClient",
+    "LocalServerParams",
+    "RemoteServerParams",
+    "OAuthParams",
+
+    "ToolFn",
+    "ToolDef",
+    "RawToolDef",
+    "ToolLike",
+    "execute_tool",
+    "execute_tool_sync",
+
+    "ChatMessage",
+    "UserMessage",
+    "SystemMessage",
+    "AssistantMessage",
+    "ToolMessage",
+
+    "MessageChunk",
+    "TextChunk",
+    "UsageChunk",
+    "ReasoningChunk",
+    "AudioChunk",
+    "ImageChunk",
+    "ToolCallChunk",
+
+    "GenerateTextResponse",
+    "StreamTextResponseSync",
+    "StreamTextResponseAsync",
+    "FullMessageQueueSync",
+    "FullMessageQueueAsync",
+
+    "enable_debugging",
+    "enable_logging",
+
+    # Exceptions
+    "AuthenticationError",
+    "PermissionDeniedError",
+    "RateLimitError",
+    "ContextWindowExceededError",
+    "BadRequestError",
+    "InvalidRequestError",
+    "InternalServerError",
+    "ServiceUnavailableError",
+    "ContentPolicyViolationError",
+    "APIError",
+    "Timeout",
+]

dais_sdk/debug.py

@@ -0,0 +1,4 @@
+import litellm
+
+def enable_debugging():
+    litellm._turn_on_debug()

dais_sdk/logger.py

@@ -0,0 +1,22 @@
+import sys
+import logging
+
+logger = logging.getLogger("LiteAI-SDK")
+logger.addHandler(logging.NullHandler())
+
+def enable_logging(level=logging.INFO):
+    """
+    Enable logging for the LiteAI SDK.
+
+    Args:
+        level: The logging level (default: logging.INFO).
+
+               Common values: logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR
+    """
+    logger.setLevel(level)
+
+    if not any(isinstance(h, logging.StreamHandler) for h in logger.handlers):
+        handler = logging.StreamHandler(sys.stderr)
+        formatter = logging.Formatter("%(asctime)s - %(name)s - %(message)s")
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)

dais_sdk/mcp_client/__init__.py

@@ -0,0 +1,15 @@
+from .base_mcp_client import McpClient, McpTool, ToolResult as McpToolResult
+from .local_mcp_client import LocalMcpClient, LocalServerParams
+from .remote_mcp_client import RemoteMcpClient, RemoteServerParams, OAuthParams
+
+__all__ = [
+    "McpClient",
+    "McpTool",
+    "McpToolResult",
+
+    "LocalMcpClient",
+    "LocalServerParams",
+    "RemoteMcpClient",
+    "RemoteServerParams",
+    "OAuthParams",
+]

dais_sdk/mcp_client/base_mcp_client.py

@@ -0,0 +1,38 @@
+from abc import ABC, abstractmethod
+from typing import Any, NamedTuple
+from mcp import Tool as McpTool
+from mcp.types import ContentBlock as ToolResultBlock
+
+Tool = McpTool
+
+class ToolResult(NamedTuple):
+    is_error: bool
+    content: list[ToolResultBlock]
+
+class McpClient(ABC):
+    @property
+    @abstractmethod
+    def name(self) -> str: ...
+
+    @abstractmethod
+    async def connect(self): ...
+    @abstractmethod
+    async def disconnect(self): ...
+    @abstractmethod
+    async def list_tools(self) -> list[Tool]:
+        """
+        Raises:
+            McpSessionNotEstablishedError: If the session is not established.
+        """
+    @abstractmethod
+    async def call_tool(
+        self, tool_name: str, arguments: dict[str, Any] | None = None
+    ) -> ToolResult:
+        """
+        Raises:
+            McpSessionNotEstablishedError: If the session is not established.
+        """
+
+class McpSessionNotEstablishedError(RuntimeError):
+    def __init__(self):
+        super().__init__("MCP Session not established, please call connect() first")

dais_sdk/mcp_client/local_mcp_client.py

@@ -0,0 +1,55 @@
+from contextlib import AsyncExitStack
+from typing import Any
+from mcp import ClientSession, StdioServerParameters as StdioServerParams
+from mcp.client.stdio import stdio_client
+from .base_mcp_client import McpClient, Tool, ToolResult, McpSessionNotEstablishedError
+
+LocalServerParams = StdioServerParams
+
+class LocalMcpClient(McpClient):
+    def __init__(self, name: str, params: LocalServerParams):
+        self._name: str = name
+        self._params: LocalServerParams = params
+        self._session: ClientSession | None = None
+        self._exit_stack: AsyncExitStack | None = None
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    async def connect(self):
+        self._exit_stack = AsyncExitStack()
+
+        try:
+            read_stream, write_stream = await self._exit_stack.enter_async_context(
+                stdio_client(self._params)
+            )
+            self._session = await self._exit_stack.enter_async_context(
+                ClientSession(read_stream, write_stream)
+            )
+            await self._session.initialize()
+        except Exception:
+            await self.disconnect()
+            raise
+
+    async def list_tools(self) -> list[Tool]:
+        if not self._session:
+            raise McpSessionNotEstablishedError()
+
+        result = await self._session.list_tools()
+        return result.tools
+
+    async def call_tool(
+        self, tool_name: str, arguments: dict[str, Any] | None = None
+    ) -> ToolResult:
+        if not self._session:
+            raise McpSessionNotEstablishedError()
+
+        response = await self._session.call_tool(tool_name, arguments)
+        return ToolResult(response.isError, response.content)
+
+    async def disconnect(self) -> None:
+        if self._exit_stack:
+            await self._exit_stack.aclose()
+            self._session = None
+            self._exit_stack = None

dais_sdk/mcp_client/oauth_server.py

@@ -0,0 +1,100 @@
+import asyncio
+import socket
+import uvicorn
+from starlette.applications import Starlette
+from starlette.responses import HTMLResponse
+from starlette.routing import Route
+from starlette.requests import Request
+from mcp.client.auth import TokenStorage
+from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
+
+CALLBACK_PAGE = """
+<html>
+    <body style="font-family: sans-serif; text-align: center; padding-top: 50px;">
+        <h2 style="color: green;">Login Successful!</h2>
+        <p>You can close this window now.</p>
+        <script>window.close();</script>
+    </body>
+</html>
+"""
+
+OAuthCode = tuple[str, str | None]
+OAuthCallbackFuture = asyncio.Future[OAuthCode]
+
+class InMemoryTokenStorage(TokenStorage):
+    def __init__(self):
+        self.tokens: OAuthToken | None = None
+        self.client_info: OAuthClientInformationFull | None = None
+
+    async def get_tokens(self) -> OAuthToken | None:
+        return self.tokens
+
+    async def set_tokens(self, tokens: OAuthToken) -> None:
+        self.tokens = tokens
+
+    async def get_client_info(self) -> OAuthClientInformationFull | None:
+        return self.client_info
+
+    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
+        self.client_info = client_info
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("0.0.0.0", 0))
+        port = s.getsockname()[1]
+        return port
+
+class LocalOAuthServer:
+    def __init__(self, timeout: int):
+        self._port = _find_free_port()
+        self._timeout = timeout
+        self._future = OAuthCallbackFuture()
+        self._server: uvicorn.Server | None = None
+        self._server_task: asyncio.Task | None = None
+
+    @property
+    def callback_url(self) -> str:
+        return f"http://localhost:{self._port}/callback"
+
+    async def _handle_callback(self, request: Request):
+        params = request.query_params
+
+        # check if error
+        if "error" in params:
+            error = params.get("error")
+            desc = params.get("error_description", "Unknown error")
+            if not self._future.done():
+                self._future.set_exception(RuntimeError(f"OAuth Error: {error} - {desc}"))
+            return HTMLResponse(f"<h3>Auth failed: {error}</h3>", status_code=400)
+
+        code = params.get("code")
+        state = params.get("state")
+
+        if not code:
+            return HTMLResponse("<h3>Missing 'code' parameter</h3>", status_code=400)
+
+        if not self._future.done():
+            self._future.set_result((code, state))
+
+        return HTMLResponse(CALLBACK_PAGE)
+
+    async def wait_for_code(self) -> OAuthCode:
+        auth_code =  await asyncio.wait_for(self._future, timeout=self._timeout)
+        self._future = OAuthCallbackFuture()
+        return auth_code
+
+    async def start(self):
+        routes = [
+            Route("/callback", self._handle_callback, methods=["GET"])
+        ]
+        app = Starlette(routes=routes)
+
+        config = uvicorn.Config(app=app, host="127.0.0.1", port=self._port, log_level="error")
+        self._server = uvicorn.Server(config)
+        self._server_task = asyncio.create_task(self._server.serve())
+
+    async def stop(self):
+        if self._server:
+            self._server.should_exit = True
+        if self._server_task:
+            await self._server_task