grasp_agents 0.5.9__py3-none-any.whl → 0.5.11__py3-none-any.whl

grasp_agents/llm.py

@@ -1,7 +1,8 @@
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Mapping, Sequence
-from typing import Any, Generic, TypeVar, cast
+from collections.abc import AsyncIterator, Mapping
+from dataclasses import dataclass, field
+from typing import Any, Generic, TypeVar
 from uuid import uuid4
 
 from pydantic import BaseModel
@@ -66,71 +67,24 @@ SettingsT_co = TypeVar("SettingsT_co", bound=LLMSettings, covariant=True)
 ConvertT_co = TypeVar("ConvertT_co", bound=Converters, covariant=True)
 
 
+@dataclass(frozen=True)
 class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
-    @abstractmethod
-    def __init__(
+    model_name: str
+    converters: ConvertT_co
+    llm_settings: SettingsT_co | None = None
+    model_id: str = field(default_factory=lambda: str(uuid4())[:8])
+
+    def _validate_response(
         self,
-        converters: ConvertT_co,
-        model_name: str | None = None,
-        model_id: str | None = None,
-        llm_settings: SettingsT_co | None = None,
-        tools: Sequence[BaseTool[BaseModel, Any, Any]] | None = None,
-        response_schema: Any | None = None,
-        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
-        **kwargs: Any,
+        completion: Completion,
+        response_schema: Any | None,
+        response_schema_by_xml_tag: Mapping[str, Any] | None,
     ) -> None:
-        super().__init__()
-
-        self._converters = converters
-        self._model_id = model_id or str(uuid4())[:8]
-        self._model_name = model_name
-        self._tools = {t.name: t for t in tools} if tools else None
-        self._llm_settings: SettingsT_co = llm_settings or cast("SettingsT_co", {})
-
         if response_schema and response_schema_by_xml_tag:
             raise ValueError(
                 "Only one of response_schema and response_schema_by_xml_tag can be "
                 "provided, but not both."
             )
-        self._response_schema = response_schema
-        self._response_schema_by_xml_tag = response_schema_by_xml_tag
-
-    @property
-    def model_id(self) -> str:
-        return self._model_id
-
-    @property
-    def model_name(self) -> str | None:
-        return self._model_name
-
-    @property
-    def llm_settings(self) -> SettingsT_co:
-        return self._llm_settings
-
-    @property
-    def response_schema(self) -> Any | None:
-        return self._response_schema
-
-    @response_schema.setter
-    def response_schema(self, response_schema: Any | None) -> None:
-        self._response_schema = response_schema
-
-    @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, Any]] | None:
-        return self._tools
-
-    @tools.setter
-    def tools(self, tools: Sequence[BaseTool[BaseModel, Any, Any]] | None) -> None:
-        self._tools = {t.name: t for t in tools} if tools else None
-
-    def __repr__(self) -> str:
-        return f"{type(self).__name__}[{self.model_id}]; model_name={self._model_name})"
-
-    def _validate_response(self, completion: Completion) -> None:
         parsing_params = {
             "from_substring": False,
             "strip_language_markdown": True,
@@ -138,17 +92,17 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         try:
             for message in completion.messages:
                 if not message.tool_calls:
-                    if self._response_schema:
+                    if response_schema:
                         validate_obj_from_json_or_py_string(
                             message.content or "",
-                            schema=self._response_schema,
+                            schema=response_schema,
                             **parsing_params,
                         )
 
-                    elif self._response_schema_by_xml_tag:
+                    elif response_schema_by_xml_tag:
                         validate_tagged_objs_from_json_or_py_string(
                             message.content or "",
-                            schema_by_xml_tag=self._response_schema_by_xml_tag,
+                            schema_by_xml_tag=response_schema_by_xml_tag,
                             **parsing_params,
                         )
         except JSONSchemaValidationError as exc:
@@ -156,7 +110,9 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
                 exc.s, exc.schema, message=str(exc)
             ) from exc
 
