grasp_agents 0.5.8__py3-none-any.whl → 0.5.10__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
@@ -25,7 +26,7 @@ from .typing.events import (
     AnnotationsEndEvent,
     AnnotationsStartEvent,
     CompletionChunkEvent,
-    CompletionEndEvent,
+    # CompletionEndEvent,
     CompletionEvent,
     CompletionStartEvent,
     LLMStateChangeEvent,
@@ -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
@@ -196,7 +152,9 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         annotations_op_evt: AnnotationsChunkEvent | None = None
         tool_calls_op_evt: ToolCallChunkEvent | None = None
 
-        def _close_open_events() -> list[LLMStateChangeEvent[Any]]:
+        def _close_open_events(
+            _event: CompletionChunkEvent[CompletionChunk] | None = None,
+        ) -> list[LLMStateChangeEvent[Any]]:
             nonlocal \
                 chunk_op_evt, \
                 thinking_op_evt, \
@@ -206,26 +164,21 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
 
             events: list[LLMStateChangeEvent[Any]] = []
 
-            if tool_calls_op_evt:
+            if not isinstance(_event, ThinkingChunkEvent) and thinking_op_evt:
+                events.append(ThinkingEndEvent.from_chunk_event(thinking_op_evt))
+                thinking_op_evt = None
+
+            if not isinstance(_event, ToolCallChunkEvent) and tool_calls_op_evt:
                 events.append(ToolCallEndEvent.from_chunk_event(tool_calls_op_evt))
+                tool_calls_op_evt = None
 
-            if response_op_evt:
+            if not isinstance(_event, ResponseChunkEvent) and response_op_evt:
                 events.append(ResponseEndEvent.from_chunk_event(response_op_evt))
+                response_op_evt = None
 
-            if thinking_op_evt:
-                events.append(ThinkingEndEvent.from_chunk_event(thinking_op_evt))
-
-            if annotations_op_evt:
+            if not isinstance(_event, AnnotationsChunkEvent) and annotations_op_evt:
                 events.append(AnnotationsEndEvent.from_chunk_event(annotations_op_evt))
-
-            if chunk_op_evt:
-                events.append(CompletionEndEvent.from_chunk_event(chunk_op_evt))
-
-            chunk_op_evt = None
-            thinking_op_evt = None
-            tool_calls_op_evt = None
-            response_op_evt = None
-            annotations_op_evt = None
+                annotations_op_evt = None
 
             return events
 
@@ -252,14 +205,14 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
                 sub_events = event.split_into_specialized()
 
                 for sub_event in sub_events:
+                    for close_event in _close_open_events(sub_event):
+                        yield close_event
+
                     if isinstance(sub_event, ThinkingChunkEvent):
                         if not thinking_op_evt:
                             thinking_op_evt = sub_event
                             yield ThinkingStartEvent.from_chunk_event(sub_event)
                         yield sub_event
-                    elif thinking_op_evt:
-                        yield ThinkingEndEvent.from_chunk_event(thinking_op_evt)
-                        thinking_op_evt = None
 
                     if isinstance(sub_event, ToolCallChunkEvent):
                         tc = sub_event.data.tool_call
@@ -273,27 +226,18 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
                             tool_calls_op_evt = sub_event
                             yield ToolCallStartEvent.from_chunk_event(sub_event)
                         yield sub_event
-                    elif tool_calls_op_evt:
-                        yield ToolCallEndEvent.from_chunk_event(tool_calls_op_evt)
-                        tool_calls_op_evt = None
 
                     if isinstance(sub_event, ResponseChunkEvent):
                         if not response_op_evt:
                             response_op_evt = sub_event
                             yield ResponseStartEvent.from_chunk_event(sub_event)
                         yield sub_event
-                    elif response_op_evt:
-                        yield ResponseEndEvent.from_chunk_event(response_op_evt)
-                        response_op_evt = None
 
                     if isinstance(sub_event, AnnotationsChunkEvent):
                         if not annotations_op_evt:
                             annotations_op_evt = sub_event
                             yield AnnotationsStartEvent.from_chunk_event(sub_event)
                         yield sub_event
-                    elif annotations_op_evt:
-                        yield AnnotationsEndEvent.from_chunk_event(annotations_op_evt)
-                        annotations_op_evt = None
 
                 prev_completion_id = chunk.id
 
@@ -308,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,
@@ -320,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,
@@ -330,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,7 @@ class OutputParser(Protocol[_InT_contra, _OutT_co, CtxT]):
         conversation: Messages,
         *,
         in_args: _InT_contra | None,
-        ctx: RunContext[CtxT] | None,
+        ctx: RunContext[CtxT],
     ) -> _OutT_co: ...
 
 
