add v2 versioning

Author: CTY-git

patchwork/common/multiturn_strategy/agentic_strategy.py

@@ -1,111 +1,173 @@
 from __future__ import annotations
 
+import abc
+import json
 import logging
+import random
+import string
 import sys
+from json import JSONDecodeError
+from pathlib import Path
 
-from pydantic import BaseModel
-from pydantic_ai import Agent
-from pydantic_ai.models.anthropic import AnthropicModel
-from typing_extensions import Any, Dict, Optional, Union
+import chevron
+from openai.types.chat import ChatCompletionMessageParam
+from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
 
-from patchwork.common.client.llm.utils import example_json_to_base_model
-from patchwork.common.tools import Tool
-from patchwork.common.utils.utils import mustache_render
+from patchwork.common.client.llm.protocol import LlmClient
+from patchwork.common.tools import CodeEditTool, Tool
+from patchwork.common.tools.agentic_tools import EndTool
 
-_COMPLETION_FLAG_ATTRIBUTE = "is_task_completed"
-_MESSAGE_ATTRIBUTE = "message"
 
+class Role(abc.ABC):
+    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] = []
 
-class AgentConfig(BaseModel):
-    class Config:
-        arbitrary_types_allowed = True
+    def generate_reply(self, message: str) -> str:
+        self.history.append(dict(role="user", content=message))
+        input_kwargs = dict(
+            messages=self.history,
+            model="claude-3-5-sonnet-latest",
+            tools=self.__get_tools_spec(),
+            max_tokens=8096,
+        )
+        is_prompt_safe = self.llm_client.is_prompt_supported(**input_kwargs)
+        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)
+        choices = response.choices or []
+
+        message_content = ""
+        for choice in choices:
+            new_message = choice.message.to_dict()
+            self.history.append(new_message)
+            if new_message.get("tool_calls") is not None:
+                self.history.extend(self.__execute_tools(new_message))
+            else:
+                message_content = new_message["content"]
+
+        return message_content
+
+    def __execute_tools(self, last_message: ChatCompletionMessageParam) -> list[ChatCompletionMessageParam]:
+        rv = []
+        for tool_call in last_message.get("tool_calls", []):
+            tool_name_to_use = tool_call.get("function", {}).get("name")
+            tool_to_use = self.tool_set.get(tool_name_to_use, None)
+            if tool_to_use is None:
+                logging.info("LLM just used an non-existent tool!")
+                continue
 
-    name: str
-    tool_set: Dict[str, Tool]
-    system_prompt: str = ""
-    example_json: Union[
-        str, Dict[str, Any]
-    ] = f'{{"{_MESSAGE_ATTRIBUTE}":"message", "{_COMPLETION_FLAG_ATTRIBUTE}": false}}'
+            logging.info(f"Running tool: {tool_name_to_use}")
+            try:
+                tool_arguments = tool_call.get("function", {}).get("arguments", "{}")
+                tool_kwargs = json.loads(tool_arguments)
+                tool_output = tool_to_use.execute(**tool_kwargs)
+            except JSONDecodeError:
+                tool_output = "Arguments must be passed through a valid JSON object"
+
+            rv.append({"tool_call_id": tool_call.get("id", ""), "role": "tool", "content": tool_output})
+
+        return rv
+
+    def __get_tools_spec(self) -> list[ChatCompletionToolParam]:
+        return [
+            dict(
+                type="function",
+                function={"name": k, **v.json_schema},
+            )
+            for k, v in self.tool_set.items()
+        ]
+
+
+class UserProxy(Role):
+    def __init__(
+        self, llm_client: LlmClient, tool_set: dict[str, Tool], system_prompt: str = None, reply_message: str = ""
+    ):
+        super().__init__(llm_client, tool_set)
+        if system_prompt is not None:
+            self.history.append(dict(role="system", content=system_prompt))
+
+        self.__reply_message = reply_message
+
+    def generate_reply(self, message: str) -> str:
+        if self.__reply_message is not None:
+            self.history.append(dict(role="user", content=message))
+            self.history.append(dict(role="assistant", content=self.__reply_message))
+            return self.__reply_message
+        else:
+            return super().generate_reply(message)
+
+
+class Assistant(Role):
+    def __init__(self, llm_client: LlmClient, tool_set: dict[str, Tool], system_prompt: str = None):
+        super().__init__(llm_client, tool_set)
+        if system_prompt is not None:
+            self.history.append(dict(role="system", content=system_prompt))
 
 
 class AgenticStrategy:
     def __init__(
         self,
-        api_key: str,
+        llm_client: LlmClient,
+        tool_set: dict[str, Tool],
         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,
+        *args,
+        **kwargs,
     ):
-        self.__limit = limit
+        self.tool_set = dict(end=EndTool(), **tool_set)
         self.__template_data = template_data
         self.__user_prompt_template = user_prompt_template
-        model = AnthropicModel("claude-3-5-sonnet-latest", api_key=api_key)
-        self.__summariser = Agent(
-            model,
-            system_prompt=mustache_render(system_prompt_template, self.__template_data),
-            result_type=example_json_to_base_model(example_json),
-            model_settings=dict(parallel_tool_calls=False),
+        self.__assistant_role = Assistant(llm_client, self.tool_set, self.__render_prompt(system_prompt_template))
+        self.__user_role = UserProxy(llm_client, dict())
+
+    def __render_prompt(self, prompt_template: str) -> str:
+        chevron.render.__globals__["_html_escape"] = lambda x: x
+        return chevron.render(
+            template=prompt_template,
+            data=self.__template_data,
+            partials_path=None,
+            partials_ext="".join(random.choices(string.ascii_uppercase + string.digits, k=32)),
+            partials_dict=dict(),
         )
-        self.__agents = []
-        for agent_config in agent_configs:
-            tools = []
-            for tool in agent_config.tool_set.values():
-                tools.append(tool.to_pydantic_ai_function_tool())
-            agent = Agent(
-                model,
-                name=agent_config.name,
-                system_prompt=mustache_render(agent_config.system_prompt, self.__template_data),
-                tools=tools,
-                result_type=example_json_to_base_model(agent_config.example_json),
-                model_settings=dict(parallel_tool_calls=False),
-            )
 
-            self.__agents.append(agent)
+    def __get_initial_prompt(self) -> list[ChatCompletionMessageParam]:
+        return [dict(role="user", content=self.__render_prompt(self.__user_prompt_template))]
 
-    def execute(self, limit: Optional[int] = None) -> dict:
-        agents_result = dict()
+    def __is_session_completed(self) -> bool:
+        for message in reversed(self.__assistant_role.history):
+            if message.get("tool") is not None:
+                continue
+            if message.get("content") == EndTool.MESSAGE:
+                return True
+
+        return False
+
+    def execute(self, limit: int | None = None) -> None:
+        message = self.__render_prompt(self.__user_prompt_template)
         try:
-            for index, agent in enumerate(self.__agents):
-                user_message = mustache_render(self.__user_prompt_template, self.__template_data)
-                message_history = None
-                agent_output = None
-                for i in range(limit or self.__limit or sys.maxsize):
-                    agent_output = agent.run_sync(user_message, message_history=message_history)
-                    message_history = agent_output.all_messages()
-                    if getattr(agent_output.data, _COMPLETION_FLAG_ATTRIBUTE, False):
-                        break
-                    user_message = "Please continue"
-                agents_result[index] = agent_output
+            for i in range(limit or self.__limit or sys.maxsize):
+                self.run_count = i + 1
+                for role in [self.__assistant_role, self.__user_role]:
+                    message = role.generate_reply(message)
+                if self.__is_session_completed():
+                    break
         except Exception as e:
             logging.error(e)
+        finally:
+            self.run_count = 0
 
-        if len(agents_result) == 0:
-            return dict()
-
-        if len(agents_result) == 1:
-            final_result = self.__summariser.run_sync(
-                "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 = self.__summariser.run_sync(
-                    "From the actions taken by the assistant. Please give me the result.",
-                    message_history=agent_result.all_messages(),
-                )
-                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 = self.__summariser.run_sync(
-                "Please give me the result from the following summary of what the assistants have done."
-                + agent_summary_list,
-            )
+    @property
+    def history(self):
+        return self.__user_role.history
 
-        return final_result.data.dict()
+    @property
+    def tool_records(self):
+        for tool in self.tool_set.values():
+            if isinstance(tool, CodeEditTool):
+                cwd = Path.cwd()
+                modified_files = [file_path.relative_to(cwd) for file_path in tool.tool_records["modified_files"]]
+                return [dict(path=str(file)) for file in modified_files]
+        return []

patchwork/patchflows/LogAnalysis/LogAnalysis.py

@@ -7,7 +7,7 @@
 from patchwork.common.utils.step_typing import validate_steps_with_inputs
 from patchwork.logger import logger
 from patchwork.step import Step
-from patchwork.steps import AgenticLLM
+from patchwork.steps import AgenticLLMV2
 
 _DEFAULT_INPUT_FILE = Path(__file__).parent / "defaults.yml"
 
@@ -16,7 +16,7 @@ class LogAnalysis(Step):
     def __init__(self, inputs: dict):
         PatchflowProgressBar(self).register_steps(
             # CallSQL,
-            AgenticLLM,
+            AgenticLLMV2,
         )
         final_inputs = yaml.safe_load(_DEFAULT_INPUT_FILE.read_text()) or dict()
         final_inputs.update(inputs)
@@ -37,7 +37,7 @@ def run(self) -> dict:
         sentry_filename = "sentry_issues.json"
         for i in range(self.inputs.get("analysis_limit") or 5):
             # for i in range(self.inputs.get("log_finding_limit") or sys.maxsize):
-            logs_detection_output = AgenticLLM(
+            logs_detection_output = AgenticLLMV2(
                 dict(
                     max_agent_calls=5,
                     agent_system_prompt="""\
@@ -67,7 +67,7 @@ def run(self) -> dict:
                 )
             ).run()
 
-            analysis_output = AgenticLLM(
+            analysis_output = AgenticLLMV2(
                 dict(
                     max_agent_calls=5,
                     prompt_value=logs_detection_output,

patchwork/steps/AgenticLLM/AgenticLLM.py

@@ -1,9 +1,7 @@
 from pathlib import Path
 
-from patchwork.common.multiturn_strategy.agentic_strategy import (
-    AgentConfig,
-    AgenticStrategy,
-)
+from patchwork.common.client.llm.aio import AioLlmClient
+from patchwork.common.multiturn_strategy.agentic_strategy import AgenticStrategy
 from patchwork.common.tools import Tool
 from patchwork.step import Step
 from patchwork.steps.AgenticLLM.typed import AgenticLLMInputs, AgenticLLMOutputs
@@ -15,21 +13,18 @@ def __init__(self, inputs):
         base_path = inputs.get("base_path")
         if base_path is None:
             base_path = str(Path.cwd())
-        self.conversation_limit = int(inputs.get("max_agent_calls", 1))
+        self.conversation_limit = int(int(inputs.get("max_llm_calls", 2)) / 2)
         self.agentic_strategy = AgenticStrategy(
-            api_key=inputs.get("anthropic_api_key"),
-            template_data=inputs.get("prompt_value", {}),
-            system_prompt_template=inputs.get("system_prompt", "Summarise from our previous conversation"),
+            llm_client=AioLlmClient.create_aio_client(inputs),
+            tool_set=Tool.get_tools(path=base_path),
+            template_data=inputs.get("prompt_value"),
+            system_prompt_template=inputs.get("system_prompt"),
             user_prompt_template=inputs.get("user_prompt"),
-            agent_configs=[
-                AgentConfig(
-                    name="Assistant",
-                    tool_set=Tool.get_tools(path=base_path),
-                    system_prompt=inputs.get("agent_system_prompt"),
-                )
-            ],
-            example_json=inputs.get("example_json"),
         )
 
     def run(self) -> dict:
-        return self.agentic_strategy.execute(limit=self.conversation_limit)
+        self.agentic_strategy.execute(limit=self.conversation_limit)
+        return dict(
+            conversation_history=self.agentic_strategy.history,
+            tool_records=self.agentic_strategy.tool_records,
+        )

patchwork/steps/AgenticLLM/typed.py

@@ -1,6 +1,7 @@
-from typing_extensions import Annotated, Any, Dict, TypedDict
+from typing_extensions import Annotated, Any, Dict, List, TypedDict
 
 from patchwork.common.utils.step_typing import StepTypeConfig
+from patchwork.steps.CallLLM.CallLLM import TOKEN_URL
 
 
 class AgenticLLMInputs(TypedDict, total=False):
@@ -9,10 +10,33 @@ class AgenticLLMInputs(TypedDict, total=False):
     system_prompt: str
     user_prompt: str
     max_llm_calls: Annotated[int, StepTypeConfig(is_config=True)]
-    anthropic_api_key: str
-    agent_system_prompt: str
-    example_json: str
+    openai_api_key: Annotated[
+        str, StepTypeConfig(is_config=True, or_op=["patched_api_key", "google_api_key", "anthropic_api_key"])
+    ]
+    anthropic_api_key: Annotated[
+        str, StepTypeConfig(is_config=True, or_op=["patched_api_key", "google_api_key", "openai_api_key"])
+    ]
+    patched_api_key: Annotated[
+        str,
+        StepTypeConfig(
+            is_config=True,
+            or_op=["openai_api_key", "google_api_key", "anthropic_api_key"],
+            msg=f"""\
+Model API key not found.
+Please login at: "{TOKEN_URL}"
+Please go to the Integration's tab and generate an API key.
+Please copy the access token that is generated, and add `--patched_api_key=<token>` to the command line.
+
+If you are using a OpenAI API Key, please set `--openai_api_key=<token>`.""",
+        ),
+    ]
+    google_api_key: Annotated[
+        str, StepTypeConfig(is_config=True, or_op=["patched_api_key", "openai_api_key", "anthropic_api_key"])
+    ]
 
 
 class AgenticLLMOutputs(TypedDict):
-    pass
+    conversation_history: List[Dict]
+    tool_records: List[Dict]
+    # request_tokens: int
+    # response_tokens: int