-    def _validate_tool_calls(self, completion: Completion) -> None:
+    def _validate_tool_calls(
+        self, completion: Completion, tools: Mapping[str, BaseTool[BaseModel, Any, Any]]
+    ) -> None:
         parsing_params = {
             "from_substring": False,
             "strip_language_markdown": True,
@@ -167,15 +123,15 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
                     tool_name = tool_call.tool_name
                     tool_arguments = tool_call.tool_arguments
 
-                    available_tool_names = list(self.tools) if self.tools else []
-                    if tool_name not in available_tool_names or not self.tools:
+                    available_tool_names = list(tools) if tools else []
+                    if tool_name not in available_tool_names or not tools:
                         raise LLMToolCallValidationError(
                             tool_name,
                             tool_arguments,
                             message=f"Tool '{tool_name}' is not available in the LLM "
                             f"tools (available: {available_tool_names})",
                         )
-                    tool = self.tools[tool_name]
+                    tool = tools[tool_name]
                     try:
                         validate_obj_from_json_or_py_string(
                             tool_arguments, schema=tool.in_type, **parsing_params
@@ -296,6 +252,9 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         self,
         conversation: Messages,
         *,
+        tools: Mapping[str, BaseTool[BaseModel, Any, Any]] | None = None,
+        response_schema: Any | None = None,
+        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
         tool_choice: ToolChoice | None = None,
         n_choices: int | None = None,
         proc_name: str | None = None,
@@ -308,6 +267,9 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         self,
         conversation: Messages,
         *,
+        tools: Mapping[str, BaseTool[BaseModel, Any, Any]] | None = None,
+        response_schema: Any | None = None,
+        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
         tool_choice: ToolChoice | None = None,
         n_choices: int | None = None,
         proc_name: str | None = None,
@@ -318,5 +280,10 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         pass
 
     @abstractmethod
-    def combine_completion_chunks(self, completion_chunks: list[Any]) -> Any:
+    def combine_completion_chunks(
+        self,
+        completion_chunks: list[Any],
+        response_schema: Any | None = None,
+        tools: Mapping[str, BaseTool[BaseModel, Any, Any]] | None = None,
+    ) -> Any:
         raise NotImplementedError

grasp_agents/llm_agent.py

@@ -1,4 +1,4 @@
-from collections.abc import AsyncIterator, Sequence
+from collections.abc import AsyncIterator, Mapping, Sequence
 from pathlib import Path
 from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
 
@@ -41,7 +41,8 @@ class OutputParser(Protocol[_InT_contra, _OutT_co, CtxT]):
         conversation: Messages,
         *,
         in_args: _InT_contra | None,
-        ctx: RunContext[CtxT] | None,
+        ctx: RunContext[CtxT],
+        call_id: str,
     ) -> _OutT_co: ...
 
 
@@ -68,16 +69,19 @@ class LLMAgent(
         # System prompt template
         sys_prompt: LLMPrompt | None = None,
         sys_prompt_path: str | Path | None = None,
+        # LLM response validation
+        response_schema: Any | None = None,
+        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
         # Agent loop settings
         max_turns: int = 100,
         react_mode: bool = False,
         final_answer_as_tool_call: bool = False,
         # Agent memory management
         reset_memory_on_run: bool = False,
-        # Retries
+        # Agent run retries
         max_retries: int = 0,
         # Multi-agent routing
-        recipients: list[ProcName] | None = None,
+        recipients: Sequence[ProcName] | None = None,
     ) -> None:
         super().__init__(name=name, recipients=recipients, max_retries=max_retries)
 
@@ -96,15 +100,6 @@ class LLMAgent(
 
         # LLM policy executor
 
-        self._used_default_llm_response_schema: bool = False
-        if (
-            llm.response_schema is None
-            and tools is None
-            and not hasattr(type(self), "output_parser")
-        ):
-            llm.response_schema = self.out_type
-            self._used_default_llm_response_schema = True
-
         if issubclass(self._out_type, BaseModel):
             final_answer_type = self._out_type
         elif not final_answer_as_tool_call:
@@ -115,10 +110,21 @@ class LLMAgent(
                 "final_answer_as_tool_call is True."
             )
 
+        self._used_default_llm_response_schema: bool = False
+        if (
+            response_schema is None
+            and tools is None
+            and not hasattr(type(self), "output_parser")
+        ):
+            response_schema = self.out_type
+            self._used_default_llm_response_schema = True
+
         self._policy_executor: LLMPolicyExecutor[CtxT] = LLMPolicyExecutor[CtxT](
             agent_name=self.name,
             llm=llm,
             tools=tools,
+            response_schema=response_schema,
+            response_schema_by_xml_tag=response_schema_by_xml_tag,
             max_turns=max_turns,
             react_mode=react_mode,
             final_answer_type=final_answer_type,
@@ -160,23 +166,33 @@ class LLMAgent(
     def _prepare_memory(
         self,
         memory: LLMAgentMemory,
+        *,
         in_args: InT | None = None,
         sys_prompt: LLMPrompt | None = None,
-        ctx: RunContext[Any] | None = None,
+        ctx: RunContext[Any],
+        call_id: str,
     ) -> None:
         if self.memory_preparator:
             return self.memory_preparator(
-                memory=memory, in_args=in_args, sys_prompt=sys_prompt, ctx=ctx
+                memory=memory,
+                in_args=in_args,
+                sys_prompt=sys_prompt,
+                ctx=ctx,
+                call_id=call_id,
             )
 
     def _memorize_inputs(
         self,
         memory: LLMAgentMemory,
+        *,
         chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
         in_args: InT | None = None,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
+        call_id: str,
     ) -> tuple[SystemMessage | None, UserMessage | None]:
-        formatted_sys_prompt = self._prompt_builder.build_system_prompt(ctx=ctx)
+        formatted_sys_prompt = self._prompt_builder.build_system_prompt(
+            ctx=ctx, call_id=call_id
+        )
 
         system_message: SystemMessage | None = None
         if self._reset_memory_on_run or memory.is_empty:
@@ -185,24 +201,22 @@ class LLMAgent(
                 system_message = cast("SystemMessage", memory.message_history[0])
         else:
             self._prepare_memory(
-                memory=memory, in_args=in_args, sys_prompt=formatted_sys_prompt, ctx=ctx
+                memory=memory,
+                in_args=in_args,
+                sys_prompt=formatted_sys_prompt,
+                ctx=ctx,
+                call_id=call_id,
             )
 
         input_message = self._prompt_builder.build_input_message(
-            chat_inputs=chat_inputs, in_args=in_args, ctx=ctx
+            chat_inputs=chat_inputs, in_args=in_args, ctx=ctx, call_id=call_id
         )
         if input_message:
             memory.update([input_message])
 
         return system_message, input_message
 
-    def _parse_output_default(
-        self,
-        conversation: Messages,
-        *,
-        in_args: InT | None = None,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> OutT:
+    def parse_output_default(self, conversation: Messages) -> OutT:
         return validate_obj_from_json_or_py_string(
             str(conversation[-1].content or ""),
             schema=self._out_type,
@@ -215,16 +229,15 @@ class LLMAgent(
         conversation: Messages,
         *,
         in_args: InT | None = None,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
+        call_id: str,
     ) -> OutT:
         if self.output_parser:
             return self.output_parser(
-                conversation=conversation, in_args=in_args, ctx=ctx
+                conversation=conversation, in_args=in_args, ctx=ctx, call_id=call_id
             )
 
-        return self._parse_output_default(
-            conversation=conversation, in_args=in_args, ctx=ctx
-        )
+        return self.parse_output_default(conversation)
 
     async def _process(
         self,
@@ -232,24 +245,28 @@ class LLMAgent(
         *,
         in_args: InT | None = None,
         memory: LLMAgentMemory,
+        ctx: RunContext[CtxT],
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
     ) -> OutT:
         system_message, input_message = self._memorize_inputs(
             memory=memory,
             chat_inputs=chat_inputs,
             in_args=in_args,
             ctx=ctx,
+            call_id=call_id,
         )
         if system_message:
-            self._print_messages([system_message], call_id=call_id, ctx=ctx)
+            self._print_messages([system_message], ctx=ctx, call_id=call_id)
         if input_message:
-            self._print_messages([input_message], call_id=call_id, ctx=ctx)
+            self._print_messages([input_message], ctx=ctx, call_id=call_id)
 
-        await self._policy_executor.execute(memory, call_id=call_id, ctx=ctx)
+        await self._policy_executor.execute(memory, ctx=ctx, call_id=call_id)
 
         return self._parse_output(
-            conversation=memory.message_history, in_args=in_args, ctx=ctx
+            conversation=memory.message_history,
+            in_args=in_args,
+            ctx=ctx,
+            call_id=call_id,
         )
 
     async def _process_stream(
@@ -258,41 +275,45 @@ class LLMAgent(
         *,
         in_args: InT | None = None,
         memory: LLMAgentMemory,
+        ctx: RunContext[CtxT],
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         system_message, input_message = self._memorize_inputs(
             memory=memory,
             chat_inputs=chat_inputs,
             in_args=in_args,
             ctx=ctx,
+            call_id=call_id,
         )
         if system_message:
-            self._print_messages([system_message], call_id=call_id, ctx=ctx)
+            self._print_messages([system_message], ctx=ctx, call_id=call_id)
             yield SystemMessageEvent(
                 data=system_message, proc_name=self.name, call_id=call_id
             )
         if input_message:
-            self._print_messages([input_message], call_id=call_id, ctx=ctx)
+            self._print_messages([input_message], ctx=ctx, call_id=call_id)
             yield UserMessageEvent(
                 data=input_message, proc_name=self.name, call_id=call_id
             )
 
         async for event in self._policy_executor.execute_stream(
-            memory, call_id=call_id, ctx=ctx
+            memory, ctx=ctx, call_id=call_id
         ):
             yield event
 
         output = self._parse_output(
-            conversation=memory.message_history, in_args=in_args, ctx=ctx
+            conversation=memory.message_history,
+            in_args=in_args,
+            ctx=ctx,
+            call_id=call_id,
         )
         yield ProcPayloadOutputEvent(data=output, proc_name=self.name, call_id=call_id)
 
     def _print_messages(
         self,
         messages: Sequence[Message],
+        ctx: RunContext[CtxT],
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
     ) -> None:
         if ctx and ctx.printer:
             ctx.printer.print_messages(messages, agent_name=self.name, call_id=call_id)
@@ -321,28 +342,31 @@ class LLMAgent(
         if cur_cls.memory_manager is not base_cls.memory_manager:
             self._policy_executor.memory_manager = self.memory_manager
 
-    def system_prompt_builder(self, ctx: RunContext[CtxT] | None = None) -> str | None:
+    def system_prompt_builder(self, ctx: RunContext[CtxT], call_id: str) -> str | None:
         if self._prompt_builder.system_prompt_builder is not None:
-            return self._prompt_builder.system_prompt_builder(ctx=ctx)
+            return self._prompt_builder.system_prompt_builder(ctx=ctx, call_id=call_id)
         raise NotImplementedError("System prompt builder is not implemented.")
 
     def input_content_builder(
-        self, in_args: InT | None = None, *, ctx: RunContext[CtxT] | None = None
+        self, in_args: InT, ctx: RunContext[CtxT], call_id: str
     ) -> Content:
         if self._prompt_builder.input_content_builder is not None:
-            return self._prompt_builder.input_content_builder(in_args=in_args, ctx=ctx)
+            return self._prompt_builder.input_content_builder(
+                in_args=in_args, ctx=ctx, call_id=call_id
+            )
         raise NotImplementedError("Input content builder is not implemented.")
 
     def tool_call_loop_terminator(
         self,
         conversation: Messages,
         *,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
+        call_id: str,
         **kwargs: Any,
     ) -> bool:
         if self._policy_executor.tool_call_loop_terminator is not None:
             return self._policy_executor.tool_call_loop_terminator(
-                conversation=conversation, ctx=ctx, **kwargs
+                conversation=conversation, ctx=ctx, call_id=call_id, **kwargs
             )
         raise NotImplementedError("Tool call loop terminator is not implemented.")
 
@@ -350,12 +374,13 @@ class LLMAgent(
         self,
         memory: LLMAgentMemory,
         *,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
+        call_id: str,
         **kwargs: Any,
     ) -> None:
         if self._policy_executor.memory_manager is not None:
             return self._policy_executor.memory_manager(
-                memory=memory, ctx=ctx, **kwargs
+                memory=memory, ctx=ctx, call_id=call_id, **kwargs
             )
         raise NotImplementedError("Memory manager is not implemented.")
 
@@ -391,12 +416,11 @@ class LLMAgent(
         self, func: OutputParser[InT, OutT, CtxT]
     ) -> OutputParser[InT, OutT, CtxT]:
         if self._used_default_llm_response_schema:
-            self._policy_executor.llm.response_schema = None
+            self._policy_executor.response_schema = None
         self.output_parser = func
 
         return func
 
     def add_memory_preparator(self, func: MemoryPreparator) -> MemoryPreparator:
         self.memory_preparator = func
-
         return func

grasp_agents/llm_agent_memory.py

@@ -13,14 +13,16 @@ class MemoryPreparator(Protocol):
     def __call__(
         self,
         memory: "LLMAgentMemory",
+        *,
         in_args: Any | None,
         sys_prompt: LLMPrompt | None,
-        ctx: RunContext[Any] | None,
+        ctx: RunContext[Any],
+        call_id: str,
     ) -> None: ...
 
 
 class LLMAgentMemory(Memory):
-    _message_history: Messages = PrivateAttr(default_factory=list)  # type: ignore
+    _message_history: Messages = PrivateAttr(default_factory=Messages)
     _sys_prompt: LLMPrompt | None = PrivateAttr(default=None)
 
     def __init__(self, sys_prompt: LLMPrompt | None = None) -> None: