dais-sdk 0.6.0__tar.gz → 0.6.1__tar.gz

{dais_sdk-0.6.0 → dais_sdk-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dais-sdk
-Version: 0.6.0
+Version: 0.6.1
 Summary: A wrapper of LiteLLM
 Author-email: BHznJNs <bhznjns@outlook.com>
 Requires-Python: >=3.10

{dais_sdk-0.6.0 → dais_sdk-0.6.1}/pyproject.toml

@@ -17,7 +17,7 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
 ]
 requires-python = ">=3.10"
-version = "0.6.0"
+version = "0.6.1"
 
 dependencies = [
     "litellm>=1.80.0",

{dais_sdk-0.6.0 → dais_sdk-0.6.1}/src/dais_sdk/__init__.py

@@ -1,5 +1,4 @@
-import asyncio
-import queue
+import json
 from typing import cast
 from collections.abc import AsyncGenerator, Generator
 from litellm import CustomStreamWrapper, completion, acompletion
@@ -12,7 +11,10 @@ from litellm.types.utils import (
 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.execute import (
+    ToolExceptionHandlerManager,
+    execute_tool_sync, execute_tool
+)
 from .tool.toolset import (
     Toolset,
     python_tool,
@@ -21,7 +23,7 @@ from .tool.toolset import (
     LocalMcpToolset,
     RemoteMcpToolset,
 )
-from .tool.utils import find_tool_by_name
+from .tool.utils import get_tool_name
 from .mcp_client import (
     McpClient,
     McpTool,
@@ -51,6 +53,10 @@ from .types.exceptions import (
     ContentPolicyViolationError,
     APIError,
     Timeout,
+    LlmToolException,
+    ToolDoesNotExistError,
+    ToolArgumentDecodeError,
+    ToolExecutionError,
 )
 from .types.message import (
     ChatMessage, UserMessage, SystemMessage, AssistantMessage, ToolMessage,
@@ -87,47 +93,49 @@ class LLM:
         self.provider = provider
         self.base_url = base_url
         self.api_key = api_key
+        self._tool_exception_handler_manager = ToolExceptionHandlerManager()
         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]:
+    @property
+    def tool_exception_handler_manager(self) -> ToolExceptionHandlerManager:
+        return self._tool_exception_handler_manager
+
+    async def execute_tool_call(self,
+                                tool_def: ToolLike,
+                                arguments: str | dict) -> tuple[str | None, str | None]:
         """
-        Receive incomplete tool messages, execute the tool calls and
-        return the result and error tuple.
+        Returns:
+            A tuple of (result, error)
         """
-        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 json.JSONDecodeError as e:
+            assert type(arguments) is str
+            _error = ToolArgumentDecodeError(get_tool_name(tool_def), arguments, e)
+            error = self._tool_exception_handler_manager.handle(_error)
         except Exception as e:
-            error = f"{type(e).__name__}: {str(e)}"
+            _error = ToolExecutionError(tool_def, arguments, e)
+            error = self._tool_exception_handler_manager.handle(_error)
         return result, error
 
-    @staticmethod
-    def execute_tool_call_sync(
-            params: LlmRequestParams,
-            incomplete_tool_message: ToolMessage
-            ) -> tuple[str | None, str | None]:
+    def execute_tool_call_sync(self,
+                               tool_def: ToolLike,
+                               arguments: str | dict
+                               ) -> 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 json.JSONDecodeError as e:
+            assert type(arguments) is str
+            _error = ToolArgumentDecodeError(get_tool_name(tool_def), arguments, e)
+            error = self._tool_exception_handler_manager.handle(_error)
         except Exception as e:
-            error = f"{type(e).__name__}: {str(e)}"
+            _error = ToolExecutionError(tool_def, arguments, e)
+            error = self._tool_exception_handler_manager.handle(_error)
         return result, error
 
     def _resolve_tool_calls_sync(self, params: LlmRequestParams, assistant_message: AssistantMessage) -> Generator[ToolMessage]:
@@ -136,17 +144,14 @@ class LLM:
             := 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.")
+            tool = params.find_tool(incomplete_tool_message.name)
+            if tool is None:
+                _error = ToolDoesNotExistError(incomplete_tool_message.name)
+                error = self._tool_exception_handler_manager.handle(_error)
+                yield incomplete_tool_message.with_result(None, error)
                 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)
+            result, error = self.execute_tool_call_sync(tool, incomplete_tool_message.arguments)
+            yield incomplete_tool_message.with_result(result, error)
 
     async def _resolve_tool_calls(self, params: LlmRequestParams, assistant_message: AssistantMessage) -> AsyncGenerator[ToolMessage]:
         if not params.execute_tools: return
@@ -154,17 +159,14 @@ class LLM:
             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.")
+            tool = params.find_tool(incomplete_tool_message.name)
+            if tool is None:
+                _error = ToolDoesNotExistError(incomplete_tool_message.name)
+                error = self._tool_exception_handler_manager.handle(_error)
+                yield incomplete_tool_message.with_result(None, error)
                 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)
+            result, error = await self.execute_tool_call(tool, incomplete_tool_message.arguments)
+            yield incomplete_tool_message.with_result(result, error)
 
     def list_models(self) -> list[str]:
         provider_config = ProviderConfigManager.get_provider_model_info(
@@ -208,8 +210,8 @@ class LLM:
             - 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
+        def stream(response: CustomStreamWrapper, full_message_queue: FullMessageQueueSync) -> Generator[MessageChunk]:
+            message_collector = AssistantMessageCollector()
             for chunk in response:
                 chunk = cast(LiteLlmModelResponseStream, chunk)
                 yield from openai_chunk_normalizer(chunk)
@@ -223,9 +225,8 @@ class LLM:
             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()
+        returned_stream = stream(cast(CustomStreamWrapper, response), full_message_queue)
         return returned_stream, full_message_queue
 
     async def stream_text(self, params: LlmRequestParams) -> StreamTextResponseAsync:
@@ -234,8 +235,8 @@ class LLM:
             - 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 def stream(response: CustomStreamWrapper, full_message_queue: FullMessageQueueAsync) -> AsyncGenerator[MessageChunk]:
+            message_collector = AssistantMessageCollector()
             async for chunk in response:
                 chunk = cast(LiteLlmModelResponseStream, chunk)
                 for normalized_chunk in openai_chunk_normalizer(chunk):
@@ -244,14 +245,14 @@ class LLM:
 
             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()
+        returned_stream = stream(cast(CustomStreamWrapper, response), full_message_queue)
         return returned_stream, full_message_queue
 
 __all__ = [
@@ -317,4 +318,8 @@ __all__ = [
     "ContentPolicyViolationError",
     "APIError",
     "Timeout",
+    "LlmToolException",
+    "ToolDoesNotExistError",
+    "ToolArgumentDecodeError",
+    "ToolExecutionError",
 ]

{dais_sdk-0.6.0 → dais_sdk-0.6.1}/src/dais_sdk/mcp_client/remote_mcp_client.py

@@ -98,7 +98,7 @@ class RemoteMcpClient(McpClient):
         try:
             webbrowser.open(url)
         except Exception as e:
-            logger.error(f"[OAuth] Not able to open browser: {e}")
+            logger.error(f"[OAuth] Not able to open browser", exc_info=e)
 
     async def _handle_oauth_callback(self) -> OAuthCode:
         if self._oauth_context is None:

{dais_sdk-0.6.0 → dais_sdk-0.6.1}/src/dais_sdk/tool/execute.py

@@ -1,12 +1,42 @@
 import asyncio
 import json
 from functools import singledispatch
-from typing import Any, Awaitable, Callable, cast
-from types import FunctionType, MethodType, CoroutineType
+from typing import Any, assert_never, Callable, cast
+from types import FunctionType, MethodType
 from ..types.tool import ToolDef, ToolLike
+from ..types.exceptions import ToolDoesNotExistError, ToolArgumentDecodeError, ToolExecutionError
+from ..logger import logger
 
-async def _coroutine_wrapper(awaitable: Awaitable[Any]) -> CoroutineType:
-    return await awaitable
+class ToolExceptionHandlerManager:
+    ToolException = ToolDoesNotExistError | ToolArgumentDecodeError | ToolExecutionError
+    Handler = Callable[[ToolException], str]
+
+    def __init__(self):
+        self._handlers: dict[
+            type[ToolExceptionHandlerManager.ToolException],
+            ToolExceptionHandlerManager.Handler
+        ] = {}
+
+    def register(self, exception_type: type[ToolException]):
+        def decorator(handler: ToolExceptionHandlerManager.Handler) -> ToolExceptionHandlerManager.Handler:
+            self.set_handler(exception_type, handler)
+            return handler
+        return decorator
+
+    def set_handler(self, exception_type: type[ToolException], handler: Handler):
+        self._handlers[exception_type] = handler
+
+    def get_handler(self, exception_type: type[ToolException]) -> Handler | None:
+        return self._handlers.get(exception_type)
+
+    def handle(self, e: ToolException) -> str:
+        handler = self.get_handler(type(e))
+        if handler is None:
+            logger.warning(f"Unhandled tool exception: {type(e).__name__}", exc_info=e)
+            return f"Unhandled tool exception | {type(e).__name__}: {e}"
+        return handler(e)
+
+# --- --- --- --- --- ---
 
 def _arguments_normalizer(arguments: str | dict) -> dict:
     if isinstance(arguments, str):
@@ -15,7 +45,7 @@ def _arguments_normalizer(arguments: str | dict) -> dict:
     elif isinstance(arguments, dict):
         return arguments
     else:
-        raise ValueError(f"Invalid arguments type: {type(arguments)}")
+        assert_never(arguments)
 
 def _result_normalizer(result: Any) -> str:
     if isinstance(result, str):
@@ -24,13 +54,18 @@ def _result_normalizer(result: Any) -> str:
 
 @singledispatch
 def execute_tool_sync(tool: ToolLike, arguments: str | dict) -> str:
+    """
+    Raises:
+        ValueError: If the tool type is not supported.
+        JSONDecodeError: If the arguments is a string but not valid JSON.
+    """
     raise ValueError(f"Invalid tool type: {type(tool)}")
 
 @execute_tool_sync.register(FunctionType)
 @execute_tool_sync.register(MethodType)
 def _(toolfn: Callable, arguments: str | dict) -> str:
     arguments = _arguments_normalizer(arguments)
-    result = (asyncio.run(_coroutine_wrapper(toolfn(**arguments)))
+    result = (asyncio.run(toolfn(**arguments))
               if asyncio.iscoroutinefunction(toolfn)
               else toolfn(**arguments))
     return _result_normalizer(result)
@@ -38,13 +73,18 @@ def _(toolfn: Callable, arguments: str | dict) -> str:
 @execute_tool_sync.register(ToolDef)
 def _(tooldef: ToolDef, arguments: str | dict) -> str:
     arguments = _arguments_normalizer(arguments)
-    result = (asyncio.run(_coroutine_wrapper(tooldef.execute(**arguments)))
+    result = (asyncio.run(tooldef.execute(**arguments))
               if asyncio.iscoroutinefunction(tooldef.execute)
               else tooldef.execute(**arguments))
     return _result_normalizer(result)
 
 @singledispatch
 async def execute_tool(tool: ToolLike, arguments: str | dict) -> str:
+    """
+    Raises:
+        ValueError: If the tool type is not supported.
+        JSONDecodeError: If the arguments is a string but not valid JSON.
+    """
     raise ValueError(f"Invalid tool type: {type(tool)}")
 
 @execute_tool.register(FunctionType)

{dais_sdk-0.6.0 → dais_sdk-0.6.1}/src/dais_sdk/tool/utils.py

@@ -1,4 +1,15 @@
-from ..types.tool import ToolFn, ToolDef, ToolLike
+from typing import assert_never
+from ..types.tool import ToolDef, ToolLike
+
+def get_tool_name(tool: ToolLike) -> str:
+    if callable(tool):
+        return tool.__name__
+    elif isinstance(tool, ToolDef):
+        return tool.name
+    elif isinstance(tool, dict):
+        return tool.get("name", "")
+    else:
+        assert_never(tool)
 
 def find_tool_by_name(tools: list[ToolLike], name: str) -> ToolLike | None:
     for tool in tools:

dais_sdk-0.6.1/src/dais_sdk/types/exceptions.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+import json
+from typing import TYPE_CHECKING
+from litellm.exceptions import (  
+    AuthenticationError,
+    PermissionDeniedError,
+    RateLimitError,
+    ContextWindowExceededError,
+    BadRequestError,
+    InvalidRequestError,
+    InternalServerError,
+    ServiceUnavailableError,
+    ContentPolicyViolationError,
+    APIError,
+    Timeout,
+)
+
+if TYPE_CHECKING:
+    from .tool import ToolLike
+
+class LlmToolException(Exception): pass
+
+class ToolDoesNotExistError(LlmToolException):
+    def __init__(self, tool_name: str):
+        self.tool_name = tool_name
+
+class ToolArgumentDecodeError(LlmToolException):
+    def __init__(self, tool_name: str, arguments: str, raw_error: json.JSONDecodeError):
+        self.tool_name = tool_name
+        self.arguments = arguments
+        self.raw_error = raw_error
+
+class ToolExecutionError(LlmToolException):
+    def __init__(self, tool: ToolLike, arguments: str | dict, raw_error: Exception):
+        self.tool = tool
+        self.arguments = arguments
+        self.raw_error = raw_error
+
+__all__ = [
+    "AuthenticationError",
+    "PermissionDeniedError",
+    "RateLimitError",
+    "ContextWindowExceededError",
+    "BadRequestError",
+    "InvalidRequestError",
+    "InternalServerError",
+    "ServiceUnavailableError",
+    "ContentPolicyViolationError",
+    "APIError",
+    "Timeout",
+
+    "LlmToolException",
+    "ToolDoesNotExistError",
+    "ToolArgumentDecodeError",
+    "ToolExecutionError",
+]

{dais_sdk-0.6.0 → dais_sdk-0.6.1}/src/dais_sdk/types/message.py

@@ -59,6 +59,14 @@ class ToolMessage(ChatMessage):
         if isinstance(v, str): return v
         return json.dumps(v, ensure_ascii=False)
 
+    def with_result(self, result: str | None, error: str | None) -> "ToolMessage":
+        return ToolMessage(
+            tool_call_id=self.tool_call_id,
+            name=self.name,
+            arguments=self.arguments,
+            result=result,
+            error=error)
+
     def to_litellm_message(self) -> ChatCompletionToolMessage:
         if self.result is None and self.error is None:
             raise ValueError(f"ToolMessage({self.id}, {self.name}) is incomplete, "

{dais_sdk-0.6.0 → dais_sdk-0.6.1}/src/dais_sdk/types/request_params.py

@@ -47,9 +47,3 @@ class LlmRequestParams:
         if (tools := self.extract_tools()) is None:
             return None
         return find_tool_by_name(tools, tool_name)
-
-    class ToolDoesNotExistError(Exception):
-        def __init__(self, tool_name: str):
-            self.tool_name = tool_name
-            self.message = f"Tool \"{tool_name}\" does not exist in the request params."
-            super().__init__(self.message)

dais_sdk-0.6.0/src/dais_sdk/types/exceptions.py

@@ -1,27 +0,0 @@
-from litellm.exceptions import (  
-    AuthenticationError,
-    PermissionDeniedError,
-    RateLimitError,
-    ContextWindowExceededError,
-    BadRequestError,
-    InvalidRequestError,
-    InternalServerError,
-    ServiceUnavailableError,
-    ContentPolicyViolationError,
-    APIError,
-    Timeout,
-)
-
-__all__ = [
-    "AuthenticationError",
-    "PermissionDeniedError",
-    "RateLimitError",
-    "ContextWindowExceededError",
-    "BadRequestError",
-    "InvalidRequestError",
-    "InternalServerError",
-    "ServiceUnavailableError",
-    "ContentPolicyViolationError",
-    "APIError",
-    "Timeout",
-]