aofire-python-agent 0.1.0__py3-none-any.whl

python_agent/coding_agent.py

@@ -0,0 +1,193 @@
+"""Coding agent: writes Python code to production standards."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import sys
+from collections.abc import AsyncIterator
+from typing import Any
+
+from claude_agent_sdk import (
+    AssistantMessage,
+    ClaudeAgentOptions,
+    ResultMessage,
+    query,
+)
+
+from python_agent.agent_utils import print_text_blocks
+from python_agent.dag_utils import load_dag
+from python_agent.rules import coding_system_prompt
+from python_agent.tool_guard import make_tool_guard
+
+ESCALATION_MODEL = "claude-opus-4-6"
+
+
+async def _prompt_stream(
+    text: str,
+) -> AsyncIterator[dict[str, Any]]:
+    """Wrap a string prompt as an async iterable.
+
+    Required by claude_agent_sdk when can_use_tool is set.
+    """
+    yield {"role": "user", "content": text}
+
+
+async def run_query(
+    task: str, options: Any,
+) -> Any:
+    """Run a single query and return the ResultMessage."""
+    result: Any = None
+    async for message in query(
+        prompt=_prompt_stream(task), options=options,
+    ):
+        if isinstance(message, AssistantMessage):
+            print_text_blocks(message)
+        elif isinstance(message, ResultMessage):
+            cost = message.total_cost_usd or 0.0
+            print(f"\nDone. Cost: ${cost:.4f}")
+            result = message
+    return result
+
+
+def should_escalate(
+    result: Any, max_turns: int | None,
+) -> bool:
+    """Decide whether to escalate to a stronger model."""
+    if result is None:
+        return False
+    if result.is_error:
+        return True
+    if max_turns is not None and result.num_turns >= max_turns:
+        return True
+    return False
+
+
+def remaining_budget(
+    result: Any, max_budget: float | None,
+) -> float | None:
+    """Calculate remaining budget after a query."""
+    if max_budget is None:
+        return None
+    if result is None or result.total_cost_usd is None:
+        return max_budget
+    return float(max_budget - result.total_cost_usd)
+
+
+def _load_ontology_json(
+    dag_file: str | None,
+) -> str | None:
+    """Load ontology JSON from a DAG file, if provided."""
+    if dag_file is None:
+        return None
+    dag = load_dag(dag_file, project_name="")
+    node = dag.get_current_node()
+    if node is None:
+        return None
+    return node.ontology.model_dump_json(indent=2)
+
+
+async def run(
+    task: str, project_dir: str, model: str,
+    max_turns: int | None, max_budget: float | None,
+    dag_file: str | None = None,
+) -> None:
+    """Run the coding agent on a task, escalating to Opus if stuck."""
+    ontology_json = _load_ontology_json(dag_file)
+    prompt = coding_system_prompt(
+        project_dir, ontology_json=ontology_json,
+    )
+    guard = make_tool_guard(project_dir)
+    options = ClaudeAgentOptions(
+        model=model,
+        system_prompt=prompt,
+        allowed_tools=["Read", "Edit", "Bash", "Glob", "Grep"],
+        permission_mode="bypassPermissions",
+        max_turns=max_turns,
+        max_budget_usd=max_budget,
+        cwd=project_dir,
+        can_use_tool=guard,
+    )
+    result = await run_query(task, options)
+    if model == ESCALATION_MODEL:
+        return
+    if not should_escalate(result, max_turns):
+        return
+    print("\nEscalating to Opus...")
+    budget = remaining_budget(result, max_budget)
+    escalation_task = (
+        "Continue this task. Partial changes may already "
+        f"exist in the working directory.\n\n{task}"
+    )
+    escalation_options = ClaudeAgentOptions(
+        model=ESCALATION_MODEL,
+        system_prompt=prompt,
+        allowed_tools=["Read", "Edit", "Bash", "Glob", "Grep"],
+        permission_mode="bypassPermissions",
+        max_turns=max_turns,
+        max_budget_usd=budget,
+        cwd=project_dir,
+        can_use_tool=guard,
+    )
+    await run_query(escalation_task, escalation_options)
+
+
+def parse_args(
+    argv: list[str] | None = None,
+) -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Run the Python coding agent on a task",
+    )
+    parser.add_argument(
+        "task",
+        help="Task description for the agent",
+    )
+    parser.add_argument(
+        "-d", "--project-dir",
+        default=".",
+        help="Project directory to work in (default: current dir)",
+    )
+    parser.add_argument(
+        "-m", "--model",
+        default="claude-sonnet-4-6",
+        help="Model to use (default: claude-sonnet-4-6)",
+    )
+    parser.add_argument(
+        "--max-turns",
+        type=int,
+        default=30,
+        help="Maximum agent turns (default: 30)",
+    )
+    parser.add_argument(
+        "--max-budget",
+        type=float,
+        default=5.0,
+        help="Maximum budget in USD (default: 5.0)",
+    )
+    parser.add_argument(
+        "--dag-file",
+        default=None,
+        help="Path to ontology DAG JSON file for design context",
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Entry point for the aofire-coding-agent CLI."""
+    args = parse_args(argv)
+    asyncio.run(
+        run(
+            args.task,
+            args.project_dir,
+            args.model,
+            args.max_turns,
+            args.max_budget,
+            dag_file=args.dag_file,
+        )
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

python_agent/convergence_agent.py

@@ -0,0 +1,362 @@
+"""Convergence agent: interactive candidate selection."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import sys
+from typing import Any
+
+from pydantic import BaseModel
+
+from claude_agent_sdk import (
+    ClaudeAgentOptions,
+    ClaudeSDKClient,
+)
+
+from python_agent.dag_utils import (
+    load_dag,
+    save_dag,
+    save_snapshot,
+)
+from python_agent.agent_utils import read_user_input
+from python_agent.discovery_agent import (
+    backtrack,
+    format_ontology_summary,
+    print_response,
+    process_response,
+)
+from python_agent.ontology import (
+    DAGNode,
+    Decision,
+    Ontology,
+    OntologyDAG,
+)
+from python_agent.rules import convergence_system_prompt
+
+
+class AgentState(BaseModel):
+    """Mutable state for the convergence agent loop."""
+
+    ontology: Ontology
+    accepted: bool = False
+
+
+def get_children_summaries(
+    dag: OntologyDAG, node_id: str,
+) -> list[tuple[int, DAGNode, str]]:
+    """Build list of (index, node, summary) for children."""
+    children = dag.children_of(node_id)
+    result: list[tuple[int, DAGNode, str]] = []
+    for i, child in enumerate(children, 1):
+        summary = format_ontology_summary(child.ontology)
+        result.append((i, child, summary))
+    return result
+
+
+def format_children_list(
+    children_summaries: list[tuple[int, DAGNode, str]],
+) -> str:
+    """Format children summaries for display."""
+    if not children_summaries:
+        return "No candidates at this node."
+    lines: list[str] = []
+    for idx, child, summary in children_summaries:
+        label = child.label or child.id
+        first_line = summary.split("\n")[0]
+        lines.append(f"  {idx}. {label}: {first_line}")
+    return "\n".join(lines)
+
+
+def navigate_to_node(
+    dag: OntologyDAG, node_id: str,
+) -> Ontology | None:
+    """Set current_node_id and return the node's ontology."""
+    node = dag.get_node(node_id)
+    if node is None:
+        return None
+    dag.current_node_id = node_id
+    return node.ontology.model_copy(deep=True)
+
+
+def is_command(text: str) -> bool:
+    """Check if text is a convergence meta-command."""
+    cmd = text.strip().lower()
+    return (
+        cmd in ("list", "show", "accept", "back")
+        or cmd.startswith("select ")
+        or cmd.startswith("save")
+    )
+
+
+_CmdResult = tuple[str, Ontology | None, bool]
+
+
+def _handle_list_cmd(
+    command: str, ontology: Ontology,
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Handle 'list' command."""
+    children = get_children_summaries(
+        dag, dag.current_node_id,
+    )
+    return (format_children_list(children), None, False)
+
+
+def _handle_show_cmd(
+    command: str, ontology: Ontology,
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Handle 'show' command."""
+    return (format_ontology_summary(ontology), None, False)
+
+
+def _handle_accept_cmd(
+    command: str, ontology: Ontology,
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Handle 'accept' command."""
+    node = dag.get_current_node()
+    label = node.label if node else "unknown"
+    decision = Decision(
+        question="candidate-selection",
+        options=[],
+        chosen=label,
+        rationale="user-accepted",
+    )
+    save_snapshot(
+        dag, ontology,
+        f"accepted-{label}", decision,
+    )
+    save_dag(dag, dag_path)
+    return (f"Accepted: {label}. You can now refine.",
+            None, True)
+
+
+def _handle_back_cmd(
+    command: str, ontology: Ontology,
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Handle 'back' command."""
+    node = backtrack(dag)
+    if node is None:
+        return ("Already at root.", None, False)
+    new_onto = node.ontology.model_copy(deep=True)
+    save_dag(dag, dag_path)
+    label = node.label or node.id
+    msg = f"Backtracked to: {node.id} ({label})"
+    return (msg, new_onto, False)
+
+
+def _handle_select_cmd(
+    command: str, ontology: Ontology,
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Handle 'select <n>' command."""
+    parts = command.strip().split()
+    if len(parts) < 2:
+        return ("Usage: select <number>", None, False)
+    try:
+        index = int(parts[1])
+    except ValueError:
+        return ("Invalid number.", None, False)
+    children = dag.children_of(dag.current_node_id)
+    if index < 1 or index > len(children):
+        msg = f"Range: 1-{len(children)}"
+        return (msg, None, False)
+    child = children[index - 1]
+    new_onto = navigate_to_node(dag, child.id)
+    label = child.label or child.id
+    return (f"Selected: {label}", new_onto, False)
+
+
+def _handle_save_cmd(
+    command: str, ontology: Ontology,
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Handle 'save [label]' command."""
+    label = command.strip()[4:].strip() or "snapshot"
+    save_snapshot(dag, ontology, label)
+    save_dag(dag, dag_path)
+    return (
+        f"Saved snapshot: {dag.current_node_id}",
+        None, False,
+    )
+
+
+_CmdHandler = Any  # Callable signature for command handlers
+
+_DISPATCH: dict[str, _CmdHandler] = {
+    "list": _handle_list_cmd,
+    "show": _handle_show_cmd,
+    "accept": _handle_accept_cmd,
+    "back": _handle_back_cmd,
+}
+
+
+def handle_command(
+    command: str, ontology: Ontology,
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Dispatch a convergence meta-command.
+
+    Returns (message, new_ontology_or_None, is_accept).
+    """
+    cmd = command.strip().lower()
+    handler = _DISPATCH.get(cmd)
+    if handler is not None:
+        result: _CmdResult = handler(
+            command, ontology, dag, dag_path,
+        )
+        return result
+    if cmd.startswith("select "):
+        return _handle_select_cmd(
+            command, ontology, dag, dag_path,
+        )
+    return _handle_save_cmd(
+        command, ontology, dag, dag_path,
+    )
+
+
+def dispatch_command(
+    command: str, state: AgentState,
+    dag: OntologyDAG, dag_path: str,
+) -> None:
+    """Dispatch command and update agent state."""
+    msg, new_onto, is_accept = handle_command(
+        command, state.ontology, dag, dag_path,
+    )
+    print(msg)
+    if new_onto is not None:
+        state.ontology = new_onto
+    if is_accept:
+        state.accepted = True
+
+
+def maybe_process(
+    text: str, state: AgentState,
+) -> None:
+    """Apply ontology updates if accepted."""
+    if state.accepted:
+        process_response(text, state.ontology)
+
+
+def build_query(
+    user_input: str, state: AgentState,
+    dag: OntologyDAG,
+) -> str:
+    """Build LLM query with current context."""
+    from python_agent.rules import frame_data
+    children = get_children_summaries(
+        dag, dag.current_node_id,
+    )
+    ctx = format_children_list(children)
+    node = dag.get_current_node()
+    label = node.label if node else "unknown"
+    context = f"node={label}, children:\n{ctx}"
+    framed = frame_data("context-data", context)
+    return f"[Context: {framed}]\n\n{user_input}"
+
+
+def _init_state(
+    dag: OntologyDAG,
+) -> AgentState | None:
+    """Initialize agent state from DAG."""
+    node = dag.get_current_node()
+    if node is None:
+        return None
+    ontology = node.ontology.model_copy(deep=True)
+    return AgentState(ontology=ontology)
+
+
+def _print_status(
+    state: AgentState, dag: OntologyDAG,
+) -> None:
+    """Print initial status."""
+    print(format_ontology_summary(state.ontology))
+    children = get_children_summaries(
+        dag, dag.current_node_id,
+    )
+    print(format_children_list(children))
+
+
+# taint: ignore[CWE-200] -- Interactive agent displays LLM output to user
+async def _main_loop(
+    client: Any, state: AgentState,
+    dag: OntologyDAG, dag_path: str,
+) -> None:
+    """Run the interactive convergence loop."""
+    while True:
+        user_input = read_user_input()
+        if user_input is None:
+            break
+        if is_command(user_input):
+            dispatch_command(
+                user_input, state, dag, dag_path,
+            )
+            continue
+        q = build_query(user_input, state, dag)
+        await client.query(q)
+        text = await print_response(client)
+        maybe_process(text, state)
+
+
+async def run(dag_path: str, model: str) -> None:
+    """Run the convergence agent interactively."""
+    dag = load_dag(dag_path, "unknown")
+    state = _init_state(dag)
+    if state is None:
+        print("Error: DAG has no current node.")
+        return
+    _print_status(state, dag)
+
+    ontology_json = json.dumps(
+        state.ontology.model_dump(), indent=2,
+    )
+    children = get_children_summaries(
+        dag, dag.current_node_id,
+    )
+    prompt = convergence_system_prompt(
+        ontology_json, format_children_list(children),
+    )
+    options = ClaudeAgentOptions(
+        model=model,
+        system_prompt=prompt,
+        allowed_tools=[],
+        permission_mode="default",
+    )
+
+    async with ClaudeSDKClient(options=options) as client:
+        await _main_loop(client, state, dag, dag_path)
+
+
+def parse_args(
+    argv: list[str] | None = None,
+) -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Interactive candidate convergence",
+    )
+    parser.add_argument(
+        "--dag-file",
+        required=True,
+        help="Path to DAG JSON file",
+    )
+    parser.add_argument(
+        "-m", "--model",
+        default="claude-opus-4-6",
+        help="Model to use (default: claude-opus-4-6)",
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Entry point for the aofire-convergence-agent CLI."""
+    args = parse_args(argv)
+    asyncio.run(run(args.dag_file, args.model))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())