grasp_agents 0.5.8__py3-none-any.whl → 0.5.10__py3-none-any.whl

grasp_agents/llm_policy_executor.py

@@ -1,6 +1,6 @@
 import asyncio
 import json
-from collections.abc import AsyncIterator, Coroutine, Sequence
+from collections.abc import AsyncIterator, Coroutine, Mapping, Sequence
 from itertools import starmap
 from logging import getLogger
 from typing import Any, Generic, Protocol, final
@@ -36,7 +36,7 @@ class ToolCallLoopTerminator(Protocol[CtxT]):
         self,
         conversation: Messages,
         *,
-        ctx: RunContext[CtxT] | None,
+        ctx: RunContext[CtxT],
         **kwargs: Any,
     ) -> bool: ...
 
@@ -46,7 +46,7 @@ class MemoryManager(Protocol[CtxT]):
         self,
         memory: LLMAgentMemory,
         *,
-        ctx: RunContext[CtxT] | None,
+        ctx: RunContext[CtxT],
         **kwargs: Any,
     ) -> None: ...
 
@@ -54,9 +54,12 @@ class MemoryManager(Protocol[CtxT]):
 class LLMPolicyExecutor(Generic[CtxT]):
     def __init__(
         self,
+        *,
         agent_name: str,
         llm: LLM[LLMSettings, Converters],
         tools: list[BaseTool[BaseModel, Any, CtxT]] | None,
+        response_schema: Any | None = None,
+        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
         max_turns: int,
         react_mode: bool = False,
         final_answer_type: type[BaseModel] = BaseModel,
@@ -70,12 +73,15 @@ class LLMPolicyExecutor(Generic[CtxT]):
         self._final_answer_as_tool_call = final_answer_as_tool_call
         self._final_answer_tool = self.get_final_answer_tool()
 
-        _tools: list[BaseTool[BaseModel, Any, CtxT]] | None = tools
+        tools_list: list[BaseTool[BaseModel, Any, CtxT]] | None = tools
         if tools and final_answer_as_tool_call:
-            _tools = tools + [self._final_answer_tool]
+            tools_list = tools + [self._final_answer_tool]
+        self._tools = {t.name: t for t in tools_list} if tools_list else None
+
+        self._response_schema = response_schema
+        self._response_schema_by_xml_tag = response_schema_by_xml_tag
 
         self._llm = llm
-        self._llm.tools = _tools
 
         self._max_turns = max_turns
         self._react_mode = react_mode
@@ -91,9 +97,21 @@ class LLMPolicyExecutor(Generic[CtxT]):
     def llm(self) -> LLM[LLMSettings, Converters]:
         return self._llm
 
+    @property
+    def response_schema(self) -> Any | None:
+        return self._response_schema
+
+    @response_schema.setter
+    def response_schema(self, value: Any | None) -> None:
+        self._response_schema = value
+
+    @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, CtxT]]:
-        return self._llm.tools or {}
+        return self._tools or {}
 
     @property
     def max_turns(self) -> int:
@@ -104,7 +122,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         self,
         conversation: Messages,
         *,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
         **kwargs: Any,
     ) -> bool:
         if self.tool_call_loop_terminator:
@@ -117,7 +135,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         self,
         memory: LLMAgentMemory,
         *,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
         **kwargs: Any,
     ) -> None:
         if self.memory_manager:
