109 lines
3.0 KiB
Python
109 lines
3.0 KiB
Python
"""
|
|
A minimal LangGraph agent implementation.
|
|
|
|
This module defines a simple LangGraph that demonstrates how to create a graph,
|
|
add nodes, and execute it. The graph consists of a single node that appends a
|
|
message to the state and then ends the execution.
|
|
|
|
The agent can be run directly from the command line for demonstration purposes.
|
|
"""
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Dict, Any
|
|
|
|
# Import LangGraph components
|
|
try:
|
|
from langgraph.graph import StateGraph, END
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"langgraph is not installed. Please add 'langgraph' to your requirements.txt "
|
|
"and run 'pip install -r requirements.txt'."
|
|
) from exc
|
|
|
|
|
|
@dataclass
|
|
class AgentState:
|
|
"""
|
|
The state that flows through the graph.
|
|
|
|
Attributes
|
|
----------
|
|
messages : List[str]
|
|
A list of messages that the agent accumulates during execution.
|
|
"""
|
|
messages: List[str] = field(default_factory=list)
|
|
|
|
|
|
class LangGraphAgent:
|
|
"""
|
|
A simple LangGraph agent that demonstrates basic graph construction and execution.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
"""
|
|
Initialize the graph and define its nodes and edges.
|
|
"""
|
|
self.graph = StateGraph(AgentState)
|
|
|
|
# Add nodes
|
|
self.graph.add_node("start", self._start_node)
|
|
self.graph.add_node("end", self._end_node)
|
|
|
|
# Define the entry point and transitions
|
|
self.graph.set_entry_point("start")
|
|
self.graph.add_edge("start", "end")
|
|
self.graph.add_edge("end", END)
|
|
|
|
# Compile the graph into a runnable function
|
|
self._graph_fn = self.graph.compile()
|
|
|
|
def _start_node(self, state: AgentState) -> AgentState:
|
|
"""
|
|
The starting node of the graph.
|
|
|
|
It appends a greeting message to the state's messages list.
|
|
"""
|
|
state.messages.append("Hello from LangGraph!")
|
|
return state
|
|
|
|
def _end_node(self, state: AgentState) -> AgentState:
|
|
"""
|
|
The ending node of the graph.
|
|
|
|
Currently, it performs no additional processing.
|
|
"""
|
|
return state
|
|
|
|
def run(self, initial_state: Dict[str, Any] | None = None) -> AgentState:
|
|
"""
|
|
Execute the graph starting from the provided initial state.
|
|
|
|
Parameters
|
|
----------
|
|
initial_state : dict or None
|
|
Optional dictionary to initialize the AgentState. If None, an empty state
|
|
is used.
|
|
|
|
Returns
|
|
-------
|
|
AgentState
|
|
The final state after graph execution.
|
|
"""
|
|
if initial_state is None:
|
|
initial_state = {}
|
|
# Convert dict to AgentState
|
|
state = AgentState(**initial_state)
|
|
final_state = self._graph_fn(state)
|
|
return final_state
|
|
|
|
|
|
if __name__ == "__main__":
|
|
"""
|
|
Example usage of the LangGraphAgent.
|
|
|
|
Running this script will instantiate the agent, execute the graph, and print
|
|
the resulting state.
|
|
"""
|
|
agent = LangGraphAgent()
|
|
result = agent.run()
|
|
print("Final state messages:", result.messages) |