SurfSense/surfsense_backend/app/agents/researcher/graph.py
2025-04-19 23:25:06 -07:00

43 lines
1.4 KiB
Python

from langgraph.graph import StateGraph
from .state import State
from .nodes import write_answer_outline, process_sections
from .configuration import Configuration
from typing import TypedDict, List, Dict, Any, Optional
# Define what keys are in our state dict
class GraphState(TypedDict):
# Intermediate data produced during workflow
answer_outline: Optional[Any]
# Final output
final_written_report: Optional[str]
def build_graph():
"""
Build and return the LangGraph workflow.
This function constructs the researcher agent graph with proper state management
and node connections following LangGraph best practices.
Returns:
A compiled LangGraph workflow
"""
# Define a new graph with state class
workflow = StateGraph(State, config_schema=Configuration)
# Add nodes to the graph
workflow.add_node("write_answer_outline", write_answer_outline)
workflow.add_node("process_sections", process_sections)
# Define the edges - create a linear flow
workflow.add_edge("__start__", "write_answer_outline")
workflow.add_edge("write_answer_outline", "process_sections")
workflow.add_edge("process_sections", "__end__")
# Compile the workflow into an executable graph
graph = workflow.compile()
graph.name = "Surfsense Researcher" # This defines the custom name in LangSmith
return graph
# Compile the graph once when the module is loaded
graph = build_graph()