grasp_agents 0.3.11__py3-none-any.whl → 0.4.2__py3-none-any.whl

grasp_agents/cloud_llm.py

@@ -2,9 +2,9 @@ import fnmatch
 import logging
 import os
 from abc import abstractmethod
-from collections.abc import AsyncIterator, Mapping, Sequence
+from collections.abc import AsyncIterator, Mapping
 from copy import deepcopy
-from typing import Any, Generic, Literal
+from typing import Any, Generic, Literal, NotRequired
 
 import httpx
 from pydantic import BaseModel
@@ -16,10 +16,9 @@ from tenacity import (
 )
 from typing_extensions import TypedDict
 
-from .http_client import AsyncHTTPClientParams, create_async_http_client
+from .http_client import AsyncHTTPClientParams, create_simple_async_httpx_client
 from .llm import LLM, ConvertT_co, LLMSettings, SettingsT_co
-from .message_history import MessageHistory
-from .rate_limiting.rate_limiter_chunked import RateLimiterC, limit_rate_chunked
+from .rate_limiting.rate_limiter_chunked import RateLimiterC, limit_rate
 from .typing.completion import Completion
 from .typing.completion_chunk import (
     CompletionChoice,
@@ -33,30 +32,30 @@ from .typing.tool import BaseTool, ToolChoice
 logger = logging.getLogger(__name__)
 
 
-APIProvider = Literal["openai", "openrouter", "google_ai_studio"]
+APIProviderName = Literal["openai", "openrouter", "google_ai_studio"]
 
 
-class APIProviderInfo(TypedDict):
-    name: APIProvider
+class APIProvider(TypedDict):
+    name: APIProviderName
     base_url: str
-    api_key: str | None
-    struct_outputs_support: tuple[str, ...]
+    api_key: NotRequired[str | None]
+    struct_outputs_support: NotRequired[tuple[str, ...]]
 
 
-PROVIDERS: dict[APIProvider, APIProviderInfo] = {
-    "openai": APIProviderInfo(
+API_PROVIDERS: dict[APIProviderName, APIProvider] = {
+    "openai": APIProvider(
         name="openai",
         base_url="https://api.openai.com/v1",
         api_key=os.getenv("OPENAI_API_KEY"),
         struct_outputs_support=("*",),
     ),
-    "openrouter": APIProviderInfo(
+    "openrouter": APIProvider(
         name="openrouter",
         base_url="https://openrouter.ai/api/v1",
         api_key=os.getenv("OPENROUTER_API_KEY"),
         struct_outputs_support=(),
     ),
-    "google_ai_studio": APIProviderInfo(
+    "google_ai_studio": APIProvider(
         name="google_ai_studio",
         base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
         api_key=os.getenv("GOOGLE_AI_STUDIO_API_KEY"),
@@ -66,18 +65,17 @@ PROVIDERS: dict[APIProvider, APIProviderInfo] = {
 
 
 def retry_error_callback(retry_state: RetryCallState) -> Completion:
-    assert retry_state.outcome is not None
-    exception = retry_state.outcome.exception()
+    exception = retry_state.outcome.exception() if retry_state.outcome else None
     if exception:
         if retry_state.attempt_number == 1:
             logger.warning(
-                f"CloudLLM completion request failed:\n{exception}",
-                exc_info=exception,
+                f"\nCloudLLM completion request failed:\n{exception}",
+                # exc_info=exception,
             )
         if retry_state.attempt_number > 1:
             logger.warning(
-                f"CloudLLM completion request failed after retrying:\n{exception}",
-                exc_info=exception,
+                f"\nCloudLLM completion request failed after retrying:\n{exception}",
+                # exc_info=exception,
             )
     failed_message = AssistantMessage(content=None, refusal=str(exception))
 
@@ -87,11 +85,12 @@ def retry_error_callback(retry_state: RetryCallState) -> Completion:
     )
 
 
-def retry_before_callback(retry_state: RetryCallState) -> None:
-    if retry_state.attempt_number > 1:
+def retry_before_sleep_callback(retry_state: RetryCallState) -> None:
+    exception = retry_state.outcome.exception() if retry_state.outcome else None
+    if exception:
         logger.info(
-            "Retrying CloudLLM completion request "
-            f"(attempt {retry_state.attempt_number - 1}) ..."
+            "\nRetrying CloudLLM completion request "
+            f"(attempt {retry_state.attempt_number}):\n{exception}"
         )
 
 
@@ -106,10 +105,13 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         model_name: str,
         converters: ConvertT_co,
         llm_settings: SettingsT_co | None = None,
-        model_id: str | None = None,
         tools: list[BaseTool[BaseModel, Any, Any]] | None = None,
         response_format: type | Mapping[str, type] | None = None,
+        model_id: str | None = None,
+        # Custom LLM provider
+        api_provider: APIProvider | None = None,
         # Connection settings
+        async_http_client: httpx.AsyncClient | None = None,
         async_http_client_params: (
             dict[str, Any] | AsyncHTTPClientParams | None
         ) = None,
@@ -120,8 +122,6 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         rate_limiter_max_concurrency: int = 300,
         # Retries
         num_generation_retries: int = 0,
-        # Disable tqdm for batch processing
-        no_tqdm: bool = True,
         **kwargs: Any,
     ) -> None:
         self.llm_settings: CloudLLMSettings | None
@@ -139,29 +139,31 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         self._model_name = model_name
         model_name_parts = model_name.split(":", 1)
 
-        if len(model_name_parts) == 2 and model_name_parts[0] in PROVIDERS:
-            api_provider, api_model_name = model_name_parts
-            if api_provider not in PROVIDERS:
+        if len(model_name_parts) == 2:
+            api_provider_name, api_model_name = model_name_parts
+            self._api_model_name: str = api_model_name
+            if api_provider_name not in API_PROVIDERS:
                 raise ValueError(
-                    f"API provider '{api_provider}' is not supported. "
-                    f"Supported providers are: {', '.join(PROVIDERS.keys())}"
+                    f"API provider '{api_provider_name}' is not supported. "
+                    f"Supported providers are: {', '.join(API_PROVIDERS.keys())}"
                 )
-
-            self._api_provider: APIProvider | None = api_provider
-            self._api_model_name: str = api_model_name
-            self._base_url: str | None = PROVIDERS[api_provider]["base_url"]
-            self._api_key: str | None = PROVIDERS[api_provider]["api_key"]
-            self._struct_outputs_support: bool = any(
-                fnmatch.fnmatch(self._model_name, pat)
-                for pat in PROVIDERS[api_provider]["struct_outputs_support"]
+            _api_provider = API_PROVIDERS[api_provider_name]
+        elif api_provider is not None:
+            self._api_model_name: str = model_name
+            _api_provider = api_provider
+        else:
+            raise ValueError(
+                "API provider must be specified either in the model name "
+                "or as a separate argument."
             )
 
-        else:
-            self._api_provider = None
-            self._api_model_name = model_name
-            self._base_url = None
-            self._api_key = None
-            self._struct_outputs_support = False
+        self._api_provider_name: APIProviderName = _api_provider["name"]
+        self._base_url: str | None = _api_provider.get("base_url")
+        self._api_key: str | None = _api_provider.get("api_key")
+        self._struct_outputs_support: bool = any(
+            fnmatch.fnmatch(self._model_name, pat)
+            for pat in _api_provider.get("struct_outputs_support", ())
+        )
 
         if (
             self._llm_settings.get("use_struct_outputs")
@@ -181,23 +183,20 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
                 max_concurrency=rate_limiter_max_concurrency,
             )
         )
-        self.no_tqdm = no_tqdm
-        self._client: Any
 
         self._async_http_client: httpx.AsyncClient | None = None
-        if async_http_client_params is not None:
-            val_async_http_client_params = AsyncHTTPClientParams.model_validate(
+        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._async_http_client = create_async_http_client(
-                val_async_http_client_params
-            )
 
         self.num_generation_retries = num_generation_retries
 
     @property
-    def api_provider(self) -> APIProvider | None:
-        return self._api_provider
+    def api_provider_name(self) -> APIProviderName | None:
+        return self._api_provider_name
 
     @property
     def rate_limiter(
@@ -353,7 +352,8 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
 
         return iterate()
 
-    async def generate_completion(
+    @limit_rate
+    async def generate_completion(  # type: ignore[override]
         self,
         conversation: Messages,
         *,
@@ -363,7 +363,7 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         wrapped_func = retry(
             wait=wait_random_exponential(min=1, max=8),
             stop=stop_after_attempt(self.num_generation_retries + 1),
-            before=retry_before_callback,
+            before_sleep=retry_before_sleep_callback,
             retry_error_callback=retry_error_callback,
         )(self.__class__.generate_completion_no_retry)
 
@@ -371,23 +371,6 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
             self, conversation, tool_choice=tool_choice, n_choices=n_choices
         )
 
-    @limit_rate_chunked  # type: ignore
-    async def _generate_completion_batch(
-        self,
-        conversation: Messages,
-        *,
-        tool_choice: ToolChoice | None = None,
-    ) -> Completion:
-        return await self.generate_completion(conversation, tool_choice=tool_choice)
-
-    async def generate_completion_batch(
-        self, message_history: MessageHistory, *, tool_choice: ToolChoice | None = None
-    ) -> Sequence[Completion]:
-        return await self._generate_completion_batch(
-            list(message_history.conversations),  # type: ignore
-            tool_choice=tool_choice,
-        )
-
     def _get_rate_limiter(
         self,
         rate_limiter: RateLimiterC[Messages, AssistantMessage] | None = None,

grasp_agents/comm_processor.py

@@ -5,12 +5,13 @@ from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast
 from pydantic import BaseModel
 from pydantic.json_schema import SkipJsonSchema
 
+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, PacketEvent
-from .typing.io import InT_contra, MemT_co, OutT_co, ProcName
+from .typing.io import InT, OutT_co, ProcName
 
 logger = logging.getLogger(__name__)
 
@@ -31,8 +32,8 @@ class ExitCommunicationHandler(Protocol[_OutT_contra, CtxT]):
 
 
 class CommProcessor(
-    Processor[InT_contra, OutT_co, MemT_co, CtxT],
-    Generic[InT_contra, OutT_co, MemT_co, CtxT],
+    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",
@@ -45,8 +46,9 @@ class CommProcessor(
         *,
         recipients: Sequence[ProcName] | None = None,
         packet_pool: PacketPool[CtxT] | None = None,
+        num_par_run_retries: int = 0,
     ) -> None:
-        super().__init__(name=name)
+        super().__init__(name=name, num_par_run_retries=num_par_run_retries)
 
         self.recipients = recipients or []
 
@@ -56,6 +58,10 @@ class CommProcessor(
             ExitCommunicationHandler[OutT_co, CtxT] | None
         ) = None
 
+    @property
+    def packet_pool(self) -> PacketPool[CtxT] | None:
+        return self._packet_pool
+
     def _validate_routing(self, payloads: Sequence[OutT_co]) -> Sequence[ProcName]:
         if all(isinstance(p, DynCommPayload) for p in payloads):
             payloads_ = cast("Sequence[DynCommPayload]", payloads)
@@ -88,9 +94,10 @@ class CommProcessor(
         self,
         chat_inputs: Any | None = None,
         *,
-        in_packet: Packet[InT_contra] | None = None,
-        in_args: InT_contra | Sequence[InT_contra] | None = None,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | Sequence[InT] | None = None,
         forgetful: bool = False,
+        run_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT_co]:
         out_packet = await super().run(
@@ -98,6 +105,7 @@ class CommProcessor(
             in_packet=in_packet,
             in_args=in_args,
             forgetful=forgetful,
+            run_id=run_id,
             ctx=ctx,
         )
         recipients = self._validate_routing(out_packet.payloads)
@@ -114,9 +122,10 @@ class CommProcessor(
         self,
         chat_inputs: Any | None = None,
         *,
-        in_packet: Packet[InT_contra] | None = None,
-        in_args: InT_contra | Sequence[InT_contra] | None = None,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | None = None,
         forgetful: bool = False,
+        run_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         out_packet: Packet[OutT_co] | None = None
@@ -125,6 +134,7 @@ class CommProcessor(
             in_packet=in_packet,
             in_args=in_args,
             forgetful=forgetful,
+            run_id=run_id,
             ctx=ctx,
         ):
             if isinstance(event, PacketEvent):
@@ -152,7 +162,7 @@ class CommProcessor(
 
         return func
 
-    def _exit_communication_fn(
+    def _exit_communication(
         self, out_packet: Packet[OutT_co], ctx: RunContext[CtxT] | None
     ) -> bool:
         if self._exit_communication_impl:
@@ -162,7 +172,7 @@ class CommProcessor(
 
     async def _packet_handler(
         self,
-        packet: Packet[InT_contra],
+        packet: Packet[InT],
         ctx: RunContext[CtxT] | None = None,
         **run_kwargs: Any,
     ) -> None:
@@ -170,7 +180,7 @@ class CommProcessor(
 
         out_packet = await self.run(ctx=ctx, in_packet=packet, **run_kwargs)
 
-        if self._exit_communication_fn(out_packet=out_packet, ctx=ctx):
+        if self._exit_communication(out_packet=out_packet, ctx=ctx):
             await self._packet_pool.stop_all()
             return
 

grasp_agents/errors.py

@@ -0,0 +1,34 @@
+class InputValidationError(Exception):
+    pass
+
+
+class StringParsingError(Exception):
+    pass
+
+
+class CompletionError(Exception):
+    pass
+
+
+class CombineCompletionChunksError(Exception):
+    pass
+
+
+class ToolValidationError(Exception):
+    pass
+
+
+class OutputValidationError(Exception):
+    pass
+
+
+class WorkflowConstructionError(Exception):
+    pass
+
+
+class SystemPromptBuilderError(Exception):
+    pass
+
+
+class InputPromptBuilderError(Exception):
+    pass

grasp_agents/http_client.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import httpx
 from pydantic import BaseModel, NonNegativeFloat, PositiveInt
 
@@ -9,10 +11,12 @@ class AsyncHTTPClientParams(BaseModel):
     keepalive_expiry: float | None = 5
 
 
-def create_async_http_client(
-    client_params: AsyncHTTPClientParams,
+def create_simple_async_httpx_client(
+    client_params: AsyncHTTPClientParams | dict[str, Any],
 ) -> httpx.AsyncClient:
-    http_client = httpx.AsyncClient(
+    if isinstance(client_params, dict):
+        client_params = AsyncHTTPClientParams(**client_params)
+    return httpx.AsyncClient(
         timeout=httpx.Timeout(client_params.timeout),
         limits=httpx.Limits(
             max_connections=client_params.max_connections,
@@ -20,5 +24,3 @@ def create_async_http_client(
             keepalive_expiry=client_params.keepalive_expiry,
         ),
     )
-
-    return http_client

grasp_agents/llm.py

@@ -1,6 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Mapping, Sequence
+from collections.abc import AsyncIterator, Mapping
 from typing import Any, Generic, TypeVar, cast
 from uuid import uuid4
 
@@ -9,7 +9,7 @@ from typing_extensions import TypedDict
 
 from grasp_agents.utils import validate_obj_from_json_or_py_string
 
-from .message_history import MessageHistory
+from .errors import ToolValidationError
 from .typing.completion import Completion
 from .typing.converters import Converters
 from .typing.events import CompletionChunkEvent, CompletionEvent
@@ -118,7 +118,7 @@ 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 ValueError(
+                        raise ToolValidationError(
                             f"Tool '{tool_name}' is not available in the LLM tools "
                             f"(available: {available_tool_names}"
                         )
@@ -146,9 +146,3 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         n_choices: int | None = None,
     ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent]:
         pass
-
-    @abstractmethod
-    async def generate_completion_batch(
-        self, message_history: MessageHistory, *, tool_choice: ToolChoice | None = None
-    ) -> Sequence[Completion]:
-        pass