CodeRouterAI/workflow.py at main · debraj7777/CodeRouterAI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
from typing import TypedDict
from langgraph.graph import StateGraph, END
from dotenv import load_dotenv

from agent.classifier import classify_complexity
from router import route_query
from local import call_local_model
from remote import call_remote_model
from cost_tracker import log_query

load_dotenv()

# ─── State Definition ─────────────────────────────────────────────────────────

class AgentState(TypedDict):
    query: str
    user_id: int
    complexity_score: int
    model_choice: str
    response: str
    model_used: str
    tokens: int
    cost_saved: float
    cost_incurred: float
    fallback_used: bool

# ─── Node Functions ───────────────────────────────────────────────────────────

def classifier_node(state: AgentState) -> AgentState:
    score = classify_complexity(state["query"])
    return {**state, "complexity_score": score}

def router_node(state: AgentState) -> AgentState:
    choice = route_query(state["complexity_score"])
    return {**state, "model_choice": choice}

def local_node(state: AgentState) -> AgentState:
    result = call_local_model(state["query"])
    log_query(state["user_id"], state["query"], state["complexity_score"], "local", result)
    return {
        **state,
        "response": result["response"],
        "model_used": result["model_used"],
        "tokens": result["tokens"],
        "cost_saved": result["cost_saved"],
        "cost_incurred": result["cost_incurred"],
        "fallback_used": False
    }

def remote_node(state: AgentState) -> AgentState:
    """Handles query with Fireworks AI."""
    try:
        result = call_remote_model(state["query"])
        log_query(state["user_id"], state["query"], state["complexity_score"], "remote", result)
        return {
            **state,
            "response": result["response"],
            "model_used": result["model_used"],
            "tokens": result["tokens"],
            "cost_saved": result["cost_saved"],
            "cost_incurred": result["cost_incurred"],
            "fallback_used": False
        }
    except Exception as e:
        error_msg = (
            f"**Error Calling Remote Model**\n\n"
            f"The task complexity ({state['complexity_score']}/5) exceeds local model capabilities, "
            f"and the remote model call failed:\n"
            f"`{str(e)}`\n\n"
            f"Please check your `FIREWORKS_API_KEY` configuration and try again."
        )
        error_result = {
            "response": error_msg,
            "model_used": "fireworks-api (failed)",
            "tokens": 0,
            "cost_saved": 0.0,
            "cost_incurred": 0.0
        }
        log_query(state["user_id"], state["query"], state["complexity_score"], "remote (failed)", error_result)
        return {
            **state,
            "response": error_msg,
            "model_used": "fireworks-api (failed)",
            "tokens": 0,
            "cost_saved": 0.0,
            "cost_incurred": 0.0,
            "fallback_used": False
        }

# ─── Routing Logic ────────────────────────────────────────────────────────────

def decide_route(state: AgentState) -> str:
    return state["model_choice"]

# ─── Build Graph ──────────────────────────────────────────────────────────────

def build_graph():
    graph = StateGraph(AgentState)

    graph.add_node("classifier", classifier_node)
    graph.add_node("router", router_node)
    graph.add_node("local", local_node)
    graph.add_node("remote", remote_node)

    graph.set_entry_point("classifier")
    graph.add_edge("classifier", "router")
    graph.add_conditional_edges(
        "router",
        decide_route,
        {"local": "local", "remote": "remote"}
    )
    graph.add_edge("local", END)
    graph.add_edge("remote", END)

    return graph.compile()

coderouter = build_graph()