Fix types on AgenticLLMV2 and DatabaseAgent

Author: whoisarpit

patchwork/steps/AgenticLLMV2/AgenticLLMV2.py

@@ -18,15 +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-5-sonnet-latest",
+            model=inputs.get("strategy_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",
+                    model=inputs.get("agent_model", "claude-3-7-sonnet-latest"),
                     tool_set=Tool.get_tools(path=base_path),
                     system_prompt=inputs.get("agent_system_prompt"),
                 )

patchwork/steps/AgenticLLMV2/typed.py

@@ -9,9 +9,13 @@ class AgenticLLMV2Inputs(TypedDict, total=False):
     system_prompt: str
     user_prompt: str
     max_agent_calls: Annotated[int, StepTypeConfig(is_config=True)]
+    google_api_key: str
     anthropic_api_key: str
+    openai_api_key: str
+    strategy_model: str
+    agent_model: str
     agent_system_prompt: str
-    example_json: str
+    example_json: Dict
 
 
 class AgenticLLMV2Outputs(TypedDict):

patchwork/steps/DatabaseAgent/typed.py

@@ -28,7 +28,7 @@ class DatabaseAgentInputs(__DatabaseAgentOutputsRequiredInputs, total=False):
     google_api_key: Annotated[
         str, StepTypeConfig(is_config=True, or_op=["patched_api_key", "openai_api_key", "anthropic_api_key"])
     ]
-    example_json: str
+    example_json: Dict
 
 
 class DatabaseAgentOutputs(TypedDict):