-
Notifications
You must be signed in to change notification settings - Fork 0
/
SourceCode.py
95 lines (71 loc) · 2.89 KB
/
SourceCode.py
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
from Agent_Creation import create_agent
from Node_Creation import agent_node
import functools
import RequestAndInsertionTool
from LLM_Initialization import llm_initialization
from PythonREPLTool import get_REPL_Tool
from Supervisor import Creation_Of_Supervisor_Chain,get_members
from langgraph.graph import StateGraph, END, START
from langchain_core.messages import BaseMessage, HumanMessage
from Graph_Display import display_graph
from QueryGenerator import get_sql_flow
from typing import Annotated,Sequence, TypedDict
import operator
# The agent state is the input to each node in the graph
class AgentState(TypedDict):
# The annotation tells the graph that new messages will always
# be added to the current states
messages: Annotated[Sequence[BaseMessage], operator.add]
# The 'next' field indicates where to route to next
next: str
## Creation of Extractor Agent
extracter_agent = create_agent(
llm_initialization(),
[RequestAndInsertionTool.RequestAndInsertTool],
"""
you're a helpful assistant responsible for extracting the information from the URL
and the extracted details will be inserted to the database of the Person Table using psycopg2"""
)
extracter_node = functools.partial(agent_node, agent=extracter_agent, name="Extractor")
## Creation of Visualizer Agent
visualizer_agent = create_agent(
llm_initialization(),
[get_REPL_Tool()],
"You may generate safe python code to visualize data and generate bar plot charts using matplotlib.",
)
visualizer_node = functools.partial(agent_node, agent=visualizer_agent, name="Visualizer")
## Creation of SQL Analyzer Agent
app = get_sql_flow()
## Creation of Supervisor
workflow = StateGraph(AgentState)
workflow.add_node("Extractor", extracter_node)
workflow.add_node("Analyzer", app)
workflow.add_node("Visualizer", visualizer_node)
workflow.add_node("supervisor", Creation_Of_Supervisor_Chain(llm_initialization()))
## Adding the Edges between the Nodes of the graph
members = get_members()
for member in members:
# We want our workers to ALWAYS "report back" to the supervisor when done
workflow.add_edge(member, "supervisor")
# The supervisor populates the "next" field in the graph state
# which routes to a node or finishes
conditional_map = {k: k for k in members}
conditional_map["FINISH"] = END
workflow.add_conditional_edges("supervisor", lambda x: x["next"], conditional_map)
# Finally, add entrypoint
workflow.set_entry_point("supervisor")
graph = workflow.compile()
## Graph Stream
for s in graph.stream(
{
"messages": [
HumanMessage(content="""
Analyze the number of people based on their gender """)
]
}
):
if "__end__" not in s:
print(s)
print("----")
## Displaying the Graph
display_graph(graph)