grasp_agents 0.3.8__tar.gz → 0.3.9__tar.gz

{grasp_agents-0.3.8 → grasp_agents-0.3.9}/PKG-INFO

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

{grasp_agents-0.3.8 → grasp_agents-0.3.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "grasp_agents"
-version = "0.3.8"
+version = "0.3.9"
 description = "Grasp Agents Library"
 readme = "README.md"
 requires-python = ">=3.11.4,<4"

{grasp_agents-0.3.8 → grasp_agents-0.3.9}/src/grasp_agents/llm_policy_executor.py

@@ -1,7 +1,7 @@
 import asyncio
 import json
 from collections.abc import AsyncIterator, Coroutine, Sequence
-from itertools import starmap
+from itertools import chain, starmap
 from logging import getLogger
 from typing import Any, ClassVar, Generic, Protocol, TypeVar
 
@@ -141,19 +141,14 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
         completion_batch = await self.llm.generate_completion_batch(
             memory.message_history, tool_choice=tool_choice
         )
-        if (
-            len(completion_batch[0].messages) > 1
-            and memory.message_history.batch_size > 1
-        ):
-            raise ValueError(
-                "Batch size must be 1 when generating completions with n>1."
-            )
-        message_batch = [c.messages[0] for c in completion_batch]
+        message_batch = list(
+            chain.from_iterable([c.messages for c in completion_batch])
+        )
         memory.update(message_batch=message_batch)
 
         if ctx is not None:
             ctx.completions[self.agent_name].extend(completion_batch)
-            self._track_usage(completion_batch, ctx=ctx)
+            self._track_usage(self.agent_name, completion_batch, ctx=ctx)
             self._print_completions(completion_batch, ctx=ctx)
 
         return message_batch
@@ -179,16 +174,14 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
 
         if completion is None:
             raise RuntimeError("No completion generated during stream.")
-        if len(completion.messages) > 1:
-            raise ValueError("Streaming completion must have n=1")
 
-        message = completion.messages[0]
-        memory.update(message_batch=[message])
+        memory.update(message_batch=completion.messages)
 
-        yield GenMessageEvent(name=self.agent_name, data=message)
+        for message in completion.messages:
+            yield GenMessageEvent(name=self.agent_name, data=message)
 
         if ctx is not None:
-            self._track_usage([completion], ctx=ctx)
+            self._track_usage(self.agent_name, [completion], ctx=ctx)
             ctx.completions[self.agent_name].append(completion)
 
     async def call_tools(
@@ -207,7 +200,7 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
         tool_messages = list(
             starmap(ToolMessage.from_tool_output, zip(outs, calls, strict=False))
         )
-        memory.update(message_list=tool_messages)
+        memory.update(tool_messages)
 
         if ctx is not None:
             ctx.printer.print_llm_messages(tool_messages, agent_name=self.agent_name)
@@ -234,7 +227,7 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
                     name=self.agent_name, content=tool_call.tool_arguments
                 )
                 gen_message.tool_calls = None
-                memory.update(message_list=[final_answer_message])
+                memory.update([final_answer_message])
                 return final_answer_message
 
         return final_answer_message
@@ -247,7 +240,7 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
         user_message = UserMessage.from_text(
             "Exceeded the maximum number of turns: provide a final answer now!"
         )
-        memory.update(message_list=[user_message])
+        memory.update([user_message])
         if ctx is not None:
             ctx.printer.print_llm_messages([user_message], agent_name=self.agent_name)
 
@@ -263,6 +256,7 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
             raise RuntimeError(
                 "Final answer tool call did not return a final answer message."
             )
+
         return final_answer_message
 
     async def _generate_final_answer_stream(
@@ -273,7 +267,7 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
         user_message = UserMessage.from_text(
             "Exceeded the maximum number of turns: provide a final answer now!",
         )
-        memory.update(message_list=[user_message])
+        memory.update([user_message])
         yield UserMessageEvent(name=self.agent_name, data=user_message)
 
         tool_choice = NamedToolChoice(name=self._final_answer_tool_name)
@@ -447,10 +441,15 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
             turns += 1
 
     def _track_usage(
-        self, completion_batch: Sequence[Completion], ctx: RunContext[CtxT]
+        self,
+        agent_name: str,
+        completion_batch: Sequence[Completion],
+        ctx: RunContext[CtxT],
     ) -> None:
         ctx.usage_tracker.update(
-            completions=completion_batch, model_name=self.llm.model_name
+            agent_name=agent_name,
+            completions=completion_batch,
+            model_name=self.llm.model_name,
         )
 
     def get_final_answer_tool(self) -> BaseTool[BaseModel, None, Any]:

{grasp_agents-0.3.8 → grasp_agents-0.3.9}/src/grasp_agents/usage_tracker.py

@@ -1,4 +1,5 @@
 import logging
+from collections import defaultdict
 from collections.abc import Sequence
 from pathlib import Path
 from typing import Any, TypeAlias
@@ -23,7 +24,7 @@ class UsageTracker(BaseModel):
     source_id: str
     costs_dict_path: str | Path = COSTS_DICT_PATH
     costs_dict: CostsDict | None = None
-    total_usage: Usage = Field(default_factory=Usage)
+    usages: dict[str, Usage] = Field(default_factory=dict)
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
@@ -58,7 +59,10 @@ class UsageTracker(BaseModel):
         usage.cost = (input_cost + output_cost + reasoning_cost + cached_cost) / 1e6
 
     def update(
-        self, completions: Sequence[Completion], model_name: str | None = None
+        self,
+        agent_name: str,
+        completions: Sequence[Completion],
+        model_name: str | None = None,
     ) -> None:
         if model_name is not None and self.costs_dict is not None:
             model_costs_dict = self.costs_dict.get(model_name.split(":", 1)[-1])
@@ -71,10 +75,16 @@ class UsageTracker(BaseModel):
                     self._add_cost_to_usage(
                         usage=completion.usage, model_costs_dict=model_costs_dict
                     )
-                self.total_usage += completion.usage
+                if agent_name not in self.usages:
+                    self.usages[agent_name] = Usage()
+                self.usages[agent_name] += completion.usage
+
+    @property
+    def total_usage(self) -> Usage:
+        return sum((usage for usage in self.usages.values()), Usage())
 
     def reset(self) -> None:
-        self.total_usage = Usage()
+        self.usages = defaultdict(Usage)
 
     def print_usage(self) -> None:
         usage = self.total_usage

{grasp_agents-0.3.8 → grasp_agents-0.3.9}/src/grasp_agents/llm_agent_memory.py

@@ -30,9 +30,9 @@ class LLMAgentMemory(Memory):
 
     def update(
         self,
+        message_list: Sequence[Message] | None = None,
         *,
         message_batch: Sequence[Message] | None = None,
-        message_list: Sequence[Message] | None = None,
         ctx: RunContext[Any] | None = None,
     ):
         if message_batch is not None and message_list is not None: