progress
This commit is contained in:
128
ai_client.py
128
ai_client.py
@@ -4,6 +4,7 @@ import json
|
||||
import datetime
|
||||
from pathlib import Path
|
||||
import file_cache
|
||||
import mcp_client
|
||||
|
||||
_provider: str = "gemini"
|
||||
_model: str = "gemini-2.0-flash"
|
||||
@@ -33,7 +34,7 @@ MAX_TOOL_ROUNDS = 5
|
||||
_ANTHROPIC_CHUNK_SIZE = 180_000
|
||||
|
||||
_ANTHROPIC_SYSTEM = (
|
||||
"You are a helpful coding assistant with access to a PowerShell tool. "
|
||||
"You are a helpful coding assistant with access to a PowerShell tool and MCP file tools (read_file, list_directory, search_files, get_file_summary). "
|
||||
"When asked to create or edit files, prefer targeted edits over full rewrites. "
|
||||
"Always explain what you are doing before invoking the tool.\n\n"
|
||||
"When writing or rewriting large files (especially those containing quotes, backticks, or special characters), "
|
||||
@@ -221,8 +222,16 @@ def _list_anthropic_models() -> list[str]:
|
||||
|
||||
TOOL_NAME = "run_powershell"
|
||||
|
||||
_ANTHROPIC_TOOLS = [
|
||||
{
|
||||
def _build_anthropic_tools() -> list[dict]:
|
||||
"""Build the full Anthropic tools list: run_powershell + MCP file tools."""
|
||||
mcp_tools = []
|
||||
for spec in mcp_client.MCP_TOOL_SPECS:
|
||||
mcp_tools.append({
|
||||
"name": spec["name"],
|
||||
"description": spec["description"],
|
||||
"input_schema": spec["parameters"],
|
||||
})
|
||||
powershell_tool = {
|
||||
"name": TOOL_NAME,
|
||||
"description": (
|
||||
"Run a PowerShell script within the project base_dir. "
|
||||
@@ -243,34 +252,57 @@ _ANTHROPIC_TOOLS = [
|
||||
},
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
]
|
||||
return mcp_tools + [powershell_tool]
|
||||
|
||||
|
||||
_ANTHROPIC_TOOLS = _build_anthropic_tools()
|
||||
|
||||
|
||||
def _gemini_tool_declaration():
|
||||
from google.genai import types
|
||||
return types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name=TOOL_NAME,
|
||||
description=(
|
||||
"Run a PowerShell script within the project base_dir. "
|
||||
"Use this to create, edit, rename, or delete files and directories. "
|
||||
"The working directory is set to base_dir automatically. "
|
||||
"stdout and stderr are returned to you as the result."
|
||||
),
|
||||
parameters=types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
properties={
|
||||
"script": types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description="The PowerShell script to execute."
|
||||
)
|
||||
},
|
||||
required=["script"]
|
||||
)
|
||||
|
||||
declarations = []
|
||||
|
||||
# MCP file tools
|
||||
for spec in mcp_client.MCP_TOOL_SPECS:
|
||||
props = {}
|
||||
for pname, pdef in spec["parameters"].get("properties", {}).items():
|
||||
props[pname] = types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description=pdef.get("description", ""),
|
||||
)
|
||||
]
|
||||
)
|
||||
declarations.append(types.FunctionDeclaration(
|
||||
name=spec["name"],
|
||||
description=spec["description"],
|
||||
parameters=types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
properties=props,
|
||||
required=spec["parameters"].get("required", []),
|
||||
),
|
||||
))
|
||||
|
||||
# PowerShell tool
|
||||
declarations.append(types.FunctionDeclaration(
|
||||
name=TOOL_NAME,
|
||||
description=(
|
||||
"Run a PowerShell script within the project base_dir. "
|
||||
"Use this to create, edit, rename, or delete files and directories. "
|
||||
"The working directory is set to base_dir automatically. "
|
||||
"stdout and stderr are returned to you as the result."
|
||||
),
|
||||
parameters=types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
properties={
|
||||
"script": types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description="The PowerShell script to execute."
|
||||
)
|
||||
},
|
||||
required=["script"]
|
||||
),
|
||||
))
|
||||
|
||||
return types.Tool(function_declarations=declarations)
|
||||
|
||||
|
||||
def _run_script(script: str, base_dir: str) -> str:
|
||||
@@ -365,6 +397,7 @@ def _send_gemini(md_content: str, user_message: str, base_dir: str, file_items:
|
||||
|
||||
try:
|
||||
_ensure_gemini_client()
|
||||
mcp_client.configure(file_items or [], [base_dir])
|
||||
|
||||
if _gemini_chat is None:
|
||||
_gemini_chat = _gemini_client.chats.create(
|
||||
@@ -407,8 +440,20 @@ def _send_gemini(md_content: str, user_message: str, base_dir: str, file_items:
|
||||
|
||||
function_responses = []
|
||||
for fc in tool_calls:
|
||||
if fc.name == TOOL_NAME:
|
||||
script = fc.args.get("script", "")
|
||||
fc_name = fc.name
|
||||
fc_args = dict(fc.args)
|
||||
if fc_name in mcp_client.TOOL_NAMES:
|
||||
_append_comms("OUT", "tool_call", {"name": fc_name, "args": fc_args})
|
||||
output = mcp_client.dispatch(fc_name, fc_args)
|
||||
_append_comms("IN", "tool_result", {"name": fc_name, "output": output})
|
||||
function_responses.append(
|
||||
types.Part.from_function_response(
|
||||
name=fc_name,
|
||||
response={"output": output}
|
||||
)
|
||||
)
|
||||
elif fc_name == TOOL_NAME:
|
||||
script = fc_args.get("script", "")
|
||||
_append_comms("OUT", "tool_call", {
|
||||
"name": TOOL_NAME,
|
||||
"script": script,
|
||||
@@ -527,6 +572,7 @@ def _repair_anthropic_history(history: list[dict]):
|
||||
def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_items: list[dict] | None = None) -> str:
|
||||
try:
|
||||
_ensure_anthropic_client()
|
||||
mcp_client.configure(file_items or [], [base_dir])
|
||||
|
||||
context_blocks = _build_chunked_context_blocks(md_content)
|
||||
|
||||
@@ -557,7 +603,7 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
tools=_ANTHROPIC_TOOLS,
|
||||
tools=_build_anthropic_tools(),
|
||||
messages=_anthropic_history,
|
||||
)
|
||||
|
||||
@@ -600,22 +646,36 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item
|
||||
|
||||
tool_results = []
|
||||
for block in response.content:
|
||||
if getattr(block, "type", None) == "tool_use" and getattr(block, "name", None) == TOOL_NAME:
|
||||
script = block.input.get("script", "")
|
||||
if getattr(block, "type", None) != "tool_use":
|
||||
continue
|
||||
b_name = getattr(block, "name", None)
|
||||
b_id = getattr(block, "id", "")
|
||||
b_input = getattr(block, "input", {})
|
||||
if b_name in mcp_client.TOOL_NAMES:
|
||||
_append_comms("OUT", "tool_call", {"name": b_name, "id": b_id, "args": b_input})
|
||||
output = mcp_client.dispatch(b_name, b_input)
|
||||
_append_comms("IN", "tool_result", {"name": b_name, "id": b_id, "output": output})
|
||||
tool_results.append({
|
||||
"type": "tool_result",
|
||||
"tool_use_id": b_id,
|
||||
"content": output,
|
||||
})
|
||||
elif b_name == TOOL_NAME:
|
||||
script = b_input.get("script", "")
|
||||
_append_comms("OUT", "tool_call", {
|
||||
"name": TOOL_NAME,
|
||||
"id": block.id,
|
||||
"id": b_id,
|
||||
"script": script,
|
||||
})
|
||||
output = _run_script(script, base_dir)
|
||||
_append_comms("IN", "tool_result", {
|
||||
"name": TOOL_NAME,
|
||||
"id": block.id,
|
||||
"id": b_id,
|
||||
"output": output,
|
||||
})
|
||||
tool_results.append({
|
||||
"type": "tool_result",
|
||||
"tool_use_id": block.id,
|
||||
"tool_use_id": b_id,
|
||||
"content": output,
|
||||
})
|
||||
|
||||
|
||||
Reference in New Issue
Block a user