finalise

Author: CTY-git

patchwork/common/multiturn_strategy/agentic_strategy.py

@@ -23,6 +23,8 @@ def __init__(self, llm_client: LlmClient, tool_set: dict[str, Tool]):
         self.llm_client = llm_client
         self.tool_set = tool_set
         self.history: list[ChatCompletionMessageParam] = []
+        self.request_tokens = 0
+        self.response_tokens = 0
 
     def generate_reply(self, message: str) -> str:
         self.history.append(dict(role="user", content=message))
@@ -36,6 +38,9 @@ def generate_reply(self, message: str) -> str:
         if is_prompt_safe < 0:
             raise ValueError("The subsequent prompt is not supported, due to large size.")
         response = self.llm_client.chat_completion(**input_kwargs)
+        self.response_tokens = response.usage.completion_tokens
+        self.request_tokens = response.usage.prompt_tokens
+
         choices = response.choices or []
 
         message_content = ""
@@ -145,6 +150,17 @@ def __is_session_completed(self) -> bool:
 
         return False
 
+    def usage(self):
+        request_tokens = 0
+        response_tokens = 0
+        for role in [self.__assistant_role, self.__user_role]:
+            request_tokens += role.request_tokens
+            response_tokens += role.response_tokens
+        return {
+            "request_tokens": request_tokens,
+            "response_tokens": response_tokens,
+        }
+
     def execute(self, limit: int | None = None) -> None:
         message = self.__render_prompt(self.__user_prompt_template)
         try:

patchwork/common/multiturn_strategy/agentic_strategy_v2.py

@@ -7,6 +7,7 @@
 from pydantic import BaseModel
 from pydantic_ai import Agent
 from pydantic_ai.models.anthropic import AnthropicModel
+from pydantic_ai.result import RunResult
 from typing_extensions import Any, Dict, Optional, Union
 
 from patchwork.common.client.llm.utils import example_json_to_base_model
@@ -31,14 +32,14 @@ class Config:
 
 class AgenticStrategyV2:
     def __init__(
-            self,
-            api_key: str,
-            template_data: dict[str, str],
-            system_prompt_template: str,
-            user_prompt_template: str,
-            agent_configs: list[AgentConfig],
-            example_json: Union[str, dict[str, Any]] = '{"output":"output text"}',
-            limit: Optional[int] = None,
+        self,
+        api_key: str,
+        template_data: dict[str, str],
+        system_prompt_template: str,
+        user_prompt_template: str,
+        agent_configs: list[AgentConfig],
+        example_json: Union[str, dict[str, Any]] = '{"output":"output text"}',
+        limit: Optional[int] = None,
     ):
         self.__limit = limit
         self.__template_data = template_data
@@ -66,6 +67,19 @@ def __init__(
 
             self.__agents.append(agent)
 
+        self.__request_tokens = 0
+        self.__response_tokens = 0
+
+    def reset(self):
+        self.__request_tokens = 0
+        self.__response_tokens = 0
+
+    def usage(self):
+        return {
+            "request_tokens": self.__request_tokens,
+            "response_tokens": self.__response_tokens,
+        }
+
     def execute(self, limit: Optional[int] = None) -> dict:
         agents_result = dict()
         loop = asyncio.new_event_loop()
@@ -75,8 +89,13 @@ def execute(self, limit: Optional[int] = None) -> dict:
                 message_history = None
                 agent_output = None
                 for i in range(limit or self.__limit or sys.maxsize):
-                    agent_output = loop.run_until_complete(agent.run(user_message, message_history=message_history))
+                    agent_output: RunResult[Any] = loop.run_until_complete(
+                        agent.run(user_message, message_history=message_history)
+                    )
                     message_history = agent_output.all_messages()
+                    self.__request_tokens += agent_output.usage().request_tokens or 0
+                    self.__response_tokens += agent_output.usage().response_tokens or 0
+
                     if getattr(agent_output.data, _COMPLETION_FLAG_ATTRIBUTE, False):
                         break
                     user_message = "Please continue"
@@ -88,27 +107,38 @@ def execute(self, limit: Optional[int] = None) -> dict:
             return dict()
 
         if len(agents_result) == 1:
-            final_result = loop.run_until_complete(self.__summariser.run(
-                "From the actions taken by the assistant. Please give me the result.",
-                message_history=next(v for _, v in agents_result.items()).all_messages(),
-            ))
+            final_result = loop.run_until_complete(
+                self.__summariser.run(
+                    "From the actions taken by the assistant. Please give me the result.",
+                    message_history=next(v for _, v in agents_result.items()).all_messages(),
+                )
+            )
         else:
             agent_summaries = []
             for agent_result in agents_result.values():
-                agent_summary_result = loop.run_until_complete(self.__summariser.run(
-                    "From the actions taken by the assistant. Please give me the result.",
-                    message_history=agent_result.all_messages(),
-                ))
+                agent_summary_result = loop.run_until_complete(
+                    self.__summariser.run(
+                        "From the actions taken by the assistant. Please give me the result.",
+                        message_history=agent_result.all_messages(),
+                    )
+                )
+                self.__request_token += agent_summary_result.usage().request_tokens or 0
+                self.__response_token += agent_summary_result.usage().response_tokens or 0
+
                 agent_summary = getattr(agent_summary_result.data, _MESSAGE_ATTRIBUTE, None)
                 if agent_summary is None:
                     continue
 
                 agent_summaries.append(agent_summary)
             agent_summary_list = "\n* " + "\n* ".join(agent_summaries)
-            final_result = loop.run_until_complete(self.__summariser.run(
-                "Please give me the result from the following summary of what the assistants have done."
-                + agent_summary_list,
-            ))
+            final_result = loop.run_until_complete(
+                self.__summariser.run(
+                    "Please give me the result from the following summary of what the assistants have done."
+                    + agent_summary_list,
+                )
+            )
+        self.__request_tokens += final_result.usage().request_tokens or 0
+        self.__response_tokens += final_result.usage().response_tokens or 0
 
         loop.close()
         return final_result.data.dict()

patchwork/common/multiturn_strategy/analyze_implement.py

@@ -40,6 +40,14 @@ def _reset(self):
         self._stage = STAGE.ANALYSIS
         self.run_count = 0
         self.stage_run_counts = defaultdict(int)
+        self.__request_tokens = 0
+        self.__response_tokens = 0
+
+    def usage(self):
+        return dict(
+            request_tokens=self.__request_tokens,
+            response_tokens=self.__response_tokens,
+        )
 
     def __run_prompt(self, messages: list[ChatCompletionMessageParam]) -> list[ChatCompletionMessageParam]:
         input_kwargs = dict(
@@ -52,6 +60,8 @@ def __run_prompt(self, messages: list[ChatCompletionMessageParam]) -> list[ChatC
         if is_prompt_safe < 0:
             raise ValueError("The subsequent prompt is not supported, due to large size.")
         response = self.llm_client.chat_completion(**input_kwargs)
+        self.__request_tokens += response.usage.prompt_tokens
+        self.__response_tokens += response.usage.response_tokens
         new_messages = [choice.message.to_dict() for choice in response.choices]
         messages.extend(new_messages)
         return messages

patchwork/steps/AgenticLLM/AgenticLLM.py

@@ -27,4 +27,5 @@ def run(self) -> dict:
         return dict(
             conversation_history=self.agentic_strategy.history,
             tool_records=self.agentic_strategy.tool_records,
+            **self.agentic_strategy.usage(),
         )

patchwork/steps/AgenticLLM/typed.py

@@ -38,5 +38,5 @@ class AgenticLLMInputs(TypedDict, total=False):
 class AgenticLLMOutputs(TypedDict):
     conversation_history: List[Dict]
     tool_records: List[Dict]
-    # request_tokens: int
-    # response_tokens: int
+    request_tokens: int
+    response_tokens: int

patchwork/steps/AgenticLLMV2/AgenticLLMV2.py

@@ -32,4 +32,5 @@ def __init__(self, inputs):
         )
 
     def run(self) -> dict:
-        return self.agentic_strategy.execute(limit=self.conversation_limit)
+        result = self.agentic_strategy.execute(limit=self.conversation_limit)
+        return {**result, **self.agentic_strategy.usage()}

patchwork/steps/AgenticLLMV2/typed.py

@@ -15,4 +15,5 @@ class AgenticLLMV2Inputs(TypedDict, total=False):
 
 
 class AgenticLLMV2Outputs(TypedDict):
-    pass
+    request_tokens: int
+    response_tokens: int

patchwork/steps/FixIssue/FixIssue.py

@@ -178,4 +178,4 @@ def run(self):
                     # Git-specific errors (untracked files, etc) - keep empty diff
                     logger.warning(f"Could not get git diff for {file}: {str(e)}")
 
-        return dict(modified_files=modified_files)
+        return dict(modified_files=modified_files, **self.multiturn_llm_call.usage())

patchwork/steps/FixIssue/typed.py

@@ -55,3 +55,5 @@ class ModifiedFile(TypedDict):
 
 class FixIssueOutputs(TypedDict):
     modified_files: List[ModifiedFile]
+    request_tokens: int
+    response_tokens: int