grasp_agents 0.4.7__tar.gz → 0.5.1__tar.gz

{grasp_agents-0.4.7 → grasp_agents-0.5.1}/PKG-INFO

@@ -1,11 +1,12 @@
 Metadata-Version: 2.4
 Name: grasp_agents
-Version: 0.4.7
+Version: 0.5.1
 Summary: Grasp Agents Library
 License-File: LICENSE.md
 Requires-Python: <4,>=3.11.4
 Requires-Dist: dotenv>=0.9.9
 Requires-Dist: httpx<1,>=0.27.0
+Requires-Dist: litellm>=1.74.4
 Requires-Dist: openai<2,>=1.68.2
 Requires-Dist: pydantic>=2
 Requires-Dist: pyyaml>=6.0.2
@@ -103,7 +104,7 @@ Ensure you have a `.env` file with your OpenAI and Google AI Studio API keys set
 
 ```
 OPENAI_API_KEY=your_openai_api_key
-GOOGLE_AI_STUDIO_API_KEY=your_google_ai_studio_api_key
+GEMINI_API_KEY=your_gemini_api_key
 ```
 
 Create a script, e.g., `problem_recommender.py`:
@@ -117,7 +118,7 @@ from dotenv import load_dotenv
 from pydantic import BaseModel, Field
 
 from grasp_agents.grasp_logging import setup_logging
-from grasp_agents.openai import OpenAILLM, OpenAILLMSettings
+from grasp_agents.litellm import LiteLLM, LiteLLMSettings
 from grasp_agents import LLMAgent, BaseTool, RunContext
 
 load_dotenv()
@@ -178,9 +179,9 @@ class Problem(BaseModel):
 
 teacher = LLMAgent[None, Problem, None](
     name="teacher",
-    llm=OpenAILLM(
-        model_name="openai:gpt-4.1",
-        llm_settings=OpenAILLMSettings(temperature=0.5),
+    llm=LiteLLM(
+        model_name="gpt-4.1",
+        llm_settings=LiteLLMSettings(temperature=0.5),
     ),
     tools=[AskStudentTool()],
     react_mode=True,

{grasp_agents-0.4.7 → grasp_agents-0.5.1}/README.md

@@ -87,7 +87,7 @@ Ensure you have a `.env` file with your OpenAI and Google AI Studio API keys set
 
 ```
 OPENAI_API_KEY=your_openai_api_key
-GOOGLE_AI_STUDIO_API_KEY=your_google_ai_studio_api_key
+GEMINI_API_KEY=your_gemini_api_key
 ```
 
 Create a script, e.g., `problem_recommender.py`:
@@ -101,7 +101,7 @@ from dotenv import load_dotenv
 from pydantic import BaseModel, Field
 
 from grasp_agents.grasp_logging import setup_logging
-from grasp_agents.openai import OpenAILLM, OpenAILLMSettings
+from grasp_agents.litellm import LiteLLM, LiteLLMSettings
 from grasp_agents import LLMAgent, BaseTool, RunContext
 
 load_dotenv()
@@ -162,9 +162,9 @@ class Problem(BaseModel):
 
 teacher = LLMAgent[None, Problem, None](
     name="teacher",
-    llm=OpenAILLM(
-        model_name="openai:gpt-4.1",
-        llm_settings=OpenAILLMSettings(temperature=0.5),
+    llm=LiteLLM(
+        model_name="gpt-4.1",
+        llm_settings=LiteLLMSettings(temperature=0.5),
     ),
     tools=[AskStudentTool()],
     react_mode=True,

{grasp_agents-0.4.7 → grasp_agents-0.5.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "grasp_agents"
-version = "0.4.7"
+version = "0.5.1"
 description = "Grasp Agents Library"
 readme = "README.md"
 requires-python = ">=3.11.4,<4"
@@ -13,6 +13,7 @@ dependencies = [
     "dotenv>=0.9.9",
     "pyyaml>=6.0.2",
     "pydantic>=2",
+    "litellm>=1.74.4",
 ]
 
 [dependency-groups]

grasp_agents-0.5.1/src/grasp_agents/cloud_llm.py

@@ -0,0 +1,365 @@
+import logging
+from abc import abstractmethod
+from collections.abc import AsyncIterator, Mapping, Sequence
+from copy import deepcopy
+from typing import Any, Generic, Required, cast
+
+import httpx
+from pydantic import BaseModel
+from typing_extensions import TypedDict
+
+from .errors import LLMResponseValidationError, LLMToolCallValidationError
+from .http_client import AsyncHTTPClientParams, create_simple_async_httpx_client
+from .llm import LLM, ConvertT_co, LLMSettings, SettingsT_co
+from .rate_limiting.rate_limiter_chunked import RateLimiterC, limit_rate
+from .typing.completion import Completion
+from .typing.completion_chunk import CompletionChoice
+from .typing.events import (
+    CompletionChunkEvent,
+    CompletionEvent,
+    LLMStreamingErrorData,
+    LLMStreamingErrorEvent,
+)
+from .typing.message import AssistantMessage, Messages
+from .typing.tool import BaseTool, ToolChoice
+
+logger = logging.getLogger(__name__)
+
+
+class APIProvider(TypedDict, total=False):
+    name: Required[str]
+    base_url: str | None
+    api_key: str | None
+    # Wildcard patterns for model names that support response schema validation:
+    response_schema_support: tuple[str, ...] | None
+
+
+def make_refusal_completion(model_name: str, err: BaseException) -> Completion:
+    failed_message = AssistantMessage(content=None, refusal=str(err))
+
+    return Completion(
+        model=model_name,
+        choices=[CompletionChoice(message=failed_message, finish_reason=None, index=0)],
+    )
+
+
+class CloudLLMSettings(LLMSettings, total=False):
+    extra_headers: dict[str, Any] | None
+    extra_body: object | None
+    extra_query: dict[str, Any] | None
+
+
+LLMRateLimiter = RateLimiterC[
+    Messages,
+    AssistantMessage
+    | AsyncIterator[CompletionChunkEvent | CompletionEvent | LLMStreamingErrorEvent],
+]
+
+
+class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co]):
+    def __init__(
+        self,
+        # Base LLM args
+        model_name: str,
+        api_provider: APIProvider,
+        converters: ConvertT_co,
+        llm_settings: SettingsT_co | 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 = True,
+        model_id: str | None = None,
+        # Connection settings
+        async_http_client: httpx.AsyncClient | None = None,
+        async_http_client_params: (
+            dict[str, Any] | AsyncHTTPClientParams | None
+        ) = None,
+        max_client_retries: int = 2,
+        # Rate limiting
+        rate_limiter: LLMRateLimiter | None = None,
+        # LLM response retries: try to regenerate to pass validation
+        max_response_retries: int = 0,
+        **kwargs: Any,
+    ) -> None:
+        self.llm_settings: CloudLLMSettings | None
+
+        super().__init__(
+            model_name=model_name,
+            llm_settings=llm_settings,
+            converters=converters,
+            model_id=model_id,
+            tools=tools,
+            response_schema=response_schema,
+            response_schema_by_xml_tag=response_schema_by_xml_tag,
+            **kwargs,
+        )
+
+        self._model_name = model_name
+        self._api_provider = api_provider
+        self._apply_response_schema_via_provider = apply_response_schema_via_provider
+
+        if (
+            apply_response_schema_via_provider
+            and response_schema_by_xml_tag is not None
+        ):
+            raise ValueError(
+                "Response schema by XML tag is not supported "
+                "when apply_response_schema_via_provider is True."
+            )
+
+        self._rate_limiter: LLMRateLimiter | None = None
+        if rate_limiter is not None:
+            self._rate_limiter = rate_limiter
+            logger.info(
+                f"[{self.__class__.__name__}] Set rate limit to {rate_limiter.rpm} RPM"
+            )
+
+        self._async_http_client: httpx.AsyncClient | None = None
+        if async_http_client is not None:
+            self._async_http_client = async_http_client
+        elif async_http_client_params is not None:
+            self._async_http_client = create_simple_async_httpx_client(
+                async_http_client_params
+            )
+
+        self.max_client_retries = max_client_retries
+        self.max_response_retries = max_response_retries
+
+    @property
+    def api_provider(self) -> APIProvider:
+        return self._api_provider
+
+    @property
+    def rate_limiter(self) -> LLMRateLimiter | None:
+        return self._rate_limiter
+
+    @property
+    def tools(self) -> dict[str, BaseTool[BaseModel, Any, Any]] | None:
+        return self._tools
+
+    @tools.setter
+    def tools(self, tools: Sequence[BaseTool[BaseModel, Any, Any]] | None) -> None:
+        if not tools:
+            self._tools = None
+            return
+        strict_value = True if self._apply_response_schema_via_provider else None
+        for t in tools:
+            t.strict = strict_value
+        self._tools = {t.name: t for t in tools}
+
+    def _make_completion_kwargs(
+        self,
+        conversation: Messages,
+        tool_choice: ToolChoice | None = None,
+        n_choices: int | None = None,
+    ) -> dict[str, Any]:
+        api_messages = [self._converters.to_message(m) for m in conversation]
+
+        api_tools = None
+        api_tool_choice = None
+        if self.tools:
+            api_tools = [self._converters.to_tool(t) for t in self.tools.values()]
+            if tool_choice is not None:
+                api_tool_choice = self._converters.to_tool_choice(tool_choice)
+
+        api_llm_settings = deepcopy(self.llm_settings or {})
+
+        return dict(
+            api_messages=api_messages,
+            api_tools=api_tools,
+            api_tool_choice=api_tool_choice,
+            api_response_schema=self._response_schema,
+            n_choices=n_choices,
+            **api_llm_settings,
+        )
+
+    @abstractmethod
+    async def _get_completion(
+        self,
+        api_messages: list[Any],
+        *,
+        api_tools: list[Any] | None = None,
+        api_tool_choice: Any | None = None,
+        api_response_schema: type | None = None,
+        n_choices: int | None = None,
+        **api_llm_settings: Any,
+    ) -> Any:
+        pass
+
+    @abstractmethod
+    async def _get_completion_stream(
+        self,
+        api_messages: list[Any],
+        *,
+        api_tools: list[Any] | None = None,
+        api_tool_choice: Any | None = None,
+        api_response_schema: type | None = None,
+        n_choices: int | None = None,
+        **api_llm_settings: Any,
+    ) -> AsyncIterator[Any]:
+        pass
+
+    @limit_rate
+    async def _generate_completion_once(
+        self,
+        conversation: Messages,
+        *,
+        tool_choice: ToolChoice | None = None,
+        n_choices: int | None = None,
+    ) -> Completion:
+        completion_kwargs = self._make_completion_kwargs(
+            conversation=conversation, tool_choice=tool_choice, n_choices=n_choices
+        )
+
+        if not self._apply_response_schema_via_provider:
+            completion_kwargs.pop("api_response_schema", None)
+        api_completion = await self._get_completion(**completion_kwargs)
+
+        completion = self._converters.from_completion(
+            api_completion, name=self.model_id
+        )
+
+        if not self._apply_response_schema_via_provider:
+            self._validate_response(completion)
+            self._validate_tool_calls(completion)
+
+        return completion
+
+    async def generate_completion(
+        self,
+        conversation: Messages,
+        *,
+        tool_choice: ToolChoice | None = None,
+        n_choices: int | None = None,
+        proc_name: str | None = None,
+        call_id: str | None = None,
+    ) -> Completion:
+        n_attempt = 0
+        while n_attempt <= self.max_response_retries:
+            try:
+                return await self._generate_completion_once(
+                    conversation,  # type: ignore[return]
+                    tool_choice=tool_choice,
+                    n_choices=n_choices,
+                )
+            except (LLMResponseValidationError, LLMToolCallValidationError) as err:
+                n_attempt += 1
+
+                if n_attempt > self.max_response_retries:
+                    if n_attempt == 1:
+                        logger.warning(f"\nCloudLLM completion request failed:\n{err}")
+                    if n_attempt > 1:
+                        logger.warning(
+                            f"\nCloudLLM completion request failed after retrying:\n{err}"
+                        )
+                    raise err
+                    # return make_refusal_completion(self._model_name, err)
+
+                logger.warning(
+                    f"\nCloudLLM completion request failed (retry attempt {n_attempt}):"
+                    f"\n{err}"
+                )
+
+        return make_refusal_completion(
+            self._model_name,
+            Exception("Unexpected error: retry loop exited without returning"),
+        )
+
+    @limit_rate
+    async def _generate_completion_stream_once(
+        self,
+        conversation: Messages,
+        *,
+        tool_choice: ToolChoice | None = None,
+        n_choices: int | None = None,
+        proc_name: str | None = None,
+        call_id: str | None = None,
+    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent]:
+        completion_kwargs = self._make_completion_kwargs(
+            conversation=conversation, tool_choice=tool_choice, n_choices=n_choices
+        )
+        if not self._apply_response_schema_via_provider:
+            completion_kwargs.pop("api_response_schema", None)
+
+        api_stream = self._get_completion_stream(**completion_kwargs)
+        api_stream = cast("AsyncIterator[Any]", api_stream)
+
+        async def iterator() -> AsyncIterator[CompletionChunkEvent | CompletionEvent]:
+            api_completion_chunks: list[Any] = []
+
+            async for api_completion_chunk in api_stream:
+                api_completion_chunks.append(api_completion_chunk)
+                completion_chunk = self._converters.from_completion_chunk(
+                    api_completion_chunk, name=self.model_id
+                )
+
+                yield CompletionChunkEvent(
+                    data=completion_chunk, proc_name=proc_name, call_id=call_id
+                )
+
+            api_completion = self.combine_completion_chunks(api_completion_chunks)
+            completion = self._converters.from_completion(
+                api_completion, name=self.model_id
+            )
+
+            yield CompletionEvent(data=completion, proc_name=proc_name, call_id=call_id)
+
+            if not self._apply_response_schema_via_provider:
+                self._validate_response(completion)
+                self._validate_tool_calls(completion)
+
+        return iterator()
+
+    async def generate_completion_stream(  # type: ignore[override]
+        self,
+        conversation: Messages,
+        *,
+        tool_choice: ToolChoice | None = None,
+        n_choices: int | None = None,
+        proc_name: str | None = None,
+        call_id: str | None = None,
+    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent | LLMStreamingErrorEvent]:
+        n_attempt = 0
+        while n_attempt <= self.max_response_retries:
+            try:
+                async for event in await self._generate_completion_stream_once(  # type: ignore[return]
+                    conversation,  # type: ignore[arg-type]
+                    tool_choice=tool_choice,
+                    n_choices=n_choices,
+                    proc_name=proc_name,
+                    call_id=call_id,
+                ):
+                    yield event
+                return
+            except (LLMResponseValidationError, LLMToolCallValidationError) as err:
+                err_data = LLMStreamingErrorData(
+                    error=err, model_name=self._model_name, model_id=self.model_id
+                )
+                yield LLMStreamingErrorEvent(
+                    data=err_data, proc_name=proc_name, call_id=call_id
+                )
+
+                n_attempt += 1
+                if n_attempt > self.max_response_retries:
+                    if n_attempt == 1:
+                        logger.warning(f"\nCloudLLM completion request failed:\n{err}")
+                    if n_attempt > 1:
+                        logger.warning(
+                            "\nCloudLLM completion request failed after "
+                            f"retrying:\n{err}"
+                        )
+                        refusal_completion = make_refusal_completion(
+                            self._model_name, err
+                        )
+                        yield CompletionEvent(
+                            data=refusal_completion,
+                            proc_name=proc_name,
+                            call_id=call_id,
+                        )
+                    raise err
+                    # return
+
+                logger.warning(
+                    "\nCloudLLM completion request failed "
+                    f"(retry attempt {n_attempt}):\n{err}"
+                )

grasp_agents-0.5.1/src/grasp_agents/comm_processor.py

@@ -0,0 +1,214 @@
+import logging
+from collections.abc import AsyncIterator, Sequence
+from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast
+
+from .errors import PacketRoutingError
+from .memory import MemT
+from .packet import Packet
+from .packet_pool import PacketPool
+from .processor import Processor
+from .run_context import CtxT, RunContext
+from .typing.events import Event, ProcPacketOutputEvent, RunResultEvent
+from .typing.io import InT, OutT_co, ProcName
+
+logger = logging.getLogger(__name__)
+
+
+_OutT_contra = TypeVar("_OutT_contra", contravariant=True)
+
+
+class ExitCommunicationHandler(Protocol[_OutT_contra, CtxT]):
+    def __call__(
+        self,
+        out_packet: Packet[_OutT_contra],
+        ctx: RunContext[CtxT],
+    ) -> bool: ...
+
+
+class SetRecipientsHandler(Protocol[_OutT_contra, CtxT]):
+    def __call__(
+        self, out_packet: Packet[_OutT_contra], ctx: RunContext[CtxT]
+    ) -> None: ...
+
+
+class CommProcessor(
+    Processor[InT, OutT_co, MemT, CtxT],
+    Generic[InT, OutT_co, MemT, CtxT],
+):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+        0: "_in_type",
+        1: "_out_type",
+    }
+
+    def __init__(
+        self,
+        name: ProcName,
+        *,
+        recipients: Sequence[ProcName] | None = None,
+        packet_pool: PacketPool[CtxT] | None = None,
+        max_retries: int = 0,
+    ) -> None:
+        super().__init__(name=name, max_retries=max_retries)
+
+        self.recipients = recipients or []
+        self._packet_pool = packet_pool
+        self._is_listening = False
+
+        self._exit_communication_impl: (
+            ExitCommunicationHandler[OutT_co, CtxT] | None
+        ) = None
+        self._set_recipients_impl: SetRecipientsHandler[OutT_co, CtxT] | None = None
+
+    @property
+    def packet_pool(self) -> PacketPool[CtxT] | None:
+        return self._packet_pool
+
+    @property
+    def is_listening(self) -> bool:
+        return self._is_listening
+
+    def _set_recipients(
+        self, out_packet: Packet[OutT_co], ctx: RunContext[CtxT]
+    ) -> None:
+        if self._set_recipients_impl:
+            self._set_recipients_impl(out_packet=out_packet, ctx=ctx)
+            return
+
+        out_packet.recipients = self.recipients
+
+    def _validate_routing(self, recipients: Sequence[ProcName]) -> Sequence[ProcName]:
+        for r in recipients:
+            if r not in self.recipients:
+                raise PacketRoutingError(
+                    selected_recipient=r,
+                    allowed_recipients=cast("list[str]", self.recipients),
+                )
+
+        return self.recipients
+
+    async def run(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | Sequence[InT] | None = None,
+        forgetful: bool = False,
+        call_id: str | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> Packet[OutT_co]:
+        call_id = self._generate_call_id(call_id)
+
+        out_packet = await super().run(
+            chat_inputs=chat_inputs,
+            in_packet=in_packet,
+            in_args=in_args,
+            forgetful=forgetful,
+            call_id=call_id,
+            ctx=ctx,
+        )
+
+        if self._packet_pool is not None:
+            if ctx is None:
+                raise ValueError("RunContext must be provided when using PacketPool")
+            if self._exit_communication(out_packet=out_packet, ctx=ctx):
+                ctx.result = out_packet
+                await self._packet_pool.stop_all()
+                return out_packet
+
+            self._set_recipients(out_packet=out_packet, ctx=ctx)
+            out_packet.recipients = self._validate_routing(out_packet.recipients)
+
+            await self._packet_pool.post(out_packet)
+
+        return out_packet
+
+    async def run_stream(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | Sequence[InT] | None = None,
+        forgetful: bool = False,
+        call_id: str | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> AsyncIterator[Event[Any]]:
+        call_id = self._generate_call_id(call_id)
+
+        out_packet: Packet[OutT_co] | None = None
+        async for event in super().run_stream(
+            chat_inputs=chat_inputs,
+            in_packet=in_packet,
+            in_args=in_args,
+            forgetful=forgetful,
+            call_id=call_id,
+            ctx=ctx,
+        ):
+            if isinstance(event, ProcPacketOutputEvent):
+                out_packet = event.data
+            else:
+                yield event
+
+        if out_packet is None:
+            return
+
+        if self._packet_pool is not None:
+            if ctx is None:
+                raise ValueError("RunContext must be provided when using PacketPool")
+            if self._exit_communication(out_packet=out_packet, ctx=ctx):
+                ctx.result = out_packet
+                yield RunResultEvent(
+                    data=out_packet, proc_name=self.name, call_id=call_id
+                )
+                await self._packet_pool.stop_all()
+                return
+
+            self._set_recipients(out_packet=out_packet, ctx=ctx)
+            out_packet.recipients = self._validate_routing(out_packet.recipients)
+
+            await self._packet_pool.post(out_packet)
+
+        yield ProcPacketOutputEvent(
+            data=out_packet, proc_name=self.name, call_id=call_id
+        )
+
+    def start_listening(self, ctx: RunContext[CtxT], **run_kwargs: Any) -> None:
+        if self._packet_pool is None:
+            raise RuntimeError("Packet pool must be initialized before listening")
+
+        if self._is_listening:
+            return
+        self._is_listening = True
+
+        self._packet_pool.register_packet_handler(
+            processor_name=self.name,
+            handler=self.run_stream if ctx.is_streaming else self.run,  # type: ignore[call-arg]
+            ctx=ctx,
+            **run_kwargs,
+        )
+
+    def _exit_communication(
+        self, out_packet: Packet[OutT_co], ctx: RunContext[CtxT]
+    ) -> bool:
+        if self._exit_communication_impl:
+            return self._exit_communication_impl(out_packet=out_packet, ctx=ctx)
+
+        return False
+
+    def exit_communication(
+        self, func: ExitCommunicationHandler[OutT_co, CtxT]
+    ) -> ExitCommunicationHandler[OutT_co, CtxT]:
+        self._exit_communication_impl = func
+
+        return func
+
+    def set_recipients(
+        self, func: SetRecipientsHandler[OutT_co, CtxT]
+    ) -> SetRecipientsHandler[OutT_co, CtxT]:
+        self._select_recipients_impl = func
+
+        return func
+
+    # async def stop_listening(self) -> None:
+    #     assert self._packet_pool is not None
+    #     self._is_listening = False
+    #     await self._packet_pool.unregister_packet_handler(self.name)