grasp_agents 0.2.10__py3-none-any.whl → 0.2.11__py3-none-any.whl

grasp_agents/llm_agent.py

@@ -18,13 +18,7 @@ from .prompt_builder import (
     FormatSystemArgsHandler,
     PromptBuilder,
 )
-from .run_context import (
-    CtxT,
-    InteractionRecord,
-    RunContextWrapper,
-    SystemRunArgs,
-    UserRunArgs,
-)
+from .run_context import CtxT, InteractionRecord, RunContextWrapper
 from .tool_orchestrator import (
     ExitToolCallLoopHandler,
     ManageAgentStateHandler,
@@ -226,8 +220,8 @@ class LLMAgent(
         **gen_kwargs: Any,  # noqa: ARG002
     ) -> AgentMessage[OutT, LLMAgentState]:
         # Get run arguments
-        sys_args: SystemRunArgs = LLMPromptArgs()
-        usr_args: UserRunArgs = LLMPromptArgs()
+        sys_args: LLMPromptArgs = LLMPromptArgs()
+        usr_args: LLMPromptArgs | Sequence[LLMPromptArgs] = LLMPromptArgs()
         if ctx is not None:
             run_args = ctx.run_args.get(self.agent_id)
             if run_args is not None:
@@ -240,6 +234,7 @@ class LLMAgent(
             in_message=in_message,
             entry_point=entry_point,
         )
+        resolved_in_args = in_message.payloads if in_message else in_args
 
         # 1. Make system prompt (can be None)
         formatted_sys_prompt = self._prompt_builder.make_sys_prompt(
@@ -264,16 +259,11 @@ class LLMAgent(
         self._print_sys_msg(state=state, prev_mh_len=prev_mh_len, ctx=ctx)
 
         # 3. Make and add user messages (can be empty)
-        _in_args_batch: Sequence[InT] | None = None
-        if in_message is not None:
-            _in_args_batch = in_message.payloads
-        elif in_args is not None:
-            _in_args_batch = in_args if isinstance(in_args, Sequence) else [in_args]  # type: ignore[assignment]
 
         user_message_batch = self._prompt_builder.make_user_messages(
             chat_inputs=chat_inputs,
+            in_args=resolved_in_args,
             usr_args=usr_args,
-            in_args_batch=_in_args_batch,
             entry_point=entry_point,
             ctx=ctx,
         )
@@ -290,18 +280,23 @@ class LLMAgent(
             await self._tool_orchestrator.run_loop(state=state, ctx=ctx)
 
         # 5. Parse outputs
-        batch_size = state.message_history.batch_size
-        in_args_batch = in_message.payloads if in_message else batch_size * [None]
-        val_output_batch = [
-            self._out_type_adapter.validate_python(
-                self._parse_output(conversation=conv, in_args=in_args, ctx=ctx)
-            )
-            for conv, in_args in zip(
-                state.message_history.batched_conversations,
-                in_args_batch,
-                strict=False,
+
+        val_output_batch: list[OutT] = []
+        for i, _conv in enumerate(state.message_history.batched_conversations):
+            if isinstance(resolved_in_args, Sequence):
+                _resolved_in_args = cast("Sequence[InT]", resolved_in_args)
+                _in_args = _resolved_in_args[min(i, len(_resolved_in_args) - 1)]
+            else:
+                _resolved_in_args = cast("InT | None", resolved_in_args)
+                _in_args = _resolved_in_args
+
+            val_output_batch.append(
+                self._out_type_adapter.validate_python(
+                    self._parse_output(
+                        conversation=_conv, in_args=_in_args, batch_idx=i, ctx=ctx
+                    )
+                )
             )
-        ]
 
         # 6. Write interaction history to context
 
@@ -316,7 +311,7 @@ class LLMAgent(
                 in_prompt=self.in_prompt,
                 sys_args=sys_args,
                 usr_args=usr_args,
-                in_args=(in_message.payloads if in_message is not None else None),
+                in_args=resolved_in_args,  # type: ignore[valid-type]
                 outputs=val_output_batch,
                 state=state,
             )

grasp_agents/prompt_builder.py

@@ -1,11 +1,11 @@
 from collections.abc import Sequence
 from copy import deepcopy
-from typing import ClassVar, Generic, Protocol
+from typing import ClassVar, Generic, Protocol, cast
 
 from pydantic import BaseModel, TypeAdapter
 
 from .generics_utils import AutoInstanceAttributesMixin
-from .run_context import CtxT, RunContextWrapper, UserRunArgs
+from .run_context import CtxT, RunContextWrapper
 from .typing.content import ImageData
 from .typing.io import (
     InT,
@@ -142,11 +142,13 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
     def _usr_messages_from_prompt_template(
         self,
         in_prompt: LLMPrompt,
-        usr_args: UserRunArgs | None = None,
         in_args_batch: Sequence[InT] | None = None,
+        usr_args_batch: Sequence[LLMPromptArgs] | None = None,
         ctx: RunContextWrapper[CtxT] | None = None,
     ) -> Sequence[UserMessage]:
-        usr_args_batch_, in_args_batch_ = self._make_batched(usr_args, in_args_batch)
+        usr_args_batch_, in_args_batch_ = self._align_in_usr_batches(
+            in_args_batch, usr_args_batch
+        )
 
         val_usr_args_batch_ = [
             self.usr_args_schema.model_validate(u) for u in usr_args_batch_
@@ -174,8 +176,8 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
     def make_user_messages(
         self,
         chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
-        usr_args: UserRunArgs | None = None,
-        in_args_batch: Sequence[InT] | None = None,
+        in_args: InT | Sequence[InT] | None = None,
+        usr_args: LLMPromptArgs | Sequence[LLMPromptArgs] | None = None,
         entry_point: bool = False,
         ctx: RunContextWrapper[CtxT] | None = None,
     ) -> Sequence[UserMessage]:
@@ -196,10 +198,24 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
             if isinstance(chat_inputs, Sequence) and chat_inputs:
                 return self._usr_messages_from_content_parts(chat_inputs)
 
+        in_args_batch = cast(
+            "Sequence[InT] | None",
+            in_args if (isinstance(in_args, Sequence) or not in_args) else [in_args],
+        )
+
         # 2) No input prompt template + received args → raw JSON messages
         if self.in_prompt is None and in_args_batch:
             return self._usr_messages_from_in_args(in_args_batch)
 
+        usr_args_batch = cast(
+            "Sequence[LLMPromptArgs] | None",
+            (
+                usr_args
+                if (isinstance(usr_args, Sequence) or not usr_args)
+                else [usr_args]
+            ),
+        )
+
         # 3) Input prompt template + any args → batch & format
         if self.in_prompt is not None:
             if in_args_batch and not isinstance(in_args_batch[0], BaseModel):
@@ -209,21 +225,19 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
                 )
             return self._usr_messages_from_prompt_template(
                 in_prompt=self.in_prompt,
-                usr_args=usr_args,
+                usr_args_batch=usr_args_batch,
                 in_args_batch=in_args_batch,
                 ctx=ctx,
             )
 
         return []
 
-    def _make_batched(
+    def _align_in_usr_batches(
         self,
-        usr_args: UserRunArgs | None = None,
         in_args_batch: Sequence[InT] | None = None,
+        usr_args_batch: Sequence[LLMPromptArgs] | None = None,
     ) -> tuple[Sequence[LLMPromptArgs | DummySchema], Sequence[InT | DummySchema]]:
-        usr_args_batch_ = (
-            usr_args if isinstance(usr_args, list) else [usr_args or DummySchema()]
-        )
+        usr_args_batch_ = usr_args_batch or [DummySchema()]
         in_args_batch_ = in_args_batch or [DummySchema()]
 
         # Broadcast singleton → match lengths

grasp_agents/run_context.py

@@ -17,13 +17,10 @@ from .typing.io import (
 )
 from .usage_tracker import UsageTracker
 
-SystemRunArgs: TypeAlias = LLMPromptArgs
-UserRunArgs: TypeAlias = LLMPromptArgs | list[LLMPromptArgs]
-
 
 class RunArgs(BaseModel):
-    sys: SystemRunArgs = Field(default_factory=LLMPromptArgs)
-    usr: UserRunArgs = Field(default_factory=LLMPromptArgs)
+    sys: LLMPromptArgs = Field(default_factory=LLMPromptArgs)
+    usr: LLMPromptArgs | Sequence[LLMPromptArgs] = Field(default_factory=LLMPromptArgs)
 
     model_config = ConfigDict(extra="forbid")
 
@@ -35,9 +32,9 @@ class InteractionRecord(BaseModel, Generic[InT, OutT, StateT]):
     chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None
     sys_prompt: LLMPrompt | None = None
     in_prompt: LLMPrompt | None = None
-    sys_args: SystemRunArgs | None = None
-    usr_args: UserRunArgs | None = None
-    in_args: Sequence[InT] | None = None
+    sys_args: LLMPromptArgs | None = None
+    usr_args: LLMPromptArgs | Sequence[LLMPromptArgs] | None = None
+    in_args: InT | Sequence[InT] | None = None
     outputs: Sequence[OutT]
 
     model_config = ConfigDict(extra="forbid", frozen=True)

grasp_agents/workflow/looped_agent.py

@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+from itertools import pairwise
 from logging import getLogger
 from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
 
@@ -49,6 +50,19 @@ class LoopedWorkflowAgent(WorkflowAgent[InT, OutT, CtxT], Generic[InT, OutT, Ctx
             dynamic_routing=dynamic_routing,
         )
 
+        for prev_agent, agent in pairwise(subagents):
+            if prev_agent.out_type != agent.in_type:
+                raise ValueError(
+                    f"Output type {prev_agent.out_type} of agent {prev_agent.agent_id} "
+                    f"does not match input type {agent.in_type} of agent "
+                    f"{agent.agent_id}"
+                )
+        if subagents[-1].out_type != subagents[0].in_type:
+            raise ValueError(
+                f"Looped workflow's last agent output type {subagents[-1].out_type} "
+                f"does not match first agent input type {subagents[0].in_type}"
+            )
+
         self._max_iterations = max_iterations
 
         self._exit_workflow_loop_impl: ExitWorkflowLoopHandler[OutT, CtxT] | None = None

grasp_agents/workflow/sequential_agent.py

@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+from itertools import pairwise
 from typing import Any, ClassVar, Generic, cast, final
 
 from ..agent_message_pool import AgentMessage, AgentMessagePool
@@ -33,6 +34,14 @@ class SequentialWorkflowAgent(WorkflowAgent[InT, OutT, CtxT], Generic[InT, OutT,
             dynamic_routing=dynamic_routing,
         )
 
+        for prev_agent, agent in pairwise(subagents):
+            if prev_agent.out_type != agent.in_type:
+                raise ValueError(
+                    f"Output type {prev_agent.out_type} of agent {prev_agent.agent_id} "
+                    f"does not match input type {agent.in_type} of agent "
+                    f"{agent.agent_id}"
+                )
+
     @final
     async def run(
         self,

grasp_agents/workflow/workflow_agent.py

@@ -27,29 +27,44 @@ class WorkflowAgent(
         dynamic_routing: bool = False,
         **kwargs: Any,  # noqa: ARG002
     ) -> None:
-        if not subagents:
-            raise ValueError("At least one step is required")
+        super().__init__(
+            agent_id=agent_id,
+            message_pool=message_pool,
+            recipient_ids=recipient_ids,
+            dynamic_routing=dynamic_routing,
+        )
+
+        if len(subagents) < 2:
+            raise ValueError("At least two steps are required")
         if start_agent not in subagents:
             raise ValueError("Start agent must be in the subagents list")
         if end_agent not in subagents:
             raise ValueError("End agent must be in the subagents list")
 
-        self.subagents = subagents
+        if start_agent.in_type != self.in_type:
+            raise ValueError(
+                f"Start agent's input type {start_agent.in_type} does not "
+                f"match workflow's input type {self._in_type}"
+            )
+        if end_agent.out_type != self.out_type:
+            raise ValueError(
+                f"End agent's output type {end_agent.out_type} does not "
+                f"match workflow's output type {self._out_type}"
+            )
 
+        self._subagents = subagents
         self._start_agent = start_agent
         self._end_agent = end_agent
 
-        super().__init__(
-            agent_id=agent_id,
-            message_pool=message_pool,
-            recipient_ids=recipient_ids,
-            dynamic_routing=dynamic_routing,
-        )
         for subagent in subagents:
             assert not subagent.recipient_ids, (
                 "Subagents must not have recipient_ids set."
             )
 
+    @property
+    def subagents(self) -> Sequence[CommunicatingAgent[Any, Any, Any, CtxT]]:
+        return self._subagents
+
     @property
     def start_agent(self) -> CommunicatingAgent[InT, Any, Any, CtxT]:
         return self._start_agent

{grasp_agents-0.2.10.dist-info → grasp_agents-0.2.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: grasp_agents
-Version: 0.2.10
+Version: 0.2.11
 Summary: Grasp Agents Library
 License-File: LICENSE.md
 Requires-Python: <4,>=3.11.4

{grasp_agents-0.2.10.dist-info → grasp_agents-0.2.11.dist-info}/RECORD

@@ -9,12 +9,12 @@ grasp_agents/generics_utils.py,sha256=kw4Odte6Nvl4c9U7-mKPgXCavWZXo009zYDHAA0BR3
 grasp_agents/grasp_logging.py,sha256=H1GYhXdQvVkmauFDZ-KDwvVmPQHZUUm9sRqX_ObK2xI,1111
 grasp_agents/http_client.py,sha256=KZva2MjJjuI5ohUeU8RdTAImUnQYaqBrV2jDH8smbJw,738
 grasp_agents/llm.py,sha256=YClNxN9GUGaFHXhTU72z1AqW_Y726OC7kyVRYCnfhZ8,3682
-grasp_agents/llm_agent.py,sha256=trKvxIS5z9GYusB9zNRv_75trGP1NwzdvK2_-rElYnE,16102
+grasp_agents/llm_agent.py,sha256=UvmcvuLl8EWe51isSw9Oa9lxA9gQBJngRXj8LUMAd4w,16097
 grasp_agents/llm_agent_state.py,sha256=L54zUW5nAT-ubvEB7XNAQ84ExOgRlUFzc-Q49mUXUT0,2390
 grasp_agents/memory.py,sha256=X1YtVX8XxP5KnGPMW8BqjID8QK4hTG2obxoyhnnZ4pU,5575
 grasp_agents/printer.py,sha256=ZENcTITCcgSizvcUXnIQiFl_NcVHg-801Z-sbT0D8rg,5030
-grasp_agents/prompt_builder.py,sha256=uN21GAYdmbNiIkv60Rs_ixJyybR1JWR0lsOM9HVkiLE,8411
-grasp_agents/run_context.py,sha256=4v6IcddHSWUAMYX8M9hQRXNMfPf7gUv45nn_Fb5lLaI,2233
+grasp_agents/prompt_builder.py,sha256=A4z8B5nIfgn_cHnZCNyYjwA904UDiOaCyBc5VPu8168,8867
+grasp_agents/run_context.py,sha256=OpTdo32-WPD7XwE16LQphh2C7yrrkZ3C4Ia80pWeBDQ,2192
 grasp_agents/tool_orchestrator.py,sha256=XgKCOytPs4y4SBpS5i4eogCB428XaDnNIB3VxzJC_k0,6095
 grasp_agents/usage_tracker.py,sha256=4gy0XtfIBAjQHblEFpQuPelmxMGDvE7vw4c8Ccr1msk,3471
 grasp_agents/utils.py,sha256=KzoInW0sq-pwwUtgjtYMW8b0ivBH6MR0Zxv6Kqq3k3M,4510
@@ -37,10 +37,10 @@ grasp_agents/typing/io.py,sha256=uxSvbD05UK5nIhPfDvXIoGuU6xRMW4USZq_4IgBeGCY,609
 grasp_agents/typing/message.py,sha256=DC24XMcUPG1MHDNZKT67yKWyaMTQNF5B0yDiUg8b54Q,3833
 grasp_agents/typing/tool.py,sha256=e0pTMnRcpMpGNVQ8muE9wnh7LdIgh92AqXDo9hMDxf0,1960
 grasp_agents/workflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-grasp_agents/workflow/looped_agent.py,sha256=5tvlkPyhPdNGVSEaxTNPX0QQfPwP01OpZrvmv4bWXsA,4046
-grasp_agents/workflow/sequential_agent.py,sha256=SbeAKHCbhMy-2CmQWs1f5VMSRj88MLiAQc76g_ivGvA,2131
-grasp_agents/workflow/workflow_agent.py,sha256=E6jzfFmdUfCoYehH9ZsAa_MTICjEaP4NGwNxhc3NR_E,2436
-grasp_agents-0.2.10.dist-info/METADATA,sha256=lBHs7V6zm6E63yspRA9OOHS80M7hCtyVxuCNkaNce5Y,7188
-grasp_agents-0.2.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-grasp_agents-0.2.10.dist-info/licenses/LICENSE.md,sha256=-nNNdWqGB8gJ2O-peFQ2Irshv5tW5pHKyTcYkwvH7CE,1201
-grasp_agents-0.2.10.dist-info/RECORD,,
+grasp_agents/workflow/looped_agent.py,sha256=f7R-yVq4cP8zMTlvAF0GCpD0Izkyz5jUvUrxaXycusw,4710
+grasp_agents/workflow/sequential_agent.py,sha256=yU3X28rqZo-xVdkbEETb9pBCBG7v1GK7oxtbl5xLdCo,2526
+grasp_agents/workflow/workflow_agent.py,sha256=T1DpSipIkcrC_WJKR6Ho_cSHmVltCcF6Ve1F7isPRp0,3031
+grasp_agents-0.2.11.dist-info/METADATA,sha256=xXvakAAM5nycFuJf70buNFqQ5BMNyMbI7rdFZHHwphc,7188
+grasp_agents-0.2.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+grasp_agents-0.2.11.dist-info/licenses/LICENSE.md,sha256=-nNNdWqGB8gJ2O-peFQ2Irshv5tW5pHKyTcYkwvH7CE,1201
+grasp_agents-0.2.11.dist-info/RECORD,,