grasp_agents 0.4.7__py3-none-any.whl → 0.5.1__py3-none-any.whl

grasp_agents/litellm/converters.py

@@ -0,0 +1,138 @@
+from collections.abc import Iterable
+from typing import Any
+
+from pydantic import BaseModel
+
+from ..openai.content_converters import from_api_content, to_api_content
+from ..openai.message_converters import (
+    from_api_system_message,
+    from_api_tool_message,
+    from_api_user_message,
+    to_api_system_message,
+    to_api_tool_message,
+    to_api_user_message,
+)
+from ..openai.tool_converters import to_api_tool, to_api_tool_choice
+from ..typing.completion import Completion, Usage
+from ..typing.completion_chunk import CompletionChunk
+from ..typing.content import Content
+from ..typing.converters import Converters
+from ..typing.message import AssistantMessage, SystemMessage, ToolMessage, UserMessage
+from ..typing.tool import BaseTool, ToolChoice
+from . import (
+    LiteLLMCompletion,
+    LiteLLMCompletionChunk,
+    LiteLLMCompletionMessage,
+    LiteLLMUsage,
+    OpenAIContentPartParam,
+    OpenAISystemMessageParam,
+    OpenAIToolChoiceOptionParam,
+    OpenAIToolMessageParam,
+    OpenAIToolParam,
+    OpenAIUserMessageParam,
+)
+from .completion_chunk_converters import from_api_completion_chunk
+from .completion_converters import (
+    from_api_completion,
+    from_api_completion_usage,
+    to_api_completion,
+)
+from .message_converters import from_api_assistant_message, to_api_assistant_message
+
+
+class LiteLLMConverters(Converters):
+    @staticmethod
+    def to_completion(completion: Completion, **kwargs: Any) -> LiteLLMCompletion:
+        return to_api_completion(completion, **kwargs)
+
+    @staticmethod
+    def from_completion(
+        raw_completion: LiteLLMCompletion, name: str | None = None, **kwargs: Any
+    ) -> Completion:
+        return from_api_completion(raw_completion, name=name, **kwargs)
+
+    @staticmethod
+    def to_completion_chunk(
+        chunk: CompletionChunk, **kwargs: Any
+    ) -> LiteLLMCompletionChunk:
+        raise NotImplementedError
+
+    @staticmethod
+    def from_completion_chunk(
+        raw_chunk: LiteLLMCompletionChunk, name: str | None = None, **kwargs: Any
+    ) -> CompletionChunk:
+        return from_api_completion_chunk(raw_chunk, name=name, **kwargs)
+
+    @staticmethod
+    def from_assistant_message(
+        raw_message: LiteLLMCompletionMessage, name: str | None = None, **kwargs: Any
+    ) -> AssistantMessage:
+        return from_api_assistant_message(raw_message, name=name, **kwargs)
+
+    @staticmethod
+    def to_assistant_message(
+        assistant_message: AssistantMessage, **kwargs: Any
+    ) -> LiteLLMCompletionMessage:
+        return to_api_assistant_message(assistant_message, **kwargs)
+
+    # The remaining converters are the same as OpenAIConverters
+
+    @staticmethod
+    def to_system_message(
+        system_message: SystemMessage, **kwargs: Any
+    ) -> OpenAISystemMessageParam:
+        return to_api_system_message(system_message, **kwargs)
+
+    @staticmethod
+    def from_system_message(
+        raw_message: OpenAISystemMessageParam, name: str | None = None, **kwargs: Any
+    ) -> SystemMessage:
+        return from_api_system_message(raw_message, name=name, **kwargs)
+
+    @staticmethod
+    def to_user_message(
+        user_message: UserMessage, **kwargs: Any
+    ) -> OpenAIUserMessageParam:
+        return to_api_user_message(user_message, **kwargs)
+
+    @staticmethod
+    def from_user_message(
+        raw_message: OpenAIUserMessageParam, name: str | None = None, **kwargs: Any
+    ) -> UserMessage:
+        return from_api_user_message(raw_message, name=name, **kwargs)
+
+    @staticmethod
+    def from_completion_usage(raw_usage: LiteLLMUsage, **kwargs: Any) -> Usage:
+        return from_api_completion_usage(raw_usage, **kwargs)
+
+    @staticmethod
+    def to_tool_message(
+        tool_message: ToolMessage, **kwargs: Any
+    ) -> OpenAIToolMessageParam:
+        return to_api_tool_message(tool_message, **kwargs)
+
+    @staticmethod
+    def from_tool_message(
+        raw_message: OpenAIToolMessageParam, name: str | None = None, **kwargs: Any
+    ) -> ToolMessage:
+        return from_api_tool_message(raw_message, name=name, **kwargs)
+
+    @staticmethod
+    def to_tool(tool: BaseTool[BaseModel, Any, Any], **kwargs: Any) -> OpenAIToolParam:
+        return to_api_tool(tool, **kwargs)
+
+    @staticmethod
+    def to_tool_choice(
+        tool_choice: ToolChoice, **kwargs: Any
+    ) -> OpenAIToolChoiceOptionParam:
+        return to_api_tool_choice(tool_choice, **kwargs)
+
+    @staticmethod
+    def to_content(content: Content, **kwargs: Any) -> Iterable[OpenAIContentPartParam]:
+        return to_api_content(content, **kwargs)
+
+    @staticmethod
+    def from_content(
+        raw_content: str | Iterable[OpenAIContentPartParam], **kwargs: Any
+    ) -> Content:
+        return from_api_content(raw_content, **kwargs)

