update agentics model usage; add location to google api

Author: CTY-git

patchwork/common/client/llm/aio.py

@@ -203,20 +203,21 @@ def chat_completion(
     def create_aio_client(inputs) -> "AioLlmClient" | None:
         clients = []
 
+        client_args = {key[len("client_") :]: value for key, value in inputs.items() if key.startswith("client_")}
+
         patched_key = inputs.get("patched_api_key")
         if patched_key is not None:
             client = OpenAiLlmClient(patched_key, DEFAULT_PATCH_URL)
             clients.append(client)
 
         openai_key = inputs.get("openai_api_key") or os.environ.get("OPENAI_API_KEY")
         if openai_key is not None:
-            client_args = {key[len("client_") :]: value for key, value in inputs.items() if key.startswith("client_")}
             client = OpenAiLlmClient(openai_key, **client_args)
             clients.append(client)
 
         google_key = inputs.get("google_api_key")
         if google_key is not None:
-            client = GoogleLlmClient(google_key)
+            client = GoogleLlmClient(google_key, **client_args)
             clients.append(client)
 
         anthropic_key = inputs.get("anthropic_api_key")

patchwork/common/client/llm/google.py

@@ -54,9 +54,10 @@ class GoogleLlmClient(LlmClient):
     ]
     __MODEL_PREFIX = "models/"
 
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str, location: Optional[str] = None):
         self.__api_key = api_key
-        self.client = genai.Client(api_key=api_key)
+        self.__location = location
+        self.client = genai.Client(api_key=api_key, location=location)
 
     @lru_cache(maxsize=1)
     def __get_models_info(self) -> list[Model]:
@@ -69,7 +70,11 @@ def __get_pydantic_model(self, model_settings: ModelSettings | None) -> Model:
         if model_name is None:
             raise ValueError("Model must be set cannot be None")
 
-        return GeminiModel(model_name, api_key=self.__api_key)
+        if self.__location is None:
+            return GeminiModel(model_name, api_key=self.__api_key)
+
+        url_template = f"https://{self.__location}-generativelanguage.googleapis.com/v1beta/models/{{model}}:"
+        return GeminiModel(model_name, api_key=self.__api_key, url_template=url_template)
 
     async def request(
         self,

patchwork/common/multiturn_strategy/agentic_strategy_v2.py

@@ -25,6 +25,7 @@ class Config:
         arbitrary_types_allowed = True
 
     name: str
+    model: str
     tool_set: Dict[str, Tool]
     system_prompt: str = ""
     example_json: Union[str, Dict[str, Any]] = DEFAULT_AGENT_EXAMPLE_JSON
@@ -41,15 +42,15 @@ def model_post_init(self, __context: Any) -> None:
 
 class AgenticStrategyV2:
     def __init__(
-        self,
-        model: str,
-        llm_client: LlmClient,
-        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,
+            model: str,
+            llm_client: LlmClient,
+            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
@@ -76,7 +77,7 @@ def __init__(
                 result_type=example_json_to_base_model(agent_config.example_json),
                 model_settings=dict(
                     parallel_tool_calls=False,
-                    model=model,
+                    model=agent_config.model,
                 ),
             )
 
@@ -151,7 +152,7 @@ def execute(self, limit: Optional[int] = None) -> dict:
                 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

patchwork/common/tools/grep_tool.py

@@ -114,7 +114,7 @@ def json_schema(self) -> dict:
         return {
             "name": "find_text",
             "description": f"""\
-Tool to find text in a file using a pattern based on the Unix shell style.
+Tool to find text in a file or files in a directory using a pattern based on the Unix shell style.
 The current working directory is always {self.__working_dir}. 
 The path provided should either be absolute or relative to the current working directory.
 
@@ -124,10 +124,6 @@ def json_schema(self) -> dict:
             "input_schema": {
                 "type": "object",
                 "properties": {
-                    "path": {
-                        "description": "The path to the file to find text in.",
-                        "type": "string",
-                    },
                     "pattern": {
                         "description": """\
 The Unix shell style pattern to match files using.
@@ -141,6 +137,14 @@ def json_schema(self) -> dict:
 Example:
 * '*macs' will match the file '.emacs'
 * '*.py' will match all files with the '.py' extension
+""",
+                        "type": "string",
+                    },
+                    "path": {
+                        "description": """\
+The path to the file to find text in. 
+If not given, will search all file content in the current working directory.
+If the path is a directory, will search all file content in the directory.
 """,
                         "type": "string",
                     },
@@ -149,22 +153,22 @@ def json_schema(self) -> dict:
                         "type": "boolean",
                     },
                 },
-                "required": ["path", "pattern"],
+                "required": ["pattern"],
             },
         }
 
     def execute(
-        self,
-        path: Optional[Path] = None,
-        pattern: Optional[str] = None,
-        is_case_sensitive: bool = False,
+            self,
+            pattern: Optional[str] = None,
+            path: Optional[Path] = None,
+            is_case_sensitive: bool = False,
     ) -> str:
-        if path is None:
-            raise ValueError("Path argument is required!")
-
         if pattern is None:
             raise ValueError("pattern argument is required!")
 
+        if path is None:
+            path = Path(self.__working_dir)
+
         matcher = fnmatch.fnmatch
         if is_case_sensitive:
             matcher = fnmatch.fnmatchcase
@@ -173,15 +177,34 @@ def execute(
         if not path.is_relative_to(self.__working_dir):
             raise ValueError("Path must be relative to working dir")
 
-        matches = []
-        with path.open("r") as f:
-            for i, line in enumerate(f.readlines()):
-                if not matcher(line, pattern):
-                    continue
+        if path.is_file():
+            paths = [path]
+        else:
+            paths = [p for p in path.iterdir() if p.is_file()]
+
+        from collections import defaultdict
+        file_matches = defaultdict(list)
+        for path in paths:
+            with path.open("r") as f:
+                for i, line in enumerate(f.readlines()):
+                    if not matcher(line, pattern):
+                        continue
+
+                    content = f"Line {i + 1}: {line}"
+                    if len(line) > self.__CHAR_LIMIT:
+                        content = f"Line {i + 1}: {self.__CHAR_LIMIT_TEXT}"
+
+                    file_matches[str(path)].append(content)
+
+        total_file_matches = ""
+        for path_str, matches in file_matches.items():
+            total_file_matches += f"\nPattern matches found in '{path}':\n" + "\n".join(matches)
+
+        if len(total_file_matches) <= 5000:
+            return total_file_matches
 
-                content = f"Line {i + 1}: {line}"
-                if len(line) > self.__CHAR_LIMIT:
-                    content = f"Line {i + 1}: {self.__CHAR_LIMIT_TEXT}"
+        total_file_matches = ""
+        for path_str, matches in file_matches.items():
+            total_file_matches += f"\n {len(matches)} Pattern matches found in '{path}': <TRUNCATED>\n"
+        return total_file_matches
 
-                matches.append(content)
-        return f"Pattern matches found in '{path}':\n" + "\n".join(matches)

patchwork/steps/AgenticLLMV2/AgenticLLMV2.py

@@ -18,14 +18,15 @@ def __init__(self, inputs):
             base_path = str(Path.cwd())
         self.conversation_limit = int(inputs.get("max_agent_calls", 1))
         self.agentic_strategy = AgenticStrategyV2(
-            model="claude-3-7-sonnet-latest",
+            model="claude-3-5-sonnet-latest",
             llm_client=AioLlmClient.create_aio_client(inputs),
             template_data=inputs.get("prompt_value", {}),
             system_prompt_template=inputs.get("system_prompt", "Summarise from our previous conversation"),
             user_prompt_template=inputs.get("user_prompt"),
             agent_configs=[
                 AgentConfig(
                     name="Assistant",
+                    model="claude-3-7-sonnet-latest",
                     tool_set=Tool.get_tools(path=base_path),
                     system_prompt=inputs.get("agent_system_prompt"),
                 )