File size: 1,023 Bytes
ec9c941 | 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 | from langgraph.graph import StateGraph, END
from graph.state import GraphState
from graph.retrieve import retrieve
from graph.grade_documents import grade_documents
from graph.rewrite_query import rewrite_query
from graph.web_search import web_search
from graph.generate import generate
workflow = StateGraph(GraphState)
workflow.add_node("retrieve", retrieve)
workflow.add_node("grade", grade_documents)
workflow.add_node("rewrite", rewrite_query)
workflow.add_node("web", web_search)
workflow.add_node("generate", generate)
workflow.set_entry_point("retrieve")
workflow.add_edge(
"retrieve",
"grade"
)
def route(state):
if state["route"] == "local":
return "generate"
return "rewrite"
workflow.add_conditional_edges(
"grade",
route,
{
"generate": "generate",
"rewrite": "rewrite"
}
)
workflow.add_edge(
"rewrite",
"web"
)
workflow.add_edge(
"web",
"generate"
)
workflow.add_edge(
"generate",
END
)
app = workflow.compile() |