grasp_agents/litellm/lite_llm.py

@@ -0,0 +1,210 @@
+import logging
+from collections.abc import AsyncIterator, Mapping
+from typing import Any, cast
+
+import litellm
+from litellm.litellm_core_utils.get_supported_openai_params import (
+    get_supported_openai_params,  # type: ignore[no-redef]
+)
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+from litellm.types.llms.anthropic import AnthropicThinkingParam
+from litellm.utils import (
+    supports_parallel_function_calling,
+    supports_prompt_caching,
+    supports_reasoning,
+    supports_response_schema,
+    supports_tool_choice,
+)
+
+# from openai.lib.streaming.chat import ChunkEvent as OpenAIChunkEvent
+from pydantic import BaseModel
+
+from ..cloud_llm import APIProvider, CloudLLM, LLMRateLimiter
+from ..openai.openai_llm import OpenAILLMSettings
+from ..typing.tool import BaseTool
+from . import (
+    LiteLLMCompletion,
+    LiteLLMCompletionChunk,
+    OpenAIMessageParam,
+    OpenAIToolChoiceOptionParam,
+    OpenAIToolParam,
+)
+from .converters import LiteLLMConverters
+
+logger = logging.getLogger(__name__)
+
+
+class LiteLLMSettings(OpenAILLMSettings, total=False):
+    thinking: AnthropicThinkingParam | None
+
+
+class LiteLLM(CloudLLM[LiteLLMSettings, LiteLLMConverters]):
+    def __init__(
+        self,
+        # Base LLM args
+        model_name: str,
+        model_id: str | None = None,
+        llm_settings: LiteLLMSettings | None = None,
+        tools: list[BaseTool[BaseModel, Any, Any]] | None = None,
+        response_schema: Any | None = None,
+        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
+        apply_response_schema_via_provider: bool = False,
+        # LLM provider
+        api_provider: APIProvider | None = None,
+        # deployment_id: str | None = None,
+        # api_version: str | None = None,
+        # Connection settings
+        timeout: float | None = None,
+        max_client_retries: int = 2,
+        # Rate limiting
+        rate_limiter: LLMRateLimiter | None = None,
+        # Drop unsupported LLM settings
+        drop_params: bool = True,
+        additional_drop_params: list[str] | None = None,
+        allowed_openai_params: list[str] | None = None,
+        # Mock LLM response for testing
+        mock_response: str | None = None,
+        # LLM response retries: try to regenerate to pass validation
+        max_response_retries: int = 1,
+    ) -> None:
+        self._lite_llm_completion_params: dict[str, Any] = {
+            "max_retries": max_client_retries,
+            "timeout": timeout,
+            "drop_params": drop_params,
+            "additional_drop_params": additional_drop_params,
+            "allowed_openai_params": allowed_openai_params,
+            "mock_response": mock_response,
+            # "deployment_id": deployment_id,
+            # "api_version": api_version,
+        }
+
+        if model_name in litellm.get_valid_models():  # type: ignore[no-untyped-call]
+            _, provider_name, _, _ = litellm.get_llm_provider(model_name)  # type: ignore[no-untyped-call]
+            api_provider = APIProvider(name=provider_name)
+        elif api_provider is not None:
+            self._lite_llm_completion_params["api_key"] = api_provider.get("api_key")
+            self._lite_llm_completion_params["api_base"] = api_provider.get("api_base")
+        elif api_provider is None:
+            raise ValueError(
+                f"Model '{model_name}' is not supported by LiteLLM and no API provider "
+                "was specified. Please provide a valid API provider or use a different "
+                "model."
+            )
+        super().__init__(
+            model_name=model_name,
+            model_id=model_id,
+            llm_settings=llm_settings,
+            converters=LiteLLMConverters(),
+            tools=tools,
+            response_schema=response_schema,
+            response_schema_by_xml_tag=response_schema_by_xml_tag,
+            apply_response_schema_via_provider=apply_response_schema_via_provider,
+            api_provider=api_provider,
+            rate_limiter=rate_limiter,
+            max_client_retries=max_client_retries,
+            max_response_retries=max_response_retries,
+        )
+
+        if self._apply_response_schema_via_provider:
+            if self._tools:
+                for tool in self._tools.values():
+                    tool.strict = True
+            if not self.supports_response_schema:
+                raise ValueError(
+                    f"Model '{self._model_name}' does not support response schema "
+                    "natively. Please set `apply_response_schema_via_provider=False`"
+                )
+
+    def get_supported_openai_params(self) -> list[Any] | None:
+        return get_supported_openai_params(  # type: ignore[no-untyped-call]
+            model=self._model_name, request_type="chat_completion"
+        )
+
+    @property
+    def supports_reasoning(self) -> bool:
+        return supports_reasoning(model=self._model_name)
+
+    @property
+    def supports_parallel_function_calling(self) -> bool:
+        return supports_parallel_function_calling(model=self._model_name)
+
+    @property
+    def supports_prompt_caching(self) -> bool:
+        return supports_prompt_caching(model=self._model_name)
+
+    @property
+    def supports_response_schema(self) -> bool:
+        return supports_response_schema(model=self._model_name)
+
+    @property
+    def supports_tool_choice(self) -> bool:
+        return supports_tool_choice(model=self._model_name)
+
+    # # client
+    # model_list: Optional[list] = (None,)  # pass in a list of api_base,keys, etc.
+
+    async def _get_completion(
+        self,
+        api_messages: list[OpenAIMessageParam],
+        api_tools: list[OpenAIToolParam] | None = None,
+        api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
+        api_response_schema: type | None = None,
+        n_choices: int | None = None,
+        **api_llm_settings: Any,
+    ) -> LiteLLMCompletion:
+        completion = await litellm.acompletion(  # type: ignore[no-untyped-call]
+            model=self._model_name,
+            messages=api_messages,
+            tools=api_tools,
+            tool_choice=api_tool_choice,  # type: ignore[arg-type]
+            response_format=api_response_schema,
+            n=n_choices,
+            stream=False,
+            **self._lite_llm_completion_params,
+            **api_llm_settings,
+        )
+        completion = cast("LiteLLMCompletion", completion)
+
+        # Should not be needed in litellm>=1.74
+        completion._hidden_params["response_cost"] = litellm.completion_cost(completion)  # type: ignore[no-untyped-call]
+
+        return completion
+
+    async def _get_completion_stream(  # type: ignore[no-untyped-def]
+        self,
+        api_messages: list[OpenAIMessageParam],
+        api_tools: list[OpenAIToolParam] | None = None,
+        api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
+        api_response_schema: type | None = None,
+        n_choices: int | None = None,
+        **api_llm_settings: Any,
+    ) -> AsyncIterator[LiteLLMCompletionChunk]:
+        stream = await litellm.acompletion(  # type: ignore[no-untyped-call]
+            model=self._model_name,
+            messages=api_messages,
+            tools=api_tools,
+            tool_choice=api_tool_choice,  # type: ignore[arg-type]
+            response_format=api_response_schema,
+            stream=True,
+            n=n_choices,
+            **self._lite_llm_completion_params,
+            **api_llm_settings,
+        )
+        stream = cast("CustomStreamWrapper", stream)
+
+        async for completion_chunk in stream:
+            yield completion_chunk
+
+    def combine_completion_chunks(
+        self, completion_chunks: list[LiteLLMCompletionChunk]
+    ) -> LiteLLMCompletion:
+        combined_chunk = cast(
+            "LiteLLMCompletion",
+            litellm.stream_chunk_builder(completion_chunks),  # type: ignore[no-untyped-call]
+        )
+        # Should not be needed in litellm>=1.74
+        combined_chunk._hidden_params["response_cost"] = litellm.completion_cost(  # type: ignore[no-untyped-call]
+            combined_chunk
+        )
+
+        return combined_chunk

