graph.py

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, str(Path(__file__).parent / "core"))
sys.path.insert(0, str(Path(__file__).parent / "core" / "tools"))

# load_dotenv must precede all agent imports — agents import tools which load OpenAIEmbeddings
from dotenv import load_dotenv
load_dotenv()

from langgraph.graph import StateGraph, START, END
from langchain_core.messages import HumanMessage

from core.state import FinanceSystemState
from agents.context_extractor import context_extractor_node
from agents.orchestrator import orchestrator_routing_node, orchestrator_synthesis_node, route_to_agents
from agents.analyst import analyst_node
from agents.budget import budget_node
from agents.anomaly import anomaly_node
from agents.search import search_node
from agents.forecast import forecast_node

AGENT_NODES = ["analyst", "budget", "anomaly", "search", "forecast"]


def build_graph():
    graph = StateGraph(FinanceSystemState)

    # --- Nodes ---
    graph.add_node("context_extractor", context_extractor_node)
    graph.add_node("orchestrator_routing", orchestrator_routing_node)
    graph.add_node("analyst", analyst_node)
    graph.add_node("budget", budget_node)
    graph.add_node("anomaly", anomaly_node)
    graph.add_node("search", search_node)
    graph.add_node("forecast", forecast_node)
    graph.add_node("orchestrator_synthesis", orchestrator_synthesis_node)

    # --- Edges ---
    # Entry: extract context → route
    graph.add_edge(START, "context_extractor")
    graph.add_edge("context_extractor", "orchestrator_routing")

    # Fan-out: routing → 1–5 agents in parallel
    graph.add_conditional_edges(
        "orchestrator_routing",
        route_to_agents,
        AGENT_NODES,
    )

    # Fan-in: all agents converge at synthesis
    for agent in AGENT_NODES:
        graph.add_edge(agent, "orchestrator_synthesis")

    graph.add_edge("orchestrator_synthesis", END)

    return graph.compile()


# Compiled singleton — reused across calls
app = build_graph()


def run(question: str, user_role: str, user_department: str = None, **state_overrides) -> str:
    """
    End-to-end entry point. Returns the synthesized final response string.

    Args:
        question:        The user's natural-language question.
        user_role:       One of: CEO, CFO, Finance Analyst, Department Head.
        user_department: Required when user_role is 'Department Head'.
        **state_overrides: Any additional FinanceSystemState fields to pre-set
                           (e.g. current_year=2024 to skip extraction).
    """
    state = FinanceSystemState(
        messages=[HumanMessage(question)],
        user_role=user_role,
        user_department=user_department,
        **state_overrides,
    )
    result = app.invoke(state)
    return result["final_response"]


if __name__ == "__main__":
    scenarios = [
        {
            "question": "How much did we spend in 2024 broken down by department?",
            "user_role": "CFO",
        },
        {
            "question": "Were we over budget last year? Highlight anything over 10%.",
            "user_role": "CFO",
        },
        {
            "question": "Are there any anomalies in Engineering spending for 2024?",
            "user_role": "Finance Analyst",
            "user_department": "Engineering",
        },
        {
            "question": "What is our projected total spend over the next 6 months?",
            "user_role": "CEO",
        },
    ]

    for s in scenarios:
        q = s.pop("question")
        print(f"\n{'='*70}")
        print(f"Q: {q}")
        print(f"Role: {s.get('user_role')}")
        print("─" * 70)
        answer = run(q, **s)
        print(answer)