src/agent/rlm_agent.py

"""RLM-on-KG agent for kinship reasoning.

Implements a ReAct-style agentic loop where a GRPO-trained Gemma 4B model
navigates a kinship knowledge graph via tool calls to answer inheritance
questions.

Architecture follows RLM-on-KG (arXiv:2604.17056):
- State = (current entity, visited set, hop index)
- Actions = 4-tool MDP (search_entities, get_entity, follow_entity_link, expand_neighbors)
- Dynamic Hop Control: model decides when to stop and submit answer

Output format follows SEOcrate:
- <reasoning>...</reasoning> for step-by-step thought
- <action tool="..." entity="..." property="..."/> for tool calls
- <answer>...</answer> for final submission
"""

from __future__ import annotations

import json
import re
import time
from dataclasses import dataclass, field
from typing import Any

from src.graph.tools import GraphToolsBase


# ---------------------------------------------------------------------------
# System prompt — the ontological rules the agent must follow
# ---------------------------------------------------------------------------

KINSHIP_SYSTEM_PROMPT = """\
You are a kinship reasoning agent. You navigate a knowledge graph to answer \
inheritance questions by making tool calls and reading the results.

## KINSHIP ONTOLOGY

The knowledge graph uses these classes and properties:

CLASSES:
- Person (has: fullName, age, isAlive)
- Asset (has: assetName, assetValue, belongsTo)
- WillClause (has: clauseText, excludesRelationType, requiresRelationType)

OBJECT PROPERTIES:
- isBiologicalParentOf / hasBiologicalParent → INVERSE PAIR. Directionality matters.
- isMarriedTo → SYMMETRIC. Marriage is mutual. A symmetric property can NEVER \
satisfy a "biological lineage" requirement.
- stipulatesCondition → WillClause conditions on a Person.
- belongsTo → Asset ownership.
- excludesRelationType → What property types the will FORBIDS.
- requiresRelationType → What property types the will REQUIRES.

## TOOLS

You have 4 tools to call via the knowledge graph:

1. search_entities(query) — Find entities by name. Start here. Takes a name string.
2. get_entity(iri) — Get all data properties for an entity (age, isAlive, type). REQUIRES an IRI.
3. follow_entity_link(iri, property?) — Follow outbound edges. Filter by property URI. REQUIRES an IRI.
4. expand_neighbors(iri) — Get ALL connected entities (inbound + outbound). REQUIRES an IRI.

IMPORTANT: get_entity, follow_entity_link, and expand_neighbors require entity IRIs (e.g. http://example.org/kinship#Person5).
If you only know a name, call search_entities first to get the IRI.
When a Tool result contains an "iri" or "target_iri" field, copy that exact value into the next action.
Never pass a person's full name as the entity argument to get_entity.

## CRITICAL ONTOLOGICAL RULES

1. READ the WillClause first. Check excludesRelationType and requiresRelationType.
2. NEVER traverse a property the will excludes (e.g., isMarriedTo when excluded).
3. isMarriedTo is a SymmetricProperty — it CANNOT establish biological lineage.
4. Always verify isAlive=true before declaring someone an heir.
5. Use the FEWEST hops possible (Dynamic Hop Control).

## OUTPUT FORMAT

For each step, output:
<reasoning>Your step-by-step ontological reasoning about what to do next</reasoning>
<action tool="tool_name" entity="entity_iri" property="property_uri"/>

When you have enough evidence, output:
<reasoning>Your final synthesis of the evidence</reasoning>
<answer>Full name of the heir</answer>

Property URIs use the namespace http://example.org/kinship#
Example: http://example.org/kinship#isBiologicalParentOf
"""


# ---------------------------------------------------------------------------
# Agent state (MDP)
# ---------------------------------------------------------------------------

@dataclass
class AgentState:
    """MDP state for the RLM-on-KG agent.

    Attributes:
        current_entity: IRI of the entity currently in focus.
        visited: Set of visited entity IRIs (loop avoidance).
        hop_index: Current hop count.
        hop_budget: Maximum allowed hops before forced termination.
        evidence: List of evidence packets collected during traversal.
    """
    current_entity: str | None = None
    visited: set[str] = field(default_factory=set)
    hop_index: int = 0
    hop_budget: int = 8
    evidence: list[dict] = field(default_factory=list)


# ---------------------------------------------------------------------------
# XML tag parsing
# ---------------------------------------------------------------------------

_ACTION_RE = re.compile(
    r'<action\s+tool=["\'](\w+)["\']\s+'
    r'entity=["\']([^"\']*)["\']'
    r'(?:\s+property=["\']([^"\']*)["\'])?\s*/?>',
    re.IGNORECASE,
)

_ANSWER_RE = re.compile(
    r'<answer>(.*?)</answer>',
    re.IGNORECASE | re.DOTALL,
)

_REASONING_RE = re.compile(
    r'<reasoning>(.*?)</reasoning>',
    re.IGNORECASE | re.DOTALL,
)


def parse_action(text: str) -> dict | None:
    """Parse the first <action> tag in *text*.

    Returns:
        Dict with keys 'tool', 'entity', 'property' (optional), or None.
    """
    m = _ACTION_RE.search(text)
    if not m:
        return None
    result: dict[str, str] = {
        "tool": m.group(1),
        "entity": m.group(2),
    }
    if m.group(3):
        result["property"] = m.group(3)
    return result


def parse_answer(text: str) -> str | None:
    """Extract the content of the first <answer> tag."""
    m = _ANSWER_RE.search(text)
    return m.group(1).strip() if m else None


def parse_reasoning(text: str) -> str | None:
    """Extract the content of the first <reasoning> tag."""
    m = _REASONING_RE.search(text)
    return m.group(1).strip() if m else None



# ---------------------------------------------------------------------------
# Tool executor
# ---------------------------------------------------------------------------

def execute_tool(tools: GraphToolsBase, action: dict) -> Any:
    """Dispatch a parsed action dict to the appropriate graph tool.

    Args:
        tools: A GraphToolsBase instance (RdflibTools or WoGraphTools).
        action: Dict with 'tool', 'entity', and optionally 'property'.

    Returns:
        The tool's result (list or dict).
    """
    tool_name = action["tool"]
    entity = action.get("entity", "")
    prop = action.get("property")

    if tool_name == "search_entities":
        # For search, the 'entity' field is used as the query string
        return tools.search_entities(entity)
    elif tool_name == "get_entity":
        return tools.get_entity(entity)
    elif tool_name == "follow_entity_link":
        return tools.follow_entity_link(entity, prop)
    elif tool_name == "expand_neighbors":
        return tools.expand_neighbors(entity)
    else:
        return {"error": f"Unknown tool: {tool_name}"}


# ---------------------------------------------------------------------------
# Agent loop
# ---------------------------------------------------------------------------

@dataclass
class AgentResult:
    """Result from a completed agent run."""
    answer: str | None
    reasoning: str | None
    trace: list[dict]
    hops: int
    elapsed_seconds: float
    terminated_by: str  # "answer" | "timeout" | "error"


def run_agent(
    generate_fn,
    tools: GraphToolsBase,
    narrative: str,
    question: str,
    *,
    hop_budget: int = 8,
    verbose: bool = False,
) -> AgentResult:
    """Run the RLM-on-KG agentic loop.

    Args:
        generate_fn: A callable ``(messages: list[dict]) -> str`` that
            generates a model completion given chat messages. This abstracts
            away the specific model backend (transformers, vLLM, API, etc.).
        tools: A GraphToolsBase instance for graph queries.
        narrative: The text narrative (with semantic distractors).
        question: The inheritance question.
        hop_budget: Maximum hops before forced termination.
        verbose: If True, print each step to stdout.

    Returns:
        An AgentResult with the answer, reasoning trace, and metadata.
    """
    state = AgentState(hop_budget=hop_budget)
    start_time = time.time()

    messages: list[dict[str, str]] = [
        {"role": "system", "content": KINSHIP_SYSTEM_PROMPT},
        {
            "role": "user",
            "content": (
                f"Narrative:\n{narrative}\n\n"
                f"Question:\n{question}"
            ),
        },
    ]

    trace: list[dict] = []

    while state.hop_index < state.hop_budget:
        # Generate model response
        response = generate_fn(messages)

        if verbose:
            print(f"\n--- Hop {state.hop_index} ---")
            print(response)

        # Check for final answer (Dynamic Hop Control)
        answer = parse_answer(response)
        if answer:
            reasoning = parse_reasoning(response)
            return AgentResult(
                answer=answer,
                reasoning=reasoning,
                trace=trace,
                hops=state.hop_index,
                elapsed_seconds=time.time() - start_time,
                terminated_by="answer",
            )

        # Parse and execute action
        action = parse_action(response)
        if not action:
            # Model didn't produce a valid action or answer — add guidance
            messages.append({"role": "assistant", "content": response})
            messages.append({
                "role": "user",
                "content": (
                    "Please respond with either an <action> tag to query the "
                    "graph, or an <answer> tag with your final answer."
                ),
            })
            state.hop_index += 1
            continue

        # Execute the tool call
        try:
            result = execute_tool(tools, action)
        except Exception as e:
            result = {"error": str(e)}

        # Record in trace
        trace_entry = {
            "hop": state.hop_index,
            "action": action,
            "result": result,
            "reasoning": parse_reasoning(response),
        }
        trace.append(trace_entry)

        if verbose:
            print(f"  Action: {action}")
            print(f"  Result: {json.dumps(result, default=str)[:200]}")

        # Update state
        entity = action.get("entity", "")
        state.visited.add(entity)
        state.current_entity = entity
        state.hop_index += 1

        # Append to conversation
        messages.append({"role": "assistant", "content": response})
        messages.append({
            "role": "user",
            "content": f"Tool result:\n{json.dumps(result, default=str, indent=2)}",
        })

    # Budget exhausted
    return AgentResult(
        answer=None,
        reasoning="Hop budget exhausted without reaching a conclusion.",
        trace=trace,
        hops=state.hop_index,
        elapsed_seconds=time.time() - start_time,
        terminated_by="timeout",
    )


# ---------------------------------------------------------------------------
# Simple generate_fn for testing (uses a mock or Gemini API)
# ---------------------------------------------------------------------------

def make_gemini_generate_fn(model_name: str = "gemini-2.5-flash"):
    """Create a generate_fn that calls the Gemini API.

    Requires the ``google-genai`` package and ``GOOGLE_API_KEY`` env var.
    """
    import os
    from google import genai

    api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
    client = genai.Client(api_key=api_key)

    def generate(messages: list[dict[str, str]]) -> str:
        # Convert our messages to Gemini format
        contents = []
        for msg in messages:
            role = "model" if msg["role"] == "assistant" else "user"
            if msg["role"] == "system":
                # Prepend system as user message
                contents.append({"role": "user", "parts": [{"text": msg["content"]}]})
                contents.append({"role": "model", "parts": [{"text": "Understood. I will follow these instructions."}]})
            else:
                contents.append({"role": role, "parts": [{"text": msg["content"]}]})

        response = client.models.generate_content(
            model=model_name,
            contents=contents,
        )
        return response.text

    return generate


# ---------------------------------------------------------------------------
# Quick demo
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys
    sys.path.insert(0, ".")

    from src.graph.graph_builder import get_arthur_scenario, build_rdflib_graph
    from src.graph.tools import RdflibTools
    from src.graph.scenario_generator import generate_narrative, generate_question

    # Build the Arthur scenario
    scenario = get_arthur_scenario()
    graph = build_rdflib_graph(scenario)
    tools = RdflibTools(graph)

    narrative = generate_narrative(scenario)
    question = generate_question(scenario)

    print("=" * 60)
    print("NARRATIVE:")
    print(narrative)
    print("\nQUESTION:", question)
    print("GOLD ANSWER:", scenario["gold_answer"])
    print("=" * 60)

    # Try with Gemini if available
    try:
        generate_fn = make_gemini_generate_fn()
        result = run_agent(generate_fn, tools, narrative, question, verbose=True)
        print(f"\n{'=' * 60}")
        print(f"ANSWER: {result.answer}")
        print(f"CORRECT: {result.answer and 'Edward' in result.answer}")
        print(f"HOPS: {result.hops}")
        print(f"TIME: {result.elapsed_seconds:.2f}s")
        print(f"TERMINATED BY: {result.terminated_by}")
    except Exception as e:
        print(f"\nCould not run with Gemini API: {e}")
        print("Set GOOGLE_API_KEY to test with a real model.")