grasp_agents/litellm/message_converters.py

@@ -0,0 +1,66 @@
+from ..typing.message import (
+    AssistantMessage,
+)
+from ..typing.tool import ToolCall
+from . import LiteLLMCompletionMessage, LiteLLMFunction, LiteLLMToolCall
+
+
+def from_api_assistant_message(
+    api_message: LiteLLMCompletionMessage, name: str | None = None
+) -> AssistantMessage:
+    tool_calls = None
+    if api_message.tool_calls is not None:
+        tool_calls = [
+            ToolCall(
+                id=tool_call.id,
+                tool_name=tool_call.function.name,  # type: ignore
+                tool_arguments=tool_call.function.arguments,
+            )
+            for tool_call in api_message.tool_calls
+        ]
+
+    return AssistantMessage(
+        content=api_message.content,
+        tool_calls=tool_calls,
+        name=name,
+        thinking_blocks=getattr(api_message, "thinking_blocks", None),
+        reasoning_content=getattr(api_message, "reasoning_content", None),
+        annotations=getattr(api_message, "annotations", None),
+        provider_specific_fields=api_message.provider_specific_fields,
+        refusal=getattr(api_message, "refusal", None),
+    )
+
+
+def to_api_assistant_message(
+    message: AssistantMessage,
+) -> LiteLLMCompletionMessage:
+    api_tool_calls = None
+    if message.tool_calls is not None:
+        api_tool_calls = [
+            LiteLLMToolCall(
+                type="function",
+                id=tool_call.id,
+                function=LiteLLMFunction(
+                    name=tool_call.tool_name,
+                    arguments=tool_call.tool_arguments,
+                ),
+            )
+            for tool_call in message.tool_calls
+        ]
+
+    api_message = LiteLLMCompletionMessage(role="assistant", content=message.content)
+
+    if api_tool_calls:
+        api_message.tool_calls = api_tool_calls
+
+    for key in [
+        "thinking_blocks",
+        "reasoning_content",
+        "annotations",
+        "provider_specific_fields",
+        "refusal",
+    ]:
+        if getattr(message, key):
+            api_message[key] = getattr(message, key)
+
+    return api_message

grasp_agents/llm.py

@@ -1,18 +1,25 @@
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Mapping
+from collections.abc import AsyncIterator, Mapping, Sequence
 from typing import Any, Generic, TypeVar, cast
 from uuid import uuid4
 
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel
 from typing_extensions import TypedDict
 
-from grasp_agents.utils import validate_obj_from_json_or_py_string
+from grasp_agents.utils import (
+    validate_obj_from_json_or_py_string,
+    validate_tagged_objs_from_json_or_py_string,
+)
 
-from .errors import ToolValidationError
+from .errors import (
+    JSONSchemaValidationError,
+    LLMResponseValidationError,
+    LLMToolCallValidationError,
+)
 from .typing.completion import Completion
 from .typing.converters import Converters
-from .typing.events import CompletionChunkEvent, CompletionEvent
+from .typing.events import CompletionChunkEvent, CompletionEvent, LLMStreamingErrorEvent
 from .typing.message import Messages
 from .typing.tool import BaseTool, ToolChoice
 
@@ -38,8 +45,9 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         model_name: str | None = None,
         model_id: str | None = None,
         llm_settings: SettingsT_co | None = None,
-        tools: list[BaseTool[BaseModel, Any, Any]] | None = None,
-        response_format: Any | Mapping[str, Any] | None = None,
+        tools: Sequence[BaseTool[BaseModel, Any, Any]] | None = None,
+        response_schema: Any | None = None,
+        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__()
@@ -50,20 +58,13 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         self._tools = {t.name: t for t in tools} if tools else None
         self._llm_settings: SettingsT_co = llm_settings or cast("SettingsT_co", {})
 
-        self._response_format = response_format
-        self._response_format_adapter: (
-            TypeAdapter[Any] | Mapping[str, TypeAdapter[Any]]
-        ) = self._get_response_format_adapter(response_format=response_format)
-
-    @staticmethod
-    def _get_response_format_adapter(
-        response_format: Any | Mapping[str, Any] | None = None,
-    ) -> TypeAdapter[Any] | Mapping[str, TypeAdapter[Any]]:
-        if response_format is None:
-            return TypeAdapter(Any)
-        if isinstance(response_format, Mapping):
-            return {k: TypeAdapter(v) for k, v in response_format.items()}  # type: ignore[return-value]
-        return TypeAdapter(response_format)
+        if response_schema and response_schema_by_xml_tag:
+            raise ValueError(
+                "Only one of response_schema and response_schema_by_xml_tag can be "
+                "provided, but not both."
+            )
+        self._response_schema = response_schema
+        self._response_schema_by_xml_tag = response_schema_by_xml_tag
 
     @property
     def model_id(self) -> str:
@@ -78,40 +79,59 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         return self._llm_settings
 
     @property
-    def response_format(self) -> Any | Mapping[str, Any] | None:
-        return self._response_format
+    def response_schema(self) -> Any | None:
+        return self._response_schema
 
-    @response_format.setter
-    def response_format(self, response_format: Any | Mapping[str, Any] | None) -> None:
-        self._response_format = response_format
-        self._response_format_adapter = self._get_response_format_adapter(
-            response_format
-        )
+    @response_schema.setter
+    def response_schema(self, response_schema: Any | None) -> None:
+        self._response_schema = response_schema
+
+    @property
+    def response_schema_by_xml_tag(self) -> Mapping[str, Any] | None:
+        return self._response_schema_by_xml_tag
 
     @property
     def tools(self) -> dict[str, BaseTool[BaseModel, Any, Any]] | None:
         return self._tools
 
     @tools.setter
-    def tools(self, tools: list[BaseTool[BaseModel, Any, Any]] | None) -> None:
+    def tools(self, tools: Sequence[BaseTool[BaseModel, Any, Any]] | None) -> None:
         self._tools = {t.name: t for t in tools} if tools else None
 
     def __repr__(self) -> str:
-        return (
-            f"{type(self).__name__}(model_id={self.model_id}; "
-            f"model_name={self._model_name})"
-        )
+        return f"{type(self).__name__}[{self.model_id}]; model_name={self._model_name})"
+
+    def _validate_response(self, completion: Completion) -> None:
+        parsing_params = {
+            "from_substring": False,
+            "strip_language_markdown": True,
+        }
+        try:
+            for message in completion.messages:
+                if not message.tool_calls:
+                    if self._response_schema:
+                        validate_obj_from_json_or_py_string(
+                            message.content or "",
+                            schema=self._response_schema,
+                            **parsing_params,
+                        )
 
-    def _validate_completion(self, completion: Completion) -> None:
-        for message in completion.messages:
-            if not message.tool_calls:
-                validate_obj_from_json_or_py_string(
-                    message.content or "",
-                    adapter=self._response_format_adapter,
-                    from_substring=True,
-                )
+                    elif self._response_schema_by_xml_tag:
+                        validate_tagged_objs_from_json_or_py_string(
+                            message.content or "",
+                            schema_by_xml_tag=self._response_schema_by_xml_tag,
+                            **parsing_params,
+                        )
+        except JSONSchemaValidationError as exc:
+            raise LLMResponseValidationError(
+                exc.s, exc.schema, message=str(exc)
+            ) from exc
 
     def _validate_tool_calls(self, completion: Completion) -> None:
+        parsing_params = {
+            "from_substring": False,
+            "strip_language_markdown": True,
+        }
         for message in completion.messages:
             if message.tool_calls:
                 for tool_call in message.tool_calls:
@@ -120,14 +140,21 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
 
                     available_tool_names = list(self.tools) if self.tools else []
                     if tool_name not in available_tool_names or not self.tools:
-                        raise ToolValidationError(
-                            f"Tool '{tool_name}' is not available in the LLM tools "
-                            f"(available: {available_tool_names}"
+                        raise LLMToolCallValidationError(
+                            tool_name,
+                            tool_arguments,
+                            message=f"Tool '{tool_name}' is not available in the LLM "
+                            f"tools (available: {available_tool_names})",
                         )
                     tool = self.tools[tool_name]
-                    validate_obj_from_json_or_py_string(
-                        tool_arguments, adapter=TypeAdapter(tool.in_type)
-                    )
+                    try:
+                        validate_obj_from_json_or_py_string(
+                            tool_arguments, schema=tool.in_type, **parsing_params
+                        )
+                    except JSONSchemaValidationError as exc:
+                        raise LLMToolCallValidationError(
+                            tool_name, tool_arguments
+                        ) from exc
 
     @abstractmethod
     async def generate_completion(
@@ -136,6 +163,8 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         *,
         tool_choice: ToolChoice | None = None,
         n_choices: int | None = None,
+        proc_name: str | None = None,
+        call_id: str | None = None,
     ) -> Completion:
         pass
 
@@ -146,5 +175,11 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         *,
         tool_choice: ToolChoice | None = None,
         n_choices: int | None = None,
-    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent]:
+        proc_name: str | None = None,
+        call_id: str | None = None,
+    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent | LLMStreamingErrorEvent]:
         pass
+
+    @abstractmethod
+    def combine_completion_chunks(self, completion_chunks: list[Any]) -> Any:
+        raise NotImplementedError