@@ -68,16 +68,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 +99,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 +109,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,9 +165,10 @@ 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],
     ) -> None:
         if self.memory_preparator:
             return self.memory_preparator(
@@ -172,9 +178,10 @@ class LLMAgent(
     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],
     ) -> tuple[SystemMessage | None, UserMessage | None]:
         formatted_sys_prompt = self._prompt_builder.build_system_prompt(ctx=ctx)
 
@@ -201,7 +208,7 @@ class LLMAgent(
         conversation: Messages,
         *,
         in_args: InT | None = None,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> OutT:
         return validate_obj_from_json_or_py_string(
             str(conversation[-1].content or ""),
@@ -215,7 +222,7 @@ class LLMAgent(
         conversation: Messages,
         *,
         in_args: InT | None = None,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> OutT:
         if self.output_parser:
             return self.output_parser(
@@ -233,7 +240,7 @@ class LLMAgent(
         in_args: InT | None = None,
         memory: LLMAgentMemory,
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> OutT:
         system_message, input_message = self._memorize_inputs(
             memory=memory,
@@ -259,7 +266,7 @@ class LLMAgent(
         in_args: InT | None = None,
         memory: LLMAgentMemory,
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> AsyncIterator[Event[Any]]:
         system_message, input_message = self._memorize_inputs(
             memory=memory,
@@ -292,7 +299,7 @@ class LLMAgent(
         self,
         messages: Sequence[Message],
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
     ) -> None:
         if ctx and ctx.printer:
             ctx.printer.print_messages(messages, agent_name=self.name, call_id=call_id)
@@ -321,24 +328,18 @@ 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]) -> str | None:
         if self._prompt_builder.system_prompt_builder is not None:
             return self._prompt_builder.system_prompt_builder(ctx=ctx)
         raise NotImplementedError("System prompt builder is not implemented.")
 
-    def input_content_builder(
-        self, in_args: InT | None = None, *, ctx: RunContext[CtxT] | None = None
-    ) -> Content:
+    def input_content_builder(self, in_args: InT, ctx: RunContext[CtxT]) -> Content:
         if self._prompt_builder.input_content_builder is not None:
             return self._prompt_builder.input_content_builder(in_args=in_args, ctx=ctx)
         raise NotImplementedError("Input content builder is not implemented.")
 
     def tool_call_loop_terminator(
-        self,
-        conversation: Messages,
-        *,
-        ctx: RunContext[CtxT] | None = None,
-        **kwargs: Any,
+        self, conversation: Messages, *, ctx: RunContext[CtxT], **kwargs: Any
     ) -> bool:
         if self._policy_executor.tool_call_loop_terminator is not None:
             return self._policy_executor.tool_call_loop_terminator(
@@ -347,11 +348,7 @@ class LLMAgent(
         raise NotImplementedError("Tool call loop terminator is not implemented.")
 
     def memory_manager(
-        self,
-        memory: LLMAgentMemory,
-        *,
-        ctx: RunContext[CtxT] | None = None,
-        **kwargs: Any,
+        self, memory: LLMAgentMemory, *, ctx: RunContext[CtxT], **kwargs: Any
     ) -> None:
         if self._policy_executor.memory_manager is not None:
             return self._policy_executor.memory_manager(
@@ -391,12 +388,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,15 @@ class MemoryPreparator(Protocol):
     def __call__(
         self,
         memory: "LLMAgentMemory",
+        *,
         in_args: Any | None,
         sys_prompt: LLMPrompt | None,
-        ctx: RunContext[Any] | None,
+        ctx: RunContext[Any],
     ) -> 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: