grasp_agents 0.4.6__py3-none-any.whl → 0.5.0__py3-none-any.whl

grasp_agents/openai/openai_llm.py

@@ -1,21 +1,22 @@
+import fnmatch
 import logging
+import os
 from collections.abc import AsyncIterator, Iterable, Mapping
 from copy import deepcopy
-from typing import Any, Literal, NamedTuple
+from typing import Any, Literal
 
 import httpx
-from openai import AsyncOpenAI
+from openai import AsyncOpenAI, AsyncStream
 from openai._types import NOT_GIVEN  # type: ignore[import]
 from openai.lib.streaming.chat import (
     AsyncChatCompletionStreamManager as OpenAIAsyncChatCompletionStreamManager,
 )
+from openai.lib.streaming.chat import ChatCompletionStreamState
 from openai.lib.streaming.chat import ChunkEvent as OpenAIChunkEvent
 from pydantic import BaseModel
 
-from ..cloud_llm import APIProvider, CloudLLM, CloudLLMSettings
+from ..cloud_llm import APIProvider, CloudLLM, CloudLLMSettings, LLMRateLimiter
 from ..http_client import AsyncHTTPClientParams
-from ..rate_limiting.rate_limiter_chunked import RateLimiterC
-from ..typing.message import AssistantMessage, Messages
 from ..typing.tool import BaseTool
 from . import (
     OpenAICompletion,
@@ -23,17 +24,40 @@ from . import (
     OpenAIMessageParam,
     OpenAIParsedCompletion,
     OpenAIPredictionContentParam,
+    OpenAIResponseFormatJSONObject,
+    OpenAIResponseFormatText,
     OpenAIStreamOptionsParam,
     OpenAIToolChoiceOptionParam,
     OpenAIToolParam,
+    OpenAIWebSearchOptions,
 )
 from .converters import OpenAIConverters
 
 logger = logging.getLogger(__name__)
 
 
-class ToolCallSettings(NamedTuple):
-    strict: bool | None = None
+def get_openai_compatible_providers() -> list[APIProvider]:
+    """Returns a dictionary of available OpenAI-compatible API providers."""
+    return [
+        APIProvider(
+            name="openai",
+            base_url="https://api.openai.com/v1",
+            api_key=os.getenv("OPENAI_API_KEY"),
+            response_schema_support=("*",),
+        ),
+        APIProvider(
+            name="gemini_openai",
+            base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
+            api_key=os.getenv("GEMINI_API_KEY"),
+            response_schema_support=("*",),
+        ),
+        APIProvider(
+            name="openrouter",
+            base_url="https://openrouter.ai/api/v1",
+            api_key=os.getenv("OPENROUTER_API_KEY"),
+            response_schema_support=(),
+        ),
+    ]
 
 
 class OpenAILLMSettings(CloudLLMSettings, total=False):
@@ -41,7 +65,7 @@ class OpenAILLMSettings(CloudLLMSettings, total=False):
 
     parallel_tool_calls: bool
 
-    modalities: list[Literal["text", "audio"]] | None
+    modalities: list[Literal["text"]] | None
 
     frequency_penalty: float | None
     presence_penalty: float | None
@@ -50,19 +74,20 @@ class OpenAILLMSettings(CloudLLMSettings, total=False):
     logprobs: bool | None
     top_logprobs: int | None
 
+    stream_options: OpenAIStreamOptionsParam | None
+
     prediction: OpenAIPredictionContentParam | None
 
-    stream_options: OpenAIStreamOptionsParam | None
+    web_search_options: OpenAIWebSearchOptions | None
 
     metadata: dict[str, str] | None
     store: bool | None
     user: str
 
-    # response_format: (
-    #     OpenAIResponseFormatText
-    #     | OpenAIResponseFormatJSONSchema
-    #     | OpenAIResponseFormatJSONObject
-    # )
+    # To support the old JSON mode without respose schemas
+    response_format: OpenAIResponseFormatJSONObject | OpenAIResponseFormatText
+
+    # TODO: support audio
 
 
 class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
@@ -72,52 +97,87 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         model_name: str,
         llm_settings: OpenAILLMSettings | None = None,
         tools: list[BaseTool[BaseModel, Any, Any]] | None = None,
-        response_format: type | Mapping[str, type] | None = None,
+        response_schema: Any | None = None,
+        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
+        apply_response_schema_via_provider: bool = False,
         model_id: str | None = None,
         # Custom LLM provider
         api_provider: APIProvider | None = None,
         # Connection settings
+        max_client_retries: int = 2,
         async_http_client: httpx.AsyncClient | None = None,
         async_http_client_params: (
             dict[str, Any] | AsyncHTTPClientParams | None
         ) = None,
         async_openai_client_params: dict[str, Any] | None = None,
         # Rate limiting
-        rate_limiter: (RateLimiterC[Messages, AssistantMessage] | None) = None,
-        rate_limiter_rpm: float | None = None,
-        rate_limiter_chunk_size: int = 1000,
-        rate_limiter_max_concurrency: int = 300,
-        # Retries
-        num_generation_retries: int = 0,
+        rate_limiter: LLMRateLimiter | None = None,
+        # LLM response retries: try to regenerate to pass validation
+        max_response_retries: int = 1,
     ) -> None:
+        openai_compatible_providers = get_openai_compatible_providers()
+
+        model_name_parts = model_name.split("/", 1)
+        if api_provider is not None:
+            provider_model_name = model_name
+        elif len(model_name_parts) == 2:
+            compat_providers_map = {
+                provider["name"]: provider for provider in openai_compatible_providers
+            }
+            provider_name, provider_model_name = model_name_parts
+            if provider_name not in compat_providers_map:
+                raise ValueError(
+                    f"OpenAI compatible API provider '{provider_name}' "
+                    "is not supported. Supported providers are: "
+                    f"{', '.join(compat_providers_map.keys())}"
+                )
+            api_provider = compat_providers_map[provider_name]
+        else:
+            raise ValueError(
+                "Model name must be in the format 'provider/model_name' or "
+                "you must provide an 'api_provider' argument."
+            )
+
         super().__init__(
-            model_name=model_name,
+            model_name=provider_model_name,
             model_id=model_id,
             llm_settings=llm_settings,
             converters=OpenAIConverters(),
             tools=tools,
-            response_format=response_format,
+            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,
             async_http_client=async_http_client,
             async_http_client_params=async_http_client_params,
             rate_limiter=rate_limiter,
-            rate_limiter_rpm=rate_limiter_rpm,
-            rate_limiter_chunk_size=rate_limiter_chunk_size,
-            rate_limiter_max_concurrency=rate_limiter_max_concurrency,
-            num_generation_retries=num_generation_retries,
+            max_client_retries=max_client_retries,
+            max_response_retries=max_response_retries,
         )
 
-        self._tool_call_settings = {
-            "strict": self._llm_settings.get("use_struct_outputs", False),
-        }
+        response_schema_support: bool = any(
+            fnmatch.fnmatch(self._model_name, pat)
+            for pat in api_provider.get("response_schema_support") or []
+        )
+        if apply_response_schema_via_provider:
+            if self._tools:
+                for tool in self._tools.values():
+                    tool.strict = True
+            if not response_schema_support:
+                raise ValueError(
+                    "Native response schema validation is not supported for model "
+                    f"'{self._model_name}' by the API provider. Please set "
+                    "apply_response_schema_via_provider=False."
+                )
 
         _async_openai_client_params = deepcopy(async_openai_client_params or {})
         if self._async_http_client is not None:
             _async_openai_client_params["http_client"] = self._async_http_client
 
         self._client: AsyncOpenAI = AsyncOpenAI(
-            base_url=self._base_url,
-            api_key=self._api_key,
+            base_url=self.api_provider.get("base_url"),
+            api_key=self.api_provider.get("api_key"),
+            max_retries=max_client_retries,
             **_async_openai_client_params,
         )
 
@@ -126,15 +186,28 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         api_messages: Iterable[OpenAIMessageParam],
         api_tools: list[OpenAIToolParam] | None = None,
         api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
+        api_response_schema: type[Any] | None = None,
         n_choices: int | None = None,
         **api_llm_settings: Any,
-    ) -> OpenAICompletion:
+    ) -> OpenAICompletion | OpenAIParsedCompletion[Any]:
         tools = api_tools or NOT_GIVEN
         tool_choice = api_tool_choice or NOT_GIVEN
+        response_format = api_response_schema or NOT_GIVEN
         n = n_choices or NOT_GIVEN
 
+        if self._apply_response_schema_via_provider:
+            return await self._client.beta.chat.completions.parse(
+                model=self._model_name,
+                messages=api_messages,
+                tools=tools,
+                tool_choice=tool_choice,
+                response_format=response_format,
+                n=n,
+                **api_llm_settings,
+            )
+
         return await self._client.chat.completions.create(
-            model=self._api_model_name,
+            model=self._model_name,
             messages=api_messages,
             tools=tools,
             tool_choice=tool_choice,
@@ -143,89 +216,68 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
             **api_llm_settings,
         )
 
-    async def _get_parsed_completion(
-        self,
-        api_messages: Iterable[OpenAIMessageParam],
-        api_tools: list[OpenAIToolParam] | None = None,
-        api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
-        api_response_format: type | None = None,
-        n_choices: int | None = None,
-        **api_llm_settings: Any,
-    ) -> OpenAIParsedCompletion[Any]:
-        tools = api_tools or NOT_GIVEN
-        tool_choice = api_tool_choice or NOT_GIVEN
-        n = n_choices or NOT_GIVEN
-        response_format = api_response_format or NOT_GIVEN
-
-        return await self._client.beta.chat.completions.parse(
-            model=self._api_model_name,
-            messages=api_messages,
-            tools=tools,
-            tool_choice=tool_choice,
-            response_format=response_format,
-            n=n,
-            **api_llm_settings,
-        )
-
-    async def _get_completion_stream(
+    async def _get_completion_stream(  # type: ignore[override]
         self,
         api_messages: Iterable[OpenAIMessageParam],
         api_tools: list[OpenAIToolParam] | None = None,
         api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
+        api_response_schema: type[Any] | None = None,
         n_choices: int | None = None,
         **api_llm_settings: Any,
     ) -> AsyncIterator[OpenAICompletionChunk]:
         tools = api_tools or NOT_GIVEN
         tool_choice = api_tool_choice or NOT_GIVEN
+        response_format = api_response_schema or NOT_GIVEN
         n = n_choices or NOT_GIVEN
 
-        stream_generator = await self._client.chat.completions.create(
-            model=self._api_model_name,
-            messages=api_messages,
-            tools=tools,
-            tool_choice=tool_choice,
-            stream=True,
-            n=n,
-            **api_llm_settings,
-        )
-
-        async def iterate() -> AsyncIterator[OpenAICompletionChunk]:
+        if self._apply_response_schema_via_provider:
+            stream_manager: OpenAIAsyncChatCompletionStreamManager[Any] = (
+                self._client.beta.chat.completions.stream(
+                    model=self._model_name,
+                    messages=api_messages,
+                    tools=tools,
+                    tool_choice=tool_choice,
+                    response_format=response_format,
+                    n=n,
+                    **api_llm_settings,
+                )
+            )
+            async with stream_manager as stream:
+                async for chunk_event in stream:
+                    if isinstance(chunk_event, OpenAIChunkEvent):
+                        yield chunk_event.chunk
+        else:
+            stream_generator: AsyncStream[
+                OpenAICompletionChunk
+            ] = await self._client.chat.completions.create(
+                model=self._model_name,
+                messages=api_messages,
+                tools=tools,
+                tool_choice=tool_choice,
+                stream=True,
+                n=n,
+                **api_llm_settings,
+            )
             async with stream_generator as stream:
                 async for completion_chunk in stream:
                     yield completion_chunk
 
-        return iterate()
-
-    async def _get_parsed_completion_stream(
-        self,
-        api_messages: Iterable[OpenAIMessageParam],
-        api_tools: list[OpenAIToolParam] | None = None,
-        api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
-        api_response_format: type | None = None,
-        n_choices: int | None = None,
-        **api_llm_settings: Any,
-    ) -> AsyncIterator[OpenAICompletionChunk]:
-        tools = api_tools or NOT_GIVEN
-        tool_choice = api_tool_choice or NOT_GIVEN
-        response_format = api_response_format or NOT_GIVEN
-        n = n_choices or NOT_GIVEN
-
-        stream_manager: OpenAIAsyncChatCompletionStreamManager[
-            OpenAICompletionChunk
-        ] = self._client.beta.chat.completions.stream(
-            model=self._api_model_name,
-            messages=api_messages,
-            tools=tools,
-            tool_choice=tool_choice,
-            response_format=response_format,
-            n=n,
-            **api_llm_settings,
+    def combine_completion_chunks(
+        self, completion_chunks: list[OpenAICompletionChunk]
+    ) -> OpenAICompletion:
+        response_format = NOT_GIVEN
+        input_tools = NOT_GIVEN
+        if self._apply_response_schema_via_provider:
+            if self._response_schema:
+                response_format = self._response_schema
+            if self._tools:
+                input_tools = [
+                    self._converters.to_tool(tool) for tool in self._tools.values()
+                ]
+        state = ChatCompletionStreamState[Any](
+            input_tools=input_tools, response_format=response_format
         )
+        for chunk in completion_chunks:
+            state.handle_chunk(chunk)
 
-        async def iterate() -> AsyncIterator[OpenAICompletionChunk]:
-            async with stream_manager as stream:
-                async for chunk_event in stream:
-                    if isinstance(chunk_event, OpenAIChunkEvent):
-                        yield chunk_event.chunk
-
-        return iterate()
+        return state.get_final_completion()

grasp_agents/openai/tool_converters.py

@@ -13,10 +13,8 @@ from . import (
 )
 
 
-def to_api_tool(
-    tool: BaseTool[BaseModel, Any, Any], strict: bool | None = None
-) -> OpenAIToolParam:
-    if strict:
+def to_api_tool(tool: BaseTool[BaseModel, Any, Any]) -> OpenAIToolParam:
+    if tool.strict:
         return pydantic_function_tool(
             model=tool.in_type, name=tool.name, description=tool.description
         )
@@ -25,9 +23,9 @@ def to_api_tool(
         name=tool.name,
         description=tool.description,
         parameters=tool.in_type.model_json_schema(),
-        strict=strict,
+        strict=tool.strict,
     )
-    if strict is None:
+    if tool.strict is None:
         function.pop("strict")
 
     return OpenAIToolParam(type="function", function=function)

grasp_agents/packet.py

@@ -15,10 +15,13 @@ class Packet(BaseModel, Generic[_PayloadT_co]):
     sender: ProcName
     recipients: Sequence[ProcName] = Field(default_factory=list)
 
-    model_config = ConfigDict(extra="forbid", frozen=True)
+    model_config = ConfigDict(extra="forbid")
 
     def __repr__(self) -> str:
         return (
-            f"From: {self.sender}, To: {', '.join(self.recipients)}, "
+            f"{self.__class__.__name__}:\n"
+            f"ID: {self.id}\n"
+            f"From: {self.sender}\n"
+            f"To: {', '.join(self.recipients)}\n"
             f"Payloads: {len(self.payloads)}"
         )

grasp_agents/packet_pool.py

@@ -1,9 +1,11 @@
 import asyncio
 import logging
+from collections.abc import AsyncIterator
 from typing import Any, Generic, Protocol, TypeVar
 
 from .packet import Packet
 from .run_context import CtxT, RunContext
+from .typing.events import Event
 from .typing.io import ProcName
 
 logger = logging.getLogger(__name__)
@@ -16,16 +18,16 @@ class PacketHandler(Protocol[_PayloadT_contra, CtxT]):
     async def __call__(
         self,
         packet: Packet[_PayloadT_contra],
-        ctx: RunContext[CtxT] | None,
+        ctx: RunContext[CtxT],
         **kwargs: Any,
-    ) -> None: ...
+    ) -> AsyncIterator[Event[Any]] | None: ...
 
 
 class PacketPool(Generic[CtxT]):
     def __init__(self) -> None:
         self._queues: dict[ProcName, asyncio.Queue[Packet[Any]]] = {}
         self._packet_handlers: dict[ProcName, PacketHandler[Any, CtxT]] = {}
-        self._tasks: dict[ProcName, asyncio.Task[None]] = {}
+        self._tasks: dict[ProcName, asyncio.Task[AsyncIterator[Event[Any]] | None]] = {}
 
     async def post(self, packet: Packet[Any]) -> None:
         for recipient_id in packet.recipients:
@@ -36,7 +38,7 @@ class PacketPool(Generic[CtxT]):
         self,
         processor_name: ProcName,
         handler: PacketHandler[Any, CtxT],
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
         **run_kwargs: Any,
     ) -> None:
         self._packet_handlers[processor_name] = handler
@@ -47,11 +49,8 @@ class PacketPool(Generic[CtxT]):
             )
 
     async def _handle_packets(
-        self,
-        processor_name: ProcName,
-        ctx: RunContext[CtxT] | None = None,
-        **run_kwargs: Any,
-    ) -> None:
+        self, processor_name: ProcName, ctx: RunContext[CtxT], **run_kwargs: Any
+    ) -> AsyncIterator[Event[Any]] | None:
         queue = self._queues[processor_name]
         while True:
             try:
@@ -59,11 +58,13 @@ class PacketPool(Generic[CtxT]):
                 handler = self._packet_handlers.get(processor_name)
                 if handler is None:
                     break
-
                 try:
-                    await self._packet_handlers[processor_name](
-                        packet, ctx=ctx, **run_kwargs
-                    )
+                    if ctx.is_streaming:
+                        async for event in handler(packet, ctx=ctx, **run_kwargs):  # type: ignore[return-value]
+                            yield event
+                    else:
+                        await handler(packet, ctx=ctx, **run_kwargs)
+
                 except Exception:
                     logger.exception(f"Error handling packet for {processor_name}")