@@ -126,12 +144,16 @@ class LLMPolicyExecutor(Generic[CtxT]):
     async def generate_message(
         self,
         memory: LLMAgentMemory,
+        *,
         call_id: str,
         tool_choice: ToolChoice | None = None,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> AssistantMessage:
         completion = await self.llm.generate_completion(
             memory.message_history,
+            response_schema=self.response_schema,
+            response_schema_by_xml_tag=self.response_schema_by_xml_tag,
+            tools=self.tools,
             tool_choice=tool_choice,
             n_choices=1,
             proc_name=self.agent_name,
@@ -147,9 +169,10 @@ class LLMPolicyExecutor(Generic[CtxT]):
     async def generate_message_stream(
         self,
         memory: LLMAgentMemory,
+        *,
         call_id: str,
         tool_choice: ToolChoice | None = None,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> AsyncIterator[
         CompletionChunkEvent[CompletionChunk]
         | CompletionEvent
@@ -160,6 +183,9 @@ class LLMPolicyExecutor(Generic[CtxT]):
 
         llm_event_stream = self.llm.generate_completion_stream(
             memory.message_history,
+            response_schema=self.response_schema,
+            response_schema_by_xml_tag=self.response_schema_by_xml_tag,
+            tools=self.tools,
             tool_choice=tool_choice,
             n_choices=1,
             proc_name=self.agent_name,
@@ -189,14 +215,14 @@ class LLMPolicyExecutor(Generic[CtxT]):
         calls: Sequence[ToolCall],
         memory: LLMAgentMemory,
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> Sequence[ToolMessage]:
         # TODO: Add image support
         corouts: list[Coroutine[Any, Any, BaseModel]] = []
         for call in calls:
             tool = self.tools[call.tool_name]
             args = json.loads(call.tool_arguments)
-            corouts.append(tool(ctx=ctx, **args))
+            corouts.append(tool(call_id=call_id, ctx=ctx, **args))
 
         outs = await asyncio.gather(*corouts)
         tool_messages = list(
@@ -217,7 +243,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         calls: Sequence[ToolCall],
         memory: LLMAgentMemory,
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> AsyncIterator[ToolMessageEvent]:
         tool_messages = await self.call_tools(
             calls, memory=memory, call_id=call_id, ctx=ctx
@@ -245,7 +271,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
                 return final_answer_message
 
     async def _generate_final_answer(
-        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT] | None = None
+        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT]
     ) -> AssistantMessage:
         user_message = UserMessage.from_text(
             "Exceeded the maximum number of turns: provide a final answer now!"
@@ -268,7 +294,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         return final_answer_message
 
     async def _generate_final_answer_stream(
-        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT] | None = None
+        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT]
     ) -> AsyncIterator[Event[Any]]:
         user_message = UserMessage.from_text(
             "Exceeded the maximum number of turns: provide a final answer now!",
@@ -296,7 +322,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         )
 
     async def execute(
-        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT] | None = None
+        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT]
     ) -> AssistantMessage | Sequence[AssistantMessage]:
         # 1. Generate the first message:
         #    In ReAct mode, we generate the first message without tool calls
@@ -379,7 +405,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         self,
         memory: LLMAgentMemory,
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> AsyncIterator[Event[Any]]:
         tool_choice: ToolChoice = "none" if self._react_mode else "auto"
         gen_message: AssistantMessage | None = None
@@ -464,7 +490,11 @@ class LLMPolicyExecutor(Generic[CtxT]):
             )
 
             async def run(
-                self, inp: BaseModel, ctx: RunContext[Any] | None = None
+                self,
+                inp: BaseModel,
+                *,
+                call_id: str | None = None,
+                ctx: RunContext[Any] | None = None,
             ) -> None:
                 return None
 
@@ -473,22 +503,22 @@ class LLMPolicyExecutor(Generic[CtxT]):
     def _process_completion(
         self,
         completion: Completion,
+        *,
         call_id: str,
         print_messages: bool = False,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> None:
-        if ctx is not None:
-            ctx.completions[self.agent_name].append(completion)
-            ctx.usage_tracker.update(
+        ctx.completions[self.agent_name].append(completion)
+        ctx.usage_tracker.update(
+            agent_name=self.agent_name,
+            completions=[completion],
+            model_name=self.llm.model_name,
+        )
+        if ctx.printer and print_messages:
+            usages = [None] * (len(completion.messages) - 1) + [completion.usage]
+            ctx.printer.print_messages(
+                completion.messages,
+                usages=usages,
                 agent_name=self.agent_name,
-                completions=[completion],
-                model_name=self.llm.model_name,
+                call_id=call_id,
             )
-            if ctx.printer and print_messages:
-                usages = [None] * (len(completion.messages) - 1) + [completion.usage]
-                ctx.printer.print_messages(
-                    completion.messages,
-                    usages=usages,
-                    agent_name=self.agent_name,
-                    call_id=call_id,
-                )

grasp_agents/openai/converters.py

@@ -96,8 +96,10 @@ class OpenAIConverters(Converters):
         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)
+    def to_tool(
+        tool: BaseTool[BaseModel, Any, Any], strict: bool | None = None, **kwargs: Any
+    ) -> OpenAIToolParam:
+        return to_api_tool(tool, strict=strict, **kwargs)
 
     @staticmethod
     def to_tool_choice(

grasp_agents/openai/openai_llm.py

@@ -3,9 +3,9 @@ import logging
 import os
 from collections.abc import AsyncIterator, Iterable, Mapping
 from copy import deepcopy
+from dataclasses import dataclass, field
 from typing import Any, Literal
 
-import httpx
 from openai import AsyncOpenAI, AsyncStream
 from openai._types import NOT_GIVEN  # type: ignore[import]
 from openai.lib.streaming.chat import (
@@ -15,8 +15,7 @@ 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, LLMRateLimiter
-from ..http_client import AsyncHTTPClientParams
+from ..cloud_llm import APIProvider, CloudLLM, CloudLLMSettings
 from ..typing.tool import BaseTool
 from . import (
     OpenAICompletion,
@@ -90,97 +89,75 @@ class OpenAILLMSettings(CloudLLMSettings, total=False):
     # TODO: support audio
 
 
+@dataclass(frozen=True)
 class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
-    def __init__(
-        self,
-        # Base LLM args
-        model_name: str,
-        llm_settings: OpenAILLMSettings | 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,
-        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: LLMRateLimiter | None = None,
-        # LLM response retries: try to regenerate to pass validation
-        max_response_retries: int = 1,
-    ) -> None:
+    converters: OpenAIConverters = field(default_factory=OpenAIConverters)
+    async_openai_client_params: dict[str, Any] | None = None
+    client: AsyncOpenAI = field(init=False)
+
+    def __post_init__(self):
+        super().__post_init__()
+
         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
+        _api_provider = self.api_provider
+
+        model_name_parts = self.model_name.split("/", 1)
+        if _api_provider is not None:
+            _model_name = self.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
+            provider_name, _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"API provider '{provider_name}' is not a supported OpenAI "
+                    f"compatible provider. Supported providers are: "
                     f"{', '.join(compat_providers_map.keys())}"
                 )
-            api_provider = compat_providers_map[provider_name]
+            _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=provider_model_name,
-            model_id=model_id,
-            llm_settings=llm_settings,
-            converters=OpenAIConverters(),
-            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,
-            async_http_client=async_http_client,
-            async_http_client_params=async_http_client_params,
-            rate_limiter=rate_limiter,
-            max_client_retries=max_client_retries,
-            max_response_retries=max_response_retries,
-        )
+        if self.llm_settings is not None:
+            stream_options = self.llm_settings.get("stream_options") or {}
+            stream_options["include_usage"] = True
+            _llm_settings = deepcopy(self.llm_settings)
+            _llm_settings["stream_options"] = stream_options
+        else:
+            _llm_settings = OpenAILLMSettings(stream_options={"include_usage": True})
 
         response_schema_support: bool = any(
-            fnmatch.fnmatch(self._model_name, pat)
-            for pat in api_provider.get("response_schema_support") or []
+            fnmatch.fnmatch(_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."
-                )
+        if self.apply_response_schema_via_provider and not response_schema_support:
+            raise ValueError(
+                "Native response schema validation is not supported for model "
+                f"'{_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
+        _async_openai_client_params = deepcopy(self.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.api_provider.get("base_url"),
-            api_key=self.api_provider.get("api_key"),
-            max_retries=max_client_retries,
+        _client = AsyncOpenAI(
+            base_url=_api_provider.get("base_url"),
+            api_key=_api_provider.get("api_key"),
+            max_retries=self.max_client_retries,
             **_async_openai_client_params,
         )
 
+        object.__setattr__(self, "model_name", _model_name)
+        object.__setattr__(self, "api_provider", _api_provider)
+        object.__setattr__(self, "llm_settings", _llm_settings)
+        object.__setattr__(self, "client", _client)
+
     async def _get_completion(
         self,
         api_messages: Iterable[OpenAIMessageParam],
@@ -195,9 +172,9 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         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,
+        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,
@@ -206,8 +183,8 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
                 **api_llm_settings,
             )
 
-        return await self._client.chat.completions.create(
-            model=self._model_name,
+        return await self.client.chat.completions.create(
+            model=self.model_name,
             messages=api_messages,
             tools=tools,
             tool_choice=tool_choice,
@@ -230,10 +207,10 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         response_format = api_response_schema or NOT_GIVEN
         n = n_choices or NOT_GIVEN
 
-        if self._apply_response_schema_via_provider:
+        if self.apply_response_schema_via_provider:
             stream_manager: OpenAIAsyncChatCompletionStreamManager[Any] = (
-                self._client.beta.chat.completions.stream(
-                    model=self._model_name,
+                self.client.beta.chat.completions.stream(
+                    model=self.model_name,
                     messages=api_messages,
                     tools=tools,
                     tool_choice=tool_choice,
@@ -249,8 +226,8 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         else:
             stream_generator: AsyncStream[
                 OpenAICompletionChunk
-            ] = await self._client.chat.completions.create(
-                model=self._model_name,
+            ] = await self.client.chat.completions.create(
+                model=self.model_name,
                 messages=api_messages,
                 tools=tools,
                 tool_choice=tool_choice,
@@ -263,16 +240,20 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
                     yield completion_chunk
 
     def combine_completion_chunks(
-        self, completion_chunks: list[OpenAICompletionChunk]
+        self,
+        completion_chunks: list[OpenAICompletionChunk],
+        response_schema: Any | None = None,
+        tools: Mapping[str, BaseTool[BaseModel, Any, Any]] | None = None,
     ) -> 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:
+        if self.apply_response_schema_via_provider:
+            if response_schema:
+                response_format = response_schema
+            if tools:
                 input_tools = [
-                    self._converters.to_tool(tool) for tool in self._tools.values()
+                    self.converters.to_tool(tool, strict=True)
+                    for tool in tools.values()
                 ]
         state = ChatCompletionStreamState[Any](
             input_tools=input_tools, response_format=response_format

grasp_agents/openai/tool_converters.py

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

grasp_agents/processors/base_processor.py

@@ -1,6 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Callable, Coroutine
+from collections.abc import AsyncIterator, Callable, Coroutine, Sequence
 from functools import wraps
 from typing import (
     Any,
@@ -37,7 +37,6 @@ from ..typing.tool import BaseTool
 
 logger = logging.getLogger(__name__)
 
-_OutT_contra = TypeVar("_OutT_contra", contravariant=True)
 
 F = TypeVar("F", bound=Callable[..., Coroutine[Any, Any, Packet[Any]]])
 F_stream = TypeVar("F_stream", bound=Callable[..., AsyncIterator[Event[Any]]])
@@ -102,10 +101,13 @@ def with_retry_stream(func: F_stream) -> F_stream:
     return cast("F_stream", wrapper)
 
 
+_OutT_contra = TypeVar("_OutT_contra", contravariant=True)
+
+
 class RecipientSelector(Protocol[_OutT_contra, CtxT]):
     def __call__(
-        self, output: _OutT_contra, ctx: RunContext[CtxT] | None
-    ) -> list[ProcName] | None: ...
+        self, output: _OutT_contra, ctx: RunContext[CtxT]
+    ) -> Sequence[ProcName] | None: ...
 
 
 class BaseProcessor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]):
@@ -118,7 +120,7 @@ class BaseProcessor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, C
         self,
         name: ProcName,
         max_retries: int = 0,
-        recipients: list[ProcName] | None = None,
+        recipients: Sequence[ProcName] | None = None,
         **kwargs: Any,
     ) -> None:
         self._in_type: type[InT]
@@ -239,7 +241,7 @@ class BaseProcessor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, C
             ) from err
 
     def _validate_recipients(
-        self, recipients: list[ProcName] | None, call_id: str
+        self, recipients: Sequence[ProcName] | None, call_id: str
     ) -> None:
         for r in recipients or []:
             if r not in (self.recipients or []):
@@ -252,8 +254,8 @@ class BaseProcessor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, C
 
     @final
     def _select_recipients(
-        self, output: OutT, ctx: RunContext[CtxT] | None = None
-    ) -> list[ProcName] | None:
+        self, output: OutT, ctx: RunContext[CtxT]
+    ) -> Sequence[ProcName] | None:
         if self.recipient_selector:
             return self.recipient_selector(output=output, ctx=ctx)
 
@@ -310,9 +312,15 @@ class BaseProcessor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, C
             name: str = tool_name
             description: str = tool_description
 
-            async def run(self, inp: InT, ctx: RunContext[CtxT] | None = None) -> OutT:
+            async def run(
+                self,
+                inp: InT,
+                *,
+                call_id: str | None = None,
+                ctx: RunContext[CtxT] | None = None,
+            ) -> OutT:
                 result = await processor_instance.run(
-                    in_args=inp, forgetful=True, ctx=ctx
+                    in_args=inp, forgetful=True, call_id=call_id, ctx=ctx
                 )
 
                 return result.payloads[0]

grasp_agents/processors/parallel_processor.py

@@ -30,7 +30,7 @@ class ParallelProcessor(
         in_args: InT | None = None,
         memory: MemT,
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> OutT:
         return cast("OutT", in_args)
 
@@ -41,7 +41,7 @@ class ParallelProcessor(
         in_args: InT | None = None,
         memory: MemT,
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> AsyncIterator[Event[Any]]:
         output = cast("OutT", in_args)
         yield ProcPayloadOutputEvent(data=output, proc_name=self.name, call_id=call_id)
@@ -67,7 +67,7 @@ class ParallelProcessor(
         in_args: InT | None = None,
         forgetful: bool = False,
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> Packet[OutT]:
         memory = self.memory.model_copy(deep=True) if forgetful else self.memory
 
@@ -86,7 +86,7 @@ class ParallelProcessor(
         return Packet(payloads=[val_output], sender=self.name, recipients=recipients)
 
     async def _run_parallel(
-        self, in_args: list[InT], call_id: str, ctx: RunContext[CtxT] | None = None
+        self, in_args: list[InT], call_id: str, ctx: RunContext[CtxT]
     ) -> Packet[OutT]:
         tasks = [
             self._run_single(
@@ -114,6 +114,7 @@ class ParallelProcessor(
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT]:
         call_id = self._generate_call_id(call_id)
+        ctx = RunContext[CtxT](state=None) if ctx is None else ctx  # type: ignore
 
         val_in_args = self._validate_inputs(
             call_id=call_id,
@@ -143,7 +144,7 @@ class ParallelProcessor(
         in_args: InT | None = None,
         forgetful: bool = False,
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> AsyncIterator[Event[Any]]:
         memory = self.memory.model_copy(deep=True) if forgetful else self.memory
 
@@ -178,7 +179,7 @@ class ParallelProcessor(
         self,
         in_args: list[InT],
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> AsyncIterator[Event[Any]]:
         streams = [
             self._run_single_stream(
@@ -222,6 +223,7 @@ class ParallelProcessor(
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         call_id = self._generate_call_id(call_id)
+        ctx = RunContext[CtxT](state=None) if ctx is None else ctx  # type: ignore
 
         val_in_args = self._validate_inputs(
             call_id=call_id,