graph.py

from typing import TypedDict

from langgraph.graph import StateGraph, END

from agents.jd_agent import run_jd_agent
from agents.resume_agent import run_resume_agent
from agents.fit_agent import run_fit_agent
from agents.writer_agent import run_writer_agent
from agents.reviewer_agent import run_reviewer_agent


class AgentState(TypedDict):
    job_description: str
    resume: str
    jd_analysis: str
    resume_analysis: str
    fit_analysis: str
    draft: str
    final_answer: str


def jd_node(state):
    result = run_jd_agent(state["job_description"])
    return {"jd_analysis": result.content}


def resume_node(state):
    result = run_resume_agent(state["resume"])
    return {"resume_analysis": result.content}


def fit_node(state):
    result = run_fit_agent(
        state["jd_analysis"],
        state["resume_analysis"]
    )
    return {"fit_analysis": result.content}


def writer_node(state):
    result = run_writer_agent(state["fit_analysis"])
    return {"draft": result.content}


def reviewer_node(state):
    result = run_reviewer_agent(
        state["draft"],
        state["job_description"],
        state["resume"],
    )
    return {"final_answer": result.content}


workflow = StateGraph(AgentState)

workflow.add_node("jd_agent", jd_node)
workflow.add_node("resume_agent", resume_node)
workflow.add_node("fit_agent", fit_node)
workflow.add_node("writer_agent", writer_node)
workflow.add_node("reviewer_agent", reviewer_node)

workflow.set_entry_point("jd_agent")

workflow.add_edge("jd_agent", "resume_agent")
workflow.add_edge("resume_agent", "fit_agent")
workflow.add_edge("fit_agent", "writer_agent")
workflow.add_edge("writer_agent", "reviewer_agent")
workflow.add_edge("reviewer_agent", END)

app = workflow.compile()