File size: 842 Bytes
2f18493
 
4861f80
 
 
2f18493
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from langgraph.graph import StateGraph, START, END
from langgraph.graph.state import CompiledStateGraph
from .asistant import Assistant
from .shared_state import SharedState
from .retrieval import Retrieval


def build_agent() -> "CompiledStateGraph":
    """
    Build the agent using the StateGraph.
    """
    builder = StateGraph(SharedState)
    assistant = Assistant()  # Initialize the assistant component
    retrieval = Retrieval()  # Initialize the retrieval component
    builder.add_node(
        "retrieval",
        retrieval,
    )

    builder.add_node(
        "assistant",
        assistant
    )

    builder.add_edge(
        START,
        "retrieval",
    )
    builder.add_edge(
        "retrieval",
        "assistant",
    )
    builder.add_edge(
        "assistant",
        END,
    )
    return builder.compile()