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

python_agent/discovery_agent.py

@@ -0,0 +1,348 @@
+"""Discovery agent: interactive ontology building."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import sys
+from typing import Any
+
+from claude_agent_sdk import (
+    AssistantMessage,
+    ClaudeAgentOptions,
+    ClaudeSDKClient,
+)
+
+from python_agent.agent_utils import (
+    collect_response_text,
+    extract_ontology_json,
+    print_text_blocks,
+    read_user_input,
+)
+from python_agent.dag_utils import (
+    load_dag,
+    save_dag,
+    save_snapshot,
+)
+from python_agent.ontology import (
+    DAGNode,
+    DomainConstraint,
+    Entity,
+    Ontology,
+    OntologyDAG,
+    OpenQuestion,
+    Relationship,
+)
+from python_agent.rules import (
+    discovery_system_prompt,
+    frame_data,
+)
+
+
+def _upsert_entities(
+    ontology: Ontology, new_entities: list[Any],
+) -> None:
+    """Replace or append entities by id."""
+    existing = {
+        e.id: i for i, e in enumerate(ontology.entities)
+    }
+    for item in new_entities:
+        entity = Entity.model_validate(item)
+        idx = existing.get(entity.id)
+        if idx is not None:
+            ontology.entities[idx] = entity
+        else:
+            ontology.entities.append(entity)
+            existing[entity.id] = len(
+                ontology.entities,
+            ) - 1
+
+
+def _append_relationships(
+    ontology: Ontology, items: list[Any],
+) -> None:
+    """Append new relationships to the ontology."""
+    for item in items:
+        ontology.relationships.append(
+            Relationship.model_validate(item),
+        )
+
+
+def _append_constraints(
+    ontology: Ontology, items: list[Any],
+) -> None:
+    """Append new domain constraints to the ontology."""
+    for item in items:
+        ontology.domain_constraints.append(
+            DomainConstraint.model_validate(item),
+        )
+
+
+def _upsert_open_questions(
+    ontology: Ontology, items: list[Any],
+) -> None:
+    """Replace or append open questions by id."""
+    existing = {
+        q.id: i
+        for i, q in enumerate(ontology.open_questions)
+    }
+    for item in items:
+        question = OpenQuestion.model_validate(item)
+        idx = existing.get(question.id)
+        if idx is not None:
+            ontology.open_questions[idx] = question
+        else:
+            ontology.open_questions.append(question)
+            existing[question.id] = len(
+                ontology.open_questions,
+            ) - 1
+
+
+_MERGE_DISPATCH = {
+    "entities": _upsert_entities,
+    "relationships": _append_relationships,
+    "domain_constraints": _append_constraints,
+    "open_questions": _upsert_open_questions,
+}
+
+
+def merge_ontology_update(
+    ontology: Ontology, update_dict: dict[str, Any],
+) -> bool:
+    """Apply a partial ontology update to an Ontology.
+
+    Returns True if any updates were applied.
+    """
+    applied = False
+    for key, handler in _MERGE_DISPATCH.items():
+        items = update_dict.get(key)
+        if items:
+            handler(ontology, items)
+            applied = True
+    return applied
+
+
+def process_response(
+    response_text: str, ontology: Ontology,
+) -> bool:
+    """Extract and apply ontology updates from response text.
+
+    Returns True if an update was applied.
+    """
+    update = extract_ontology_json(response_text)
+    if update is None:
+        return False
+    return merge_ontology_update(ontology, update)
+
+
+def _format_entities(ontology: Ontology) -> list[str]:
+    """Format entity lines for summary."""
+    lines = [f"Entities ({len(ontology.entities)}):"]
+    for e in ontology.entities:
+        props = ", ".join(p.name for p in e.properties)
+        lines.append(f"  {e.id}: {e.name} [{props}]")
+    return lines
+
+
+def _format_relationships(
+    ontology: Ontology,
+) -> list[str]:
+    """Format relationship lines for summary."""
+    lines = [
+        f"Relationships ({len(ontology.relationships)}):",
+    ]
+    for r in ontology.relationships:
+        lines.append(
+            f"  {r.source_entity_id} --{r.name}--> "
+            f"{r.target_entity_id} ({r.cardinality})",
+        )
+    return lines
+
+
+def _format_questions(ontology: Ontology) -> list[str]:
+    """Format open question lines for summary."""
+    lines = [
+        f"Open Questions ({len(ontology.open_questions)}):",
+    ]
+    for q in ontology.open_questions:
+        status = "RESOLVED" if q.resolved else "open"
+        lines.append(f"  [{status}] {q.id}: {q.text}")
+    return lines
+
+
+def format_ontology_summary(ontology: Ontology) -> str:
+    """Format a human-readable summary of the ontology."""
+    lines = _format_entities(ontology)
+    lines += _format_relationships(ontology)
+    n = len(ontology.domain_constraints)
+    lines.append(f"Constraints ({n}):")
+    for c in ontology.domain_constraints:
+        lines.append(f"  {c.name}: {c.description}")
+    lines += _format_questions(ontology)
+    return "\n".join(lines)
+
+
+def backtrack(dag: OntologyDAG) -> DAGNode | None:
+    """Move to the parent of the current node.
+
+    Returns the new current node, or None if at root.
+    """
+    parents = dag.parents_of(dag.current_node_id)
+    if not parents:
+        return None
+    dag.current_node_id = parents[0].id
+    return dag.get_current_node()
+
+
+_CmdResult = tuple[str, Ontology | None]
+
+
+def _handle_show(ontology: Ontology) -> _CmdResult:
+    """Handle the 'show' command."""
+    return format_ontology_summary(ontology), None
+
+
+def _handle_save(
+    command: str, ontology: Ontology,
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Handle the 'save' 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
+
+
+def _handle_back(
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Handle the 'back' command."""
+    node = backtrack(dag)
+    if node is None:
+        return "Already at root.", None
+    new_onto = node.ontology.model_copy(deep=True)
+    save_dag(dag, dag_path)
+    return (
+        f"Backtracked to: {node.id} ({node.label})",
+        new_onto,
+    )
+
+
+def is_command(text: str) -> bool:
+    """Check if user input is a meta-command."""
+    cmd = text.strip().lower()
+    return cmd in ("show", "back") or cmd.startswith(
+        "save",
+    )
+
+
+def handle_command(
+    command: str, ontology: Ontology,
+    dag: OntologyDAG, dag_path: str,
+) -> _CmdResult:
+    """Dispatch a user meta-command.
+
+    Returns (message, new_ontology_or_None).
+    """
+    cmd = command.strip().lower()
+    if cmd == "show":
+        return _handle_show(ontology)
+    if cmd == "back":
+        return _handle_back(dag, dag_path)
+    return _handle_save(command, ontology, dag, dag_path)
+
+
+async def print_response(client: Any) -> str:
+    """Receive and print response. Return full text."""
+    parts: list[str] = []
+    async for message in client.receive_response():
+        if isinstance(message, AssistantMessage):
+            print_text_blocks(message)
+            parts.append(collect_response_text(message))
+    return "\n".join(parts)
+
+
+def _init_ontology(dag: OntologyDAG) -> Ontology:
+    """Get ontology from current DAG node, or a new one."""
+    node = dag.get_current_node()
+    if node is None:
+        return Ontology()
+    return node.ontology.model_copy(deep=True)
+
+
+# taint: ignore[CWE-200] -- Interactive agent displays LLM output to user
+async def run(
+    description: str, model: str, dag_path: str,
+) -> None:
+    """Run the discovery agent interactively."""
+    prompt = discovery_system_prompt()
+    dag = load_dag(dag_path, description)
+    ontology = _init_ontology(dag)
+
+    options = ClaudeAgentOptions(
+        model=model,
+        system_prompt=prompt,
+        allowed_tools=[],
+        permission_mode="default",
+    )
+
+    async with ClaudeSDKClient(options=options) as client:
+        framed = frame_data("user-input", description)
+        await client.query(framed)
+        text = await print_response(client)
+        process_response(text, ontology)
+
+        while True:
+            user_input = read_user_input()
+            if user_input is None:
+                break
+            if is_command(user_input):
+                msg, new_onto = handle_command(
+                    user_input, ontology, dag, dag_path,
+                )
+                print(msg)
+                if new_onto is not None:
+                    ontology = new_onto
+                continue
+            framed = frame_data("user-input", user_input)
+            await client.query(framed)
+            text = await print_response(client)
+            process_response(text, ontology)
+
+
+def parse_args(
+    argv: list[str] | None = None,
+) -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Interactive ontology discovery agent",
+    )
+    parser.add_argument(
+        "description",
+        help="Project description to start discovery",
+    )
+    parser.add_argument(
+        "--dag-file",
+        default="ontology.json",
+        help="Path to DAG JSON file "
+        "(default: ontology.json)",
+    )
+    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-discovery-agent CLI."""
+    args = parse_args(argv)
+    asyncio.run(
+        run(args.description, args.model, args.dag_file),
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

python_agent/divergence_agent.py

@@ -0,0 +1,302 @@
+"""Divergence agent: generates multiple solution candidates."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import re
+import sys
+import warnings
+from typing import Any
+
+from claude_agent_sdk import (
+    AssistantMessage,
+    ClaudeAgentOptions,
+    ResultMessage,
+    query,
+)
+
+from python_agent.agent_utils import (
+    collect_response_text,
+    extract_ontology_json,
+)
+from python_agent.dag_utils import (
+    load_dag,
+    make_node_id,
+    save_dag,
+)
+from python_agent.ontology import (
+    DAGEdge,
+    DAGNode,
+    Decision,
+    Ontology,
+    OntologyDAG,
+)
+from python_agent.rules import (
+    divergence_system_prompt,
+    strategy_system_prompt,
+)
+
+_STRATEGIES_BLOCK_RE = re.compile(
+    r"```strategies\s*\n(.*?)\n```",
+    re.DOTALL,
+)
+
+
+def extract_strategies(
+    text: str,
+) -> list[dict[str, Any]] | None:
+    """Extract strategies JSON block from text.
+
+    Returns a list of strategy dicts, or None.
+    """
+    match = _STRATEGIES_BLOCK_RE.search(text)
+    if match is None:
+        return None
+    try:
+        result = json.loads(match.group(1))
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(result, list):
+        return None
+    return result
+
+
+async def run_query(
+    task: str, options: Any,
+) -> tuple[str, float]:
+    """Run a single query. Return (response_text, cost)."""
+    parts: list[str] = []
+    cost = 0.0
+    async for message in query(
+        prompt=task, options=options,
+    ):
+        if isinstance(message, AssistantMessage):
+            parts.append(collect_response_text(message))
+        elif isinstance(message, ResultMessage):
+            cost = message.total_cost_usd or 0.0
+    return "\n".join(parts), cost
+
+
+def build_decision(
+    strategy: dict[str, Any],
+) -> Decision:
+    """Create a Decision from a strategy dict."""
+    options = strategy.get("options", [])
+    if not isinstance(options, list):
+        warnings.warn(
+            f"Non-list options replaced: {type(options)}",
+            stacklevel=2,
+        )
+        options = []
+    return Decision(
+        question=str(strategy.get("question", "architecture")),
+        options=options,
+        chosen=str(strategy.get("chosen", "")),
+        rationale=str(strategy.get("strategy", "")),
+    )
+
+
+def add_candidate_node(
+    dag: OntologyDAG, parent_id: str,
+    ontology_dict: dict[str, Any],
+    strategy: dict[str, Any],
+) -> str:
+    """Create a DAG node for a candidate and link it."""
+    from datetime import datetime, timezone
+    now = datetime.now(timezone.utc)
+    node_id = make_node_id()
+    label: str = strategy.get("label", "candidate")
+    ontology = Ontology.model_validate(ontology_dict)
+    node = DAGNode(
+        id=node_id,
+        ontology=ontology,
+        created_at=now.isoformat(),
+        label=label,
+    )
+    dag.nodes.append(node)
+    edge = DAGEdge(
+        parent_id=parent_id,
+        child_id=node_id,
+        decision=build_decision(strategy),
+        created_at=now.isoformat(),
+    )
+    dag.edges.append(edge)
+    return node_id
+
+
+def remaining_budget(
+    spent: float, max_budget: float | None,
+) -> float | None:
+    """Calculate remaining budget after spending."""
+    if max_budget is None:
+        return None
+    return max_budget - spent
+
+
+async def identify_strategies(
+    ontology_json: str, num_candidates: int,
+    model: str, max_budget: float | None,
+) -> tuple[list[dict[str, Any]], float]:
+    """Identify distinct architectural strategies."""
+    prompt = strategy_system_prompt(
+        ontology_json, num_candidates,
+    )
+    options = ClaudeAgentOptions(
+        model=model,
+        system_prompt=prompt,
+        allowed_tools=[],
+        permission_mode="default",
+        max_turns=1,
+        max_budget_usd=max_budget,
+    )
+    task = (
+        f"Identify {num_candidates} distinct architectural "
+        "strategies for this problem domain."
+    )
+    text, cost = await run_query(task, options)
+    strategies = extract_strategies(text)
+    if strategies is None:
+        return [], cost
+    return strategies[:num_candidates], cost
+
+
+async def generate_candidate(
+    ontology_json: str, strategy: dict[str, Any],
+    model: str, max_budget: float | None,
+) -> tuple[dict[str, Any] | None, float]:
+    """Generate one solution candidate for a strategy."""
+    label: str = strategy.get("label", "candidate")
+    description: str = strategy.get("strategy", "")
+    prompt = divergence_system_prompt(
+        ontology_json, description,
+    )
+    options = ClaudeAgentOptions(
+        model=model,
+        system_prompt=prompt,
+        allowed_tools=[],
+        permission_mode="default",
+        max_turns=1,
+        max_budget_usd=max_budget,
+    )
+    task = (
+        f"Generate a complete solution architecture "
+        f"following the '{label}' strategy: {description}"
+    )
+    text, cost = await run_query(task, options)
+    result = extract_ontology_json(text)
+    return result, cost
+
+
+async def run(
+    dag_file: str, num_candidates: int,
+    model: str, max_budget: float | None,
+) -> int:
+    """Run the divergence agent. Returns candidate count."""
+    dag = load_dag(dag_file, "unknown")
+    node = dag.get_current_node()
+    if node is None:
+        print("Error: DAG has no current node.")
+        return 0
+    ontology_json = json.dumps(
+        node.ontology.model_dump(), indent=2,
+    )
+    parent_id = dag.current_node_id
+    total_cost = 0.0
+
+    print(f"Identifying {num_candidates} strategies...")
+    strategies, cost = await identify_strategies(
+        ontology_json, num_candidates, model, max_budget,
+    )
+    total_cost += cost
+    if not strategies:
+        print("Error: Could not identify strategies.")
+        return 0
+
+    generated = await _generate_all(
+        dag, parent_id, ontology_json, strategies,
+        model, total_cost, max_budget,
+    )
+    save_dag(dag, dag_file)
+    return generated
+
+
+async def _generate_all(
+    dag: OntologyDAG, parent_id: str,
+    ontology_json: str,
+    strategies: list[dict[str, Any]],
+    model: str, total_cost: float,
+    max_budget: float | None,
+) -> int:
+    """Generate all candidates and add to DAG."""
+    generated = 0
+    for strategy in strategies:
+        label: str = strategy.get("label", "candidate")
+        print(f"Generating candidate: {label}...")
+        budget = remaining_budget(total_cost, max_budget)
+        result, cost = await generate_candidate(
+            ontology_json, strategy, model, budget,
+        )
+        total_cost += cost
+        if result is None:
+            print(f"  Failed: {label}")
+            continue
+        add_candidate_node(
+            dag, parent_id, result, strategy,
+        )
+        generated += 1
+        print(f"  Created: {label}")
+    cost_str = f"${total_cost:.4f}"
+    print(f"\nDone. {generated} candidates. Cost: {cost_str}")
+    return generated
+
+
+def parse_args(
+    argv: list[str] | None = None,
+) -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Generate divergent solution candidates",
+    )
+    parser.add_argument(
+        "--dag-file",
+        required=True,
+        help="Path to DAG JSON file",
+    )
+    parser.add_argument(
+        "-n", "--num-candidates",
+        type=int,
+        default=3,
+        help="Number of candidates (default: 3)",
+    )
+    parser.add_argument(
+        "-m", "--model",
+        default="claude-sonnet-4-6",
+        help="Model to use (default: claude-sonnet-4-6)",
+    )
+    parser.add_argument(
+        "--max-budget",
+        type=float,
+        default=5.0,
+        help="Maximum budget in USD (default: 5.0)",
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Entry point for the aofire-divergence-agent CLI."""
+    args = parse_args(argv)
+    count = asyncio.run(
+        run(
+            args.dag_file,
+            args.num_candidates,
+            args.model,
+            args.max_budget,
+        ),
+    )
+    return 0 if count > 0 else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())