Final_Assignment_Template / graph_builder.py
Umer797's picture
Update graph_builder.py
b48a433 verified
raw
history blame
904 Bytes
from langgraph.graph import StateGraph
def build_graph():
graph = StateGraph(dict)
def llm_step(state):
question = state.get("question")
if not question:
raise ValueError("State is missing 'question'")
# This is where you integrate your LLM call.
llm_output = f"Dummy answer to: {question}"
state["llm_output"] = llm_output
return state
def formatter_step(state):
llm_output = state.get("llm_output")
if not llm_output:
raise ValueError("State is missing 'llm_output'")
# For now, we just pass it through.
state["final_answer"] = llm_output.strip()
return state
graph.add_node("llm", llm_step)
graph.add_node("formatter", formatter_step)
graph.set_entry_point("llm")
graph.add_edge("llm", "formatter")
graph.set_finish_point("formatter")
return graph