From 5030bd848fc45bdb8605f2f83f24e265c505a027 Mon Sep 17 00:00:00 2001 From: Ed_ Date: Sat, 13 Jun 2026 20:49:37 -0400 Subject: [PATCH] ai client pass (in gemini region) --- src/ai_client.py | 503 +++++++++++++++++++++-------------------------- 1 file changed, 228 insertions(+), 275 deletions(-) diff --git a/src/ai_client.py b/src/ai_client.py index d4230b30..35506709 100644 --- a/src/ai_client.py +++ b/src/ai_client.py @@ -500,29 +500,29 @@ def set_tool_preset(preset_name: Optional[str]) -> None: _tool_approval_modes = {} if not preset_name or preset_name == "None": # Enable all tools if no preset - _agent_tools = {name: True for name in mcp_client.TOOL_NAMES} + _agent_tools = {name: True for name in mcp_client.TOOL_NAMES} _agent_tools[TOOL_NAME] = True - _active_tool_preset = None + _active_tool_preset = None else: try: manager = ToolPresetManager() presets = manager.load_all() if preset_name in presets: preset = presets[preset_name] - _active_tool_preset = preset - new_tools = {name: False for name in mcp_client.TOOL_NAMES} + _active_tool_preset = preset + new_tools = {name: False for name in mcp_client.TOOL_NAMES} new_tools[TOOL_NAME] = False for cat in preset.categories.values(): for tool in cat: - name = tool.name - new_tools[name] = True + name = tool.name + new_tools[name] = True _tool_approval_modes[name] = tool.approval _agent_tools = new_tools except Exception as e: sys.stderr.write(f"[ERROR] Failed to set tool preset '{preset_name}': {e}\n") sys.stderr.flush() _CACHED_ANTHROPIC_TOOLS = None - _CACHED_DEEPSEEK_TOOLS = None + _CACHED_DEEPSEEK_TOOLS = None def set_bias_profile(profile_name: Optional[str]) -> None: """Sets the active tool bias profile for tuning model behavior.""" @@ -531,7 +531,7 @@ def set_bias_profile(profile_name: Optional[str]) -> None: _active_bias_profile = None else: try: - manager = ToolPresetManager() + manager = ToolPresetManager() profiles = manager.load_all_bias_profiles() if profile_name in profiles: _active_bias_profile = profiles[profile_name] @@ -551,8 +551,8 @@ def _build_anthropic_tools() -> list[dict[str, Any]]: for spec in mcp_client.get_tool_schemas(): if _agent_tools.get(spec["name"], True): raw_tools.append({ - "name": spec["name"], - "description": spec["description"], + "name": spec["name"], + "description": spec["description"], "input_schema": spec["parameters"], }) if _agent_tools.get(TOOL_NAME, True): @@ -569,7 +569,7 @@ def _build_anthropic_tools() -> list[dict[str, Any]]: "type": "object", "properties": { "script": { - "type": "string", + "type": "string", "description": "The PowerShell script to execute." } }, @@ -608,9 +608,9 @@ def _gemini_tool_declaration() -> Optional[types.Tool]: for spec in mcp_client.get_tool_schemas(): if _agent_tools.get(spec["name"], True): raw_tools.append({ - "name": spec["name"], + "name": spec["name"], "description": spec["description"], - "parameters": spec["parameters"] + "parameters": spec["parameters"] }) if _agent_tools.get(TOOL_NAME, True): raw_tools.append({ @@ -626,7 +626,7 @@ def _gemini_tool_declaration() -> Optional[types.Tool]: "type": "object", "properties": { "script": { - "type": "string", + "type": "string", "description": "The PowerShell script to execute." } }, @@ -637,22 +637,22 @@ def _gemini_tool_declaration() -> Optional[types.Tool]: _BIAS_ENGINE.apply_semantic_nudges(raw_tools, _active_tool_preset) declarations: list[types.FunctionDeclaration] = [] for tool_def in raw_tools: - props = {} + props = {} params = tool_def.get("parameters", {}) for pname, pdef in params.get("properties", {}).items(): - ptype_str = pdef.get("type", "string").upper() - ptype = getattr(types.Type, ptype_str, types.Type.STRING) + ptype_str = pdef.get("type", "string").upper() + ptype = getattr(types.Type, ptype_str, types.Type.STRING) props[pname] = types.Schema( type=ptype, description=pdef.get("description", ""), ) declarations.append(types.FunctionDeclaration( - name=tool_def["name"], - description=tool_def["description"], - parameters=types.Schema( - type=types.Type.OBJECT, - properties=props, - required=params.get("required", []), + name = tool_def["name"], + description = tool_def["description"], + parameters = types.Schema( + type = types.Type.OBJECT, + properties = props, + required = params.get("required", []), ), )) return types.Tool(function_declarations=declarations) if declarations else None @@ -662,13 +662,13 @@ def _gemini_tool_declaration() -> Optional[types.Tool]: #region: Tool Execution async def _execute_tool_calls_concurrently( - calls: list[Any], - base_dir: str, + calls: list[Any], + base_dir: str, pre_tool_callback: Optional[Callable[[str, str, Optional[Callable[[str], str]]], Optional[str]]], - qa_callback: Optional[Callable[[str], str]], - r_idx: int, - provider: str, - patch_callback: Optional[Callable[[str, str], Optional[str]]] = None + qa_callback: Optional[Callable[[str], str]], + r_idx: int, + provider: str, + patch_callback: Optional[Callable[[str, str], Optional[str]]] = None ) -> list[tuple[str, str, str, str]]: # tool_name, call_id, output, original_name """ Executes tool calls concurrently using asyncio.gather. @@ -702,32 +702,29 @@ async def _execute_tool_calls_concurrently( """ monitor = performance_monitor.get_monitor() if monitor.enabled: monitor.start_component("ai_client._execute_tool_calls_concurrently") - tier = get_current_tier() + tier = get_current_tier() tasks = [] for fc in calls: - if provider == "gemini": - name, args, call_id = fc.name, dict(fc.args), fc.name # Gemini 1.0.0 doesn't have call IDs in types.Part - elif provider == "gemini_cli": - name, args, call_id = cast(str, fc.get("name")), cast(dict[str, Any], fc.get("args", {})), cast(str, fc.get("id")) - elif provider == "anthropic": - name, args, call_id = cast(str, getattr(fc, "name")), cast(dict[str, Any], getattr(fc, "input")), cast(str, getattr(fc, "id")) - elif provider == "deepseek": - tool_info = fc.get("function", {}) - name = cast(str, tool_info.get("name")) + if provider == "gemini": name, args, call_id = fc.name, dict(fc.args), fc.name # Gemini 1.0.0 doesn't have call IDs in types.Part + elif provider == "gemini_cli": name, args, call_id = cast(str, fc.get("name")), cast(dict[str, Any], fc.get("args", {})), cast(str, fc.get("id")) + elif provider == "anthropic": name, args, call_id = cast(str, getattr(fc, "name")), cast(dict[str, Any], getattr(fc, "input")), cast(str, getattr(fc, "id")) + elif provider == "deepseek": + tool_info = fc.get("function", {}) + name = cast(str, tool_info.get("name")) tool_args_str = cast(str, tool_info.get("arguments", "{}")) - call_id = cast(str, fc.get("id")) - try: args = json.loads(tool_args_str) + call_id = cast(str, fc.get("id")) + try: args = json.loads(tool_args_str) except: args = {} elif provider == "minimax": - tool_info = fc.get("function", {}) - name = cast(str, tool_info.get("name")) + tool_info = fc.get("function", {}) + name = cast(str, tool_info.get("name")) tool_args_str = cast(str, tool_info.get("arguments", "{}")) - call_id = cast(str, fc.get("id")) - try: args = json.loads(tool_args_str) + call_id = cast(str, fc.get("id")) + try: args = json.loads(tool_args_str) except: args = {} else: continue - + tasks.append(_execute_single_tool_call_async(name, args, call_id, base_dir, pre_tool_callback, qa_callback, r_idx, tier, patch_callback)) results = await asyncio.gather(*tasks) @@ -735,22 +732,22 @@ async def _execute_tool_calls_concurrently( return results def run_with_tool_loop( - client: Any, + client: Any, request: Union[OpenAICompatibleRequest, Callable[[int], OpenAICompatibleRequest]], *, - capabilities: Optional[VendorCapabilities] = None, - pre_tool_callback: Optional[Callable[[str, str, Optional[Callable[[str], str]]], Optional[str]]] = None, - qa_callback: Optional[Callable[[str], str]] = None, - stream_callback: Optional[Callable[[str], None]] = None, - patch_callback: Optional[Callable[[str, str], Optional[str]]] = None, - base_dir: str, - vendor_name: str, - history_lock: Optional[threading.Lock] = None, - history: Optional[list[dict[str, Any]]] = None, - trim_func: Optional[Callable[[list[dict[str, Any]]], None]] = None, + capabilities: Optional[VendorCapabilities] = None, + pre_tool_callback: Optional[Callable[[str, str, Optional[Callable[[str], str]]], Optional[str]]] = None, + qa_callback: Optional[Callable[[str], str]] = None, + stream_callback: Optional[Callable[[str], None]] = None, + patch_callback: Optional[Callable[[str, str], Optional[str]]] = None, + base_dir: str, + vendor_name: str, + history_lock: Optional[threading.Lock] = None, + history: Optional[list[dict[str, Any]]] = None, + trim_func: Optional[Callable[[list[dict[str, Any]]], None]] = None, reasoning_extractor: Optional[Callable[[Any], str]] = None, - send_func: Optional[Callable[[int], NormalizedResponse]] = None, - on_pre_dispatch: Optional[Callable[[int, list[dict[str, Any]]], list[dict[str, Any]]]] = None, + send_func: Optional[Callable[[int], NormalizedResponse]] = None, + on_pre_dispatch: Optional[Callable[[int, list[dict[str, Any]]], list[dict[str, Any]]]] = None, ) -> str: """ Orchestrates the LLM conversation loop, executing tool calls and updating history. @@ -800,28 +797,23 @@ def run_with_tool_loop( raise RuntimeError(res.errors[0].message if res.errors else "Unknown OpenAI error") return res.data request_builder: Callable[[int], OpenAICompatibleRequest] = (request if callable(request) else (lambda _i: request)) - dispatch_send: Callable[[int], NormalizedResponse] = send_func or _default_send - response_text: str = "" + dispatch_send: Callable[[int], NormalizedResponse] = send_func or _default_send + response_text: str = "" for _round_idx in range(MAX_TOOL_ROUNDS + 2): - response = dispatch_send(_round_idx) + response = dispatch_send(_round_idx) reasoning_content: str = reasoning_extractor(response.raw_response) if reasoning_extractor else "" - response_text = response.text or "" + response_text = response.text or "" if history_lock is not None and history is not None: with history_lock: - msg: dict[str, Any] = {"role": "assistant", "content": response.text or None} - if reasoning_content: - msg["reasoning_content"] = reasoning_content - if response.tool_calls: - msg["tool_calls"] = response.tool_calls + msg: dict[str, Any] = {"role": "assistant", "content": response.text or None} + if reasoning_content: msg["reasoning_content"] = reasoning_content + if response.tool_calls: msg["tool_calls"] = response.tool_calls history.append(msg) - if not response.tool_calls: - break - if on_pre_dispatch is not None: - _adjusted_calls = on_pre_dispatch(_round_idx, response.tool_calls) - else: - _adjusted_calls = response.tool_calls + if not response.tool_calls: break + if on_pre_dispatch is not None: _adjusted_calls = on_pre_dispatch(_round_idx, response.tool_calls) + else: _adjusted_calls = response.tool_calls try: - loop = asyncio.get_running_loop() + loop = asyncio.get_running_loop() results = asyncio.run_coroutine_threadsafe( _execute_tool_calls_concurrently( _adjusted_calls, base_dir, pre_tool_callback, qa_callback, _round_idx, vendor_name, patch_callback, @@ -836,24 +828,23 @@ def run_with_tool_loop( with history_lock: for _i, (tool_name, call_id, out, _err) in enumerate(results): history.append({ - "role": "tool", + "role": "tool", "tool_call_id": call_id, - "content": str(out) if out else "", + "content": str(out) if out else "", }) - if trim_func is not None: - trim_func(history) + if trim_func is not None: trim_func(history) return response_text async def _execute_single_tool_call_async( - name: str, - args: dict[str, Any], - call_id: str, - base_dir: str, + name: str, + args: dict[str, Any], + call_id: str, + base_dir: str, pre_tool_callback: Optional[Callable[[str, str, Optional[Callable[[str], str]]], Optional[str]]], - qa_callback: Optional[Callable[[str], str]], - r_idx: int, - tier: str | None = None, - patch_callback: Optional[Callable[[str, str], Optional[str]]] = None + qa_callback: Optional[Callable[[str], str]], + r_idx: int, + tier: str | None = None, + patch_callback: Optional[Callable[[str, str], Optional[str]]] = None ) -> tuple[str, str, str, str]: """ Executes a single tool call asynchronously, checking the approval clutch. @@ -889,9 +880,9 @@ async def _execute_single_tool_call_async( (like pre_tool_callback and _run_script) to separate worker threads using asyncio.to_thread. """ set_current_tier(tier) - out = "" + out = "" tool_executed = False - events.emit("tool_execution", payload={"status": "started", "tool": name, "args": args, "round": r_idx}) + events.emit("tool_execution", payload = {"status": "started", "tool": name, "args": args, "round": r_idx}) # Check for auto approval mode approval_mode = _tool_approval_modes.get(name, "ask") @@ -906,24 +897,22 @@ async def _execute_single_tool_call_async( elif pre_tool_callback: # pre_tool_callback is synchronous and might block for HITL res = await asyncio.to_thread(pre_tool_callback, scr, base_dir, qa_callback) - if res is None: - out = "USER REJECTED: tool execution cancelled" - else: - out = res + if res is None: out = "USER REJECTED: tool execution cancelled" + else: out = res tool_executed = True if not tool_executed: - is_native = name in mcp_client.TOOL_NAMES - ext_tools = mcp_client.get_external_mcp_manager().get_all_tools() + is_native = name in mcp_client.TOOL_NAMES + ext_tools = mcp_client.get_external_mcp_manager().get_all_tools() is_external = name in ext_tools if name and (is_native or is_external): _append_comms("OUT", "tool_call", {"name": name, "id": call_id, "args": args}) should_approve = (name in mcp_client.MUTATING_TOOLS or is_external) and approval_mode != "auto" and pre_tool_callback if should_approve: label = "MCP MUTATING" if is_native else "EXTERNAL MCP" - desc = f"# {label} TOOL: {name}\n" + "\n".join(f"# {k}: {repr(v)}" for k, v in args.items()) - _res = await asyncio.to_thread(pre_tool_callback, desc, base_dir, qa_callback) - out = "USER REJECTED: tool execution cancelled" if _res is None else await mcp_client.async_dispatch(name, args) + desc = f"# {label} TOOL: {name}\n" + "\n".join(f"# {k}: {repr(v)}" for k, v in args.items()) + _res = await asyncio.to_thread(pre_tool_callback, desc, base_dir, qa_callback) + out = "USER REJECTED: tool execution cancelled" if _res is None else await mcp_client.async_dispatch(name, args) else: out = await mcp_client.async_dispatch(name, args) if tool_log_callback: @@ -936,19 +925,16 @@ async def _execute_single_tool_call_async( out = f"ERROR: unknown tool '{name}'" if tool_log_callback: tool_log_callback(f"ERROR: {name}", out) - + return (name, call_id, out, name) def _run_script(script: str, base_dir: str, qa_callback: Optional[Callable[[str], str]] = None, patch_callback: Optional[Callable[[str, str], Optional[str]]] = None) -> str: if confirm_and_run_callback is None: return "ERROR: no confirmation handler registered" result = confirm_and_run_callback(script, base_dir, qa_callback, patch_callback) - if result is None: - output = "USER REJECTED: command was not executed" - else: - output = result - if tool_log_callback is not None: - tool_log_callback(script, output) + if result is None: output = "USER REJECTED: command was not executed" + else: output = result + if tool_log_callback is not None: tool_log_callback(script, output) return output def _truncate_tool_output(output: str) -> str: @@ -963,30 +949,25 @@ def _truncate_tool_output(output: str) -> str: def _reread_file_items(file_items: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: """ Re-reads file items from the filesystem if their modification times have changed. - Functional Purpose: Iterates through context files, compares current filesystem mtime against cached mtime, and reads file contents if changes are detected, returning both the full refreshed set and the subset of changed items. - Parameters & Inputs: - file_items (list[dict[str, Any]]): List of file dictionaries containing keys "path" and optionally "mtime", "content". + Parameters & Inputs: file_items (list[dict[str, Any]]): List of file dictionaries containing keys "path" and optionally "mtime", "content". - Returns: - tuple[list[dict[str, Any]], list[dict[str, Any]]]: A tuple containing (refreshed_items, changed_items). + Returns: tuple[list[dict[str, Any]], list[dict[str, Any]]]: A tuple containing (refreshed_items, changed_items). - Immediate-Mode DAG / Thread Context: + Immediate-Mode DAG / Thread Context: Called by: _send_gemini Calls: pathlib.Path.stat, pathlib.Path.read_text - SSDL: - `o-> [I:get_mtime] -> [B:changed?] -> [I:read_file] -> [T:diff_text]` + SSDL: `o-> [I:get_mtime] -> [B:changed?] -> [I:read_file] -> [T:diff_text]` - Thread Boundaries: - Runs synchronously in the caller thread. Does synchronous blocking file system I/O. + Thread Boundaries: Runs synchronously in the caller thread. Does synchronous blocking file system I/O. """ refreshed: list[dict[str, Any]] = [] - changed: list[dict[str, Any]] = [] + changed: list[dict[str, Any]] = [] for item in file_items: path = item.get("path") if path is None: @@ -995,11 +976,11 @@ def _reread_file_items(file_items: list[dict[str, Any]]) -> tuple[list[dict[str, p = path if isinstance(path, _P) else _P(path) try: current_mtime = p.stat().st_mtime - prev_mtime = cast(float, item.get("mtime", 0.0)) + prev_mtime = cast(float, item.get("mtime", 0.0)) if current_mtime == prev_mtime: refreshed.append(item) continue - content = p.read_text(encoding="utf-8") + content = p.read_text(encoding="utf-8") new_item = {**item, "old_content": item.get("content", ""), "content": content, "error": False, "mtime": current_mtime} refreshed.append(new_item) changed.append(new_item) @@ -1014,8 +995,8 @@ def _build_file_context_text(file_items: list[dict[str, Any]]) -> str: return "" parts: list[str] = [] for item in file_items: - path = item.get("path") or item.get("entry", "unknown") - suffix = str(path).rsplit(".", 1)[-1] if "." in str(path) else "text" + path = item.get("path") or item.get("entry", "unknown") + suffix = str(path).rsplit(".", 1)[-1] if "." in str(path) else "text" content = item.get("content", "") parts.append(f"### `{path}`\n\n```{suffix}\n{content}\n```") return "\n\n---\n\n".join(parts) @@ -1050,34 +1031,29 @@ def _build_file_diff_text(changed_items: list[dict[str, Any]]) -> str: return "" parts: list[str] = [] for item in changed_items: - path = item.get("path") or item.get("entry", "unknown") - content = cast(str, item.get("content", "")) + path = item.get("path") or item.get("entry", "unknown") + content = cast(str, item.get("content", "")) old_content = cast(str, item.get("old_content", "")) - new_lines = content.splitlines(keepends=True) + new_lines = content.splitlines(keepends=True) if len(new_lines) <= _DIFF_LINE_THRESHOLD or not old_content: suffix = str(path).rsplit(".", 1)[-1] if "." in str(path) else "text" parts.append(f"### `{path}` (full)\n\n```{suffix}\n{content}\n```") else: old_lines = old_content.splitlines(keepends=True) - diff = difflib.unified_diff(old_lines, new_lines, fromfile=str(path), tofile=str(path), lineterm="") + diff = difflib.unified_diff(old_lines, new_lines, fromfile=str(path), tofile=str(path), lineterm="") diff_text = "\n".join(diff) - if diff_text: - parts.append(f"### `{path}` (diff)\n\n```diff\n{diff_text}\n```") - else: - parts.append(f"### `{path}` (no changes detected)") + if diff_text: parts.append(f"### `{path}` (diff)\n\n```diff\n{diff_text}\n```") + else: parts.append(f"### `{path}` (no changes detected)") return "\n\n---\n\n".join(parts) def _build_deepseek_tools() -> list[dict[str, Any]]: - """ - [C: tests/test_tool_access_exclusion.py:test_build_deepseek_tools_excludes_disabled] - """ raw_tools: list[dict[str, Any]] = [] for spec in mcp_client.get_tool_schemas(): if _agent_tools.get(spec["name"], True): raw_tools.append({ - "name": spec["name"], + "name": spec["name"], "description": spec["description"], - "parameters": spec["parameters"] + "parameters": spec["parameters"] }) if _agent_tools.get(TOOL_NAME, True): raw_tools.append({ @@ -1093,7 +1069,7 @@ def _build_deepseek_tools() -> list[dict[str, Any]]: "type": "object", "properties": { "script": { - "type": "string", + "type": "string", "description": "The PowerShell script to execute." } }, @@ -1107,9 +1083,9 @@ def _build_deepseek_tools() -> list[dict[str, Any]]: tools_list.append({ "type": "function", "function": { - "name": tool_def["name"], + "name": tool_def["name"], "description": tool_def["description"], - "parameters": tool_def["parameters"], + "parameters": tool_def["parameters"], } }) return tools_list @@ -1123,35 +1099,29 @@ def _get_deepseek_tools() -> list[dict[str, Any]]: return _CACHED_DEEPSEEK_TOOLS def _content_block_to_dict(block: Any) -> dict[str, Any]: - if isinstance(block, dict): - return block - if hasattr(block, "model_dump"): - return cast(dict[str, Any], block.model_dump()) - if hasattr(block, "to_dict"): - return cast(dict[str, Any], block.to_dict()) + if isinstance(block, dict): return block + if hasattr(block, "model_dump"): return cast(dict[str, Any], block.model_dump()) + if hasattr(block, "to_dict"): return cast(dict[str, Any], block.to_dict()) block_type = getattr(block, "type", None) - if block_type == "text": - return {"type": "text", "text": block.text} - if block_type == "tool_use": - return {"type": "tool_use", "id": getattr(block, "id"), "name": getattr(block, "name"), "input": getattr(block, "input")} + if block_type == "text": return {"type": "text", "text": block.text} + if block_type == "tool_use": return {"type": "tool_use", "id": getattr(block, "id"), "name": getattr(block, "name"), "input": getattr(block, "input")} return {"type": "text", "text": str(block)} #endregion: File Context Building #region: Token Estimation -_CHARS_PER_TOKEN: float = 3.5 -_ANTHROPIC_MAX_PROMPT_TOKENS: int = 180_000 -_GEMINI_MAX_INPUT_TOKENS: int = 900_000 -_FILE_REFRESH_MARKER: str = _project_context_marker if _project_context_marker.strip() else "[SYSTEM: FILES UPDATED]" +_CHARS_PER_TOKEN: float = 3.5 +_ANTHROPIC_MAX_PROMPT_TOKENS: int = 180_000 +_GEMINI_MAX_INPUT_TOKENS: int = 900_000 +_FILE_REFRESH_MARKER: str = _project_context_marker if _project_context_marker.strip() else "[SYSTEM: FILES UPDATED]" def _estimate_message_tokens(msg: dict[str, Any]) -> int: cached = msg.get("_est_tokens") - if cached is not None: - return cast(int, cached) + if cached is not None: return cast(int, cached) total_chars = 0 - content = msg.get("content", "") - if isinstance(content, str): + content = msg.get("content", "") + if isinstance(content, str): total_chars += len(content) elif isinstance(content, list): for block in content: @@ -1174,7 +1144,7 @@ def _invalidate_token_estimate(msg: dict[str, Any]) -> None: def _estimate_prompt_tokens(system_blocks: list[dict[str, Any]], history: list[dict[str, Any]]) -> int: total = 0 for block in system_blocks: - text = cast(str, block.get("text", "")) + text = cast(str, block.get("text", "")) total += max(1, int(len(text) / _CHARS_PER_TOKEN)) total += 2500 for msg in history: @@ -1207,9 +1177,6 @@ def _strip_stale_file_refreshes(history: list[dict[str, Any]]) -> None: _invalidate_token_estimate(msg) def _chunk_text(text: str, chunk_size: int) -> list[str]: - """ - [C: src/rag_engine.py:RAGEngine._chunk_code, src/rag_engine.py:RAGEngine.index_file] - """ return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)] def _build_chunked_context_blocks(md_content: str) -> list[dict[str, Any]]: @@ -1232,10 +1199,9 @@ def _strip_cache_controls(history: list[dict[str, Any]]) -> None: def _add_history_cache_breakpoint(history: list[dict[str, Any]]) -> None: user_indices = [i for i, m in enumerate(history) if m.get("role") == "user"] - if len(user_indices) < 2: - return + if len(user_indices) < 2: return target_idx = user_indices[-2] - content = history[target_idx].get("content") + content = history[target_idx].get("content") if isinstance(content, list) and content: last_block = content[-1] if isinstance(last_block, dict): @@ -1252,11 +1218,10 @@ def _add_history_cache_breakpoint(history: list[dict[str, Any]]) -> None: def _list_anthropic_models() -> list[str]: try: anthropic = _require_warmed("anthropic") - creds = _load_credentials() - client = anthropic.Anthropic(api_key=creds["anthropic"]["api_key"]) + creds = _load_credentials() + client = anthropic.Anthropic(api_key=creds["anthropic"]["api_key"]) models: list[str] = [] - for m in client.models.list(): - models.append(m.id) + for m in client.models.list(): models.append(m.id) return sorted(models) except Exception as exc: raise _classify_anthropic_error(exc) from exc @@ -1267,23 +1232,22 @@ def _ensure_anthropic_client() -> None: if _anthropic_client is None: creds = _load_credentials() _anthropic_client = anthropic.Anthropic( - api_key=creds["anthropic"]["api_key"], - default_headers={"anthropic-beta": "prompt-caching-2024-07-31"} + api_key = creds["anthropic"]["api_key"], + default_headers = {"anthropic-beta": "prompt-caching-2024-07-31"} ) def _trim_anthropic_history(system_blocks: list[dict[str, Any]], history: list[dict[str, Any]]) -> int: _strip_stale_file_refreshes(history) est = _estimate_prompt_tokens(system_blocks, history) - if est <= _ANTHROPIC_MAX_PROMPT_TOKENS: - return 0 + if est <= _ANTHROPIC_MAX_PROMPT_TOKENS: return 0 dropped = 0 while len(history) > 3 and est > _ANTHROPIC_MAX_PROMPT_TOKENS: if history[1].get("role") == "assistant" and len(history) > 2 and history[2].get("role") == "user": removed_asst = history.pop(1) removed_user = history.pop(1) - dropped += 2 - est -= _estimate_message_tokens(removed_asst) - est -= _estimate_message_tokens(removed_user) + dropped += 2 + est -= _estimate_message_tokens(removed_asst) + est -= _estimate_message_tokens(removed_user) while len(history) > 2 and history[1].get("role") == "assistant" and history[2].get("role") == "user": content = history[2].get("content", []) if isinstance(content, list) and content and isinstance(content[0], dict) and content[0].get("type") == "tool_result": @@ -1295,17 +1259,15 @@ def _trim_anthropic_history(system_blocks: list[dict[str, Any]], history: list[d else: break else: - removed = history.pop(1) + removed = history.pop(1) dropped += 1 - est -= _estimate_message_tokens(removed) + est -= _estimate_message_tokens(removed) return dropped def _repair_anthropic_history(history: list[dict[str, Any]]) -> None: - if not history: - return + if not history: return last = history[-1] - if last.get("role") != "assistant": - return + if last.get("role") != "assistant": return content = last.get("content", []) tool_use_ids: list[str] = [] for block in content: @@ -1326,10 +1288,18 @@ def _repair_anthropic_history(history: list[dict[str, Any]]) -> None: ], }) -def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_items: list[dict[str, Any]] | None = None, discussion_history: str = "", pre_tool_callback: Optional[Callable[[str, str, Optional[Callable[[str], str]]], Optional[str]]] = None, qa_callback: Optional[Callable[[str], str]] = None, stream_callback: Optional[Callable[[str], None]] = None, patch_callback: Optional[Callable[[str, str], Optional[str]]] = None) -> Result[str]: +def _send_anthropic( + md_content: str, + user_message: str, + base_dir: str, + file_items: list[dict[str, Any]] | None = None, + discussion_history: str = "", + pre_tool_callback: Optional[Callable[[str, str, Optional[Callable[[str], str]]], Optional[str]]] = None, + qa_callback: Optional[Callable[[str], str]] = None, + stream_callback: Optional[Callable[[str], None]] = None, + patch_callback: Optional[Callable[[str, str], Optional[str]]] = None + ) -> Result[str]: """ - [C: src/ai_server.py:_handle_send] - Functional Purpose: Sends requests to Anthropic models, managing conversation history, prompt caching, token limits, and executing tool loops. Parameters & Inputs: @@ -1344,18 +1314,18 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item Runs on whichever thread calls send (typically an async worker thread). """ anthropic = _require_warmed("anthropic") - genai = _require_warmed("google.genai") - types = genai.types - monitor = performance_monitor.get_monitor() + genai = _require_warmed("google.genai") + types = genai.types + monitor = performance_monitor.get_monitor() if monitor.enabled: monitor.start_component("ai_client._send_anthropic") try: _ensure_anthropic_client() mcp_client.configure(file_items or [], [base_dir]) stable_prompt = _get_combined_system_prompt() stable_blocks: list[dict[str, Any]] = [{"type": "text", "text": stable_prompt, "cache_control": {"type": "ephemeral"}}] - context_text = f"\n\n\n{md_content}\n" + context_text = f"\n\n\n{md_content}\n" context_blocks = _build_chunked_context_blocks(context_text) - system_blocks = stable_blocks + context_blocks + system_blocks = stable_blocks + context_blocks if discussion_history and not _anthropic_history: user_content: list[dict[str, Any]] = [{"type": "text", "text": f"[DISCUSSION HISTORY]\n\n{discussion_history}\n\n---\n\n{user_message}"}] else: @@ -1369,21 +1339,20 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item if _history_trunc_limit > 0 and isinstance(t_content, str) and len(t_content) > _history_trunc_limit: block["content"] = t_content[:_history_trunc_limit] + "\n\n... [TRUNCATED BY SYSTEM TO SAVE TOKENS. Original output was too large.]" modified = True - if modified: - _invalidate_token_estimate(msg) + if modified: _invalidate_token_estimate(msg) _strip_cache_controls(_anthropic_history) _repair_anthropic_history(_anthropic_history) _anthropic_history.append({"role": "user", "content": user_content}) _add_history_cache_breakpoint(_anthropic_history) all_text_parts: list[str] = [] - _cumulative_tool_bytes = 0 - + _cumulative_tool_bytes = 0 + def _strip_private_keys(history: list[dict[str, Any]]) -> list[dict[str, Any]]: return [{k: v for k, v in m.items() if not k.startswith("_")} for m in history] - + for round_idx in range(MAX_TOOL_ROUNDS + 2): response: Any = None - dropped = _trim_anthropic_history(system_blocks, _anthropic_history) + dropped = _trim_anthropic_history(system_blocks, _anthropic_history) if dropped > 0: est_tokens = _estimate_prompt_tokens(system_blocks, _anthropic_history) _append_comms("OUT", "request", { @@ -1392,18 +1361,18 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item f"Estimated {est_tokens} tokens remaining. {len(_anthropic_history)} messages in history.]" ), }) - + events.emit("request_start", payload={"provider": "anthropic", "model": _model, "round": round_idx}) assert _anthropic_client is not None if stream_callback: with _anthropic_client.messages.stream( - model=_model, - max_tokens=_max_tokens, - temperature=_temperature, - top_p=_top_p, - system=cast(Iterable[anthropic.types.TextBlockParam], system_blocks), - tools=cast(Iterable[anthropic.types.ToolParam], _get_anthropic_tools()), - messages=cast(Iterable[anthropic.types.MessageParam], _strip_private_keys(_anthropic_history)), + model = _model, + max_tokens = _max_tokens, + temperature = _temperature, + top_p = _top_p, + system = cast(Iterable[anthropic.types.TextBlockParam], system_blocks), + tools = cast(Iterable[anthropic.types.ToolParam], _get_anthropic_tools()), + messages = cast(Iterable[anthropic.types.MessageParam], _strip_private_keys(_anthropic_history)), ) as stream: for event in stream: if isinstance(event, anthropic.types.ContentBlockDeltaEvent) and event.delta.type == "text_delta": @@ -1411,17 +1380,17 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item response = stream.get_final_message() else: response = _anthropic_client.messages.create( - model=_model, - max_tokens=_max_tokens, - temperature=_temperature, - top_p=_top_p, - system=cast(Iterable[anthropic.types.TextBlockParam], system_blocks), - tools=cast(Iterable[anthropic.types.ToolParam], _get_anthropic_tools()), - messages=cast(Iterable[anthropic.types.MessageParam], _strip_private_keys(_anthropic_history)), + model = _model, + max_tokens = _max_tokens, + temperature = _temperature, + top_p = _top_p, + system = cast(Iterable[anthropic.types.TextBlockParam], system_blocks), + tools = cast(Iterable[anthropic.types.ToolParam], _get_anthropic_tools()), + messages = cast(Iterable[anthropic.types.MessageParam], _strip_private_keys(_anthropic_history)), ) serialised_content = [_content_block_to_dict(b) for b in response.content] _anthropic_history.append({ - "role": "assistant", + "role": "assistant", "content": serialised_content, }) text_blocks = [b.text for b in response.content if hasattr(b, "text") and b.text] @@ -1436,12 +1405,10 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item if response.usage: usage_dict["input_tokens"] = response.usage.input_tokens usage_dict["output_tokens"] = response.usage.output_tokens - cache_creation = getattr(response.usage, "cache_creation_input_tokens", None) - cache_read = getattr(response.usage, "cache_read_input_tokens", None) - if cache_creation is not None: - usage_dict["cache_creation_input_tokens"] = cache_creation - if cache_read is not None: - usage_dict["cache_read_input_tokens"] = cache_read + cache_creation = getattr(response.usage, "cache_creation_input_tokens", None) + cache_read = getattr(response.usage, "cache_read_input_tokens", None) + if cache_creation is not None: usage_dict["cache_creation_input_tokens"] = cache_creation + if cache_read is not None: usage_dict["cache_read_input_tokens"] = cache_read events.emit("response_received", payload={"provider": "anthropic", "model": _model, "usage": usage_dict, "round": round_idx}) _append_comms("IN", "response", { "round": round_idx, @@ -1450,21 +1417,19 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item "tool_calls": tool_use_blocks, "usage": usage_dict, }) - if response.stop_reason != "tool_use" or not tool_use_blocks: - break - if round_idx > MAX_TOOL_ROUNDS: - break + if response.stop_reason != "tool_use" or not tool_use_blocks: break + if round_idx > MAX_TOOL_ROUNDS: break # Execute tools concurrently try: - loop = asyncio.get_running_loop() + loop = asyncio.get_running_loop() results = asyncio.run_coroutine_threadsafe( _execute_tool_calls_concurrently(response.content, base_dir, pre_tool_callback, qa_callback, round_idx, "anthropic", patch_callback), loop ).result() except RuntimeError: results = asyncio.run(_execute_tool_calls_concurrently(response.content, base_dir, pre_tool_callback, qa_callback, round_idx, "anthropic", patch_callback)) - + tool_results: list[dict[str, Any]] = [] for i, (name, call_id, out, _) in enumerate(results): truncated = _truncate_tool_output(out) @@ -1476,7 +1441,7 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item }) _append_comms("IN", "tool_result", {"name": name, "id": call_id, "output": out}) events.emit("tool_execution", payload={"status": "completed", "tool": name, "result": out, "round": round_idx}) - + if _cumulative_tool_bytes > _MAX_TOOL_OUTPUT_BYTES: tool_results.append({ "type": "text", @@ -1485,7 +1450,7 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item _append_comms("OUT", "request", {"message": f"[TOOL OUTPUT BUDGET EXCEEDED: {_cumulative_tool_bytes} bytes]"}) if file_items: file_items, changed = _reread_file_items(file_items) - refreshed_ctx = _build_file_diff_text(changed) + refreshed_ctx = _build_file_diff_text(changed) if refreshed_ctx: tool_results.append({ "type": "text", @@ -1510,7 +1475,7 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item ], }) final_text = "\n\n".join(all_text_parts) - res = final_text if final_text.strip() else "(No text returned by the model)" + res = final_text if final_text.strip() else "(No text returned by the model)" if monitor.enabled: monitor.end_component("ai_client._send_anthropic") return Result(data=res) except Exception as exc: @@ -1522,19 +1487,15 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str, file_item #region: Gemini Provider def get_gemini_cache_stats() -> dict[str, Any]: - """ - [C: src/app_controller.py:AppController._recalculate_session_usage, src/app_controller.py:AppController._update_cached_stats, tests/test_ai_cache_tracking.py:test_gemini_cache_tracking, tests/test_gemini_metrics.py:test_get_gemini_cache_stats_with_mock_client] - """ _ensure_gemini_client() - if not _gemini_client: - return {"cache_count": 0, "total_size_bytes": 0, "cached_files": []} - caches_iterator = _gemini_client.caches.list() - caches = list(caches_iterator) + if not _gemini_client: return {"cache_count": 0, "total_size_bytes": 0, "cached_files": []} + caches_iterator = _gemini_client.caches.list() + caches = list(caches_iterator) total_size_bytes = sum(getattr(c, 'size_bytes', 0) for c in caches) return { - "cache_count": len(caches), + "cache_count": len(caches), "total_size_bytes": total_size_bytes, - "cached_files": _gemini_cached_file_paths, + "cached_files": _gemini_cached_file_paths, } def _list_gemini_cli_models() -> list[str]: @@ -1549,98 +1510,90 @@ def _list_gemini_cli_models() -> list[str]: def _list_gemini_models(api_key: str) -> list[str]: try: - genai = _require_warmed("google.genai") + genai = _require_warmed("google.genai") client = genai.Client(api_key=api_key) models: list[str] = [] for m in client.models.list(): name = m.name - if name and name.startswith("models/"): - name = name[len("models/"):] - if name and "gemini" in name.lower(): - models.append(name) + if name and name.startswith("models/"): name = name[len("models/"):] + if name and "gemini" in name.lower(): models.append(name) return sorted(models) except Exception as exc: raise _classify_gemini_error(exc) from exc def _ensure_gemini_client() -> None: - """ - [C: src/rag_engine.py:GeminiEmbeddingProvider.embed] - """ global _gemini_client genai = _require_warmed("google.genai") if _gemini_client is None: - creds = _load_credentials() + creds = _load_credentials() _gemini_client = genai.Client(api_key=creds["gemini"]["api_key"]) def _get_gemini_history_list(chat: Any | None) -> list[Any]: if not chat: return [] - if hasattr(chat, "_history"): - return cast(list[Any], chat._history) - if hasattr(chat, "history"): - return cast(list[Any], chat.history) - if hasattr(chat, "get_history"): - return cast(list[Any], chat.get_history()) + if hasattr(chat, "_history"): return cast(list[Any], chat._history) + if hasattr(chat, "history"): return cast(list[Any], chat.history) + if hasattr(chat, "get_history"): return cast(list[Any], chat.get_history()) return [] def _send_gemini(md_content: str, user_message: str, base_dir: str, - file_items: list[dict[str, Any]] | None = None, - discussion_history: str = "", - pre_tool_callback: Optional[Callable[[str, str, Optional[Callable[[str], str]]], Optional[str]]] = None, - qa_callback: Optional[Callable[[str], str]] = None, - enable_tools: bool = True, - stream_callback: Optional[Callable[[str], None]] = None, - patch_callback: Optional[Callable[[str, str], Optional[str]]] = None) -> Result[str]: + file_items: list[dict[str, Any]] | None = None, + discussion_history: str = "", + pre_tool_callback: Optional[Callable[[str, str, Optional[Callable[[str], str]]], Optional[str]]] = None, + qa_callback: Optional[Callable[[str], str]] = None, + enable_tools: bool = True, + stream_callback: Optional[Callable[[str], None]] = None, + patch_callback: Optional[Callable[[str, str], Optional[str]]] = None + ) -> Result[str]: """ - [C: src/ai_server.py:_handle_send, tests/test_tier4_interceptor.py:test_gemini_provider_passes_qa_callback_to_run_script] Functional Purpose: Sends requests to Gemini via google-genai SDK, handling context caching, chat history, and tools. Parameters & Inputs: md_content, user_message, base_dir, file_items, discussion_history, callbacks, enable_tools. Immediate-Mode DAG / Thread Context: Called by: send; Calls: _ensure_gemini_client, client.caches.create, client.chats.create, run_with_tool_loop - SSDL: - [I:_ensure_gemini_client] -> [B:Cache Changed?] -> [I:client.caches.create] -> [I:client.chats.create] -> [T:Result] + SSDL: [I:_ensure_gemini_client] -> [B:Cache Changed?] -> [I:client.caches.create] -> [I:client.chats.create] -> [T:Result] Thread Boundaries: Runs on caller thread (typically an async worker thread). """ global _gemini_chat, _gemini_cache, _gemini_cache_md_hash, _gemini_cache_created_at, _gemini_cached_file_paths - genai = _require_warmed("google.genai") - types = genai.types + genai = _require_warmed("google.genai") + types = genai.types monitor = performance_monitor.get_monitor() if monitor.enabled: monitor.start_component("ai_client._send_gemini") try: _ensure_gemini_client(); mcp_client.configure(file_items or [], [base_dir]) - sys_instr = f"{_get_combined_system_prompt()}\n\n\n{md_content}\n" - td = _gemini_tool_declaration() if enable_tools else None - tools_decl = [td] if td else None + sys_instr = f"{_get_combined_system_prompt()}\n\n\n{md_content}\n" + td = _gemini_tool_declaration() if enable_tools else None + tools_decl = [td] if td else None current_md_hash = hashlib.md5(md_content.encode()).hexdigest() - old_history = None + old_history = None assert _gemini_client is not None if _gemini_chat and _gemini_cache_md_hash != current_md_hash: old_history = list(_get_gemini_history_list(_gemini_chat)) if _get_gemini_history_list(_gemini_chat) else [] if _gemini_cache: try: _gemini_client.caches.delete(name=_gemini_cache.name) except Exception as e: _append_comms("OUT", "request", {"message": f"[CACHE DELETE WARN] {e}"}) - _gemini_chat = None - _gemini_cache = None - _gemini_cache_created_at = None + _gemini_chat = None + _gemini_cache = None + _gemini_cache_created_at = None _gemini_cached_file_paths = [] _append_comms("OUT", "request", {"message": "[CONTEXT CHANGED] Rebuilding cache and chat session..."}) if _gemini_chat and _gemini_cache and _gemini_cache_created_at: elapsed = time.time() - _gemini_cache_created_at if elapsed > _GEMINI_CACHE_TTL * 0.9: old_history = list(_get_gemini_history_list(_gemini_chat)) if _get_gemini_history_list(_gemini_chat) else [] + #TODO(Ed): Review(Exception) try: _gemini_client.caches.delete(name=_gemini_cache.name) except Exception as e: _append_comms("OUT", "request", {"message": f"[CACHE DELETE WARN] {e}"}) - _gemini_chat = None - _gemini_cache = None - _gemini_cache_created_at = None + _gemini_chat = None + _gemini_cache = None + _gemini_cache_created_at = None _gemini_cached_file_paths = [] _append_comms("OUT", "request", {"message": f"[CACHE TTL] Rebuilding cache (expired after {int(elapsed)}s)..."}) if not _gemini_chat: chat_config = types.GenerateContentConfig( - system_instruction=sys_instr, - tools=cast(Any, tools_decl), - temperature=_temperature, - top_p=_top_p, - max_output_tokens=_max_tokens, - safety_settings=[types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH)] + system_instruction = sys_instr, + tools = cast(Any, tools_decl), + temperature = _temperature, + top_p = _top_p, + max_output_tokens = _max_tokens, + safety_settings = [types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH)] ) should_cache = False try: