File size: 9,362 Bytes
7042c3c ba8e285 7042c3c ba8e285 7042c3c ba8e285 7042c3c ba8e285 7042c3c |
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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 |
"""Main file for constructing the EA4ALL hierarchical graph"""
"""
EA4ALL Hierarchical Graph
This module defines the main file for constructing the EA4ALL hierarchical graph. It contains functions and classes for creating and managing the graph structure.
Functions:
- make_supervisor_node: Creates a supervisor node for managing a conversation between architect workers.
- call_landscape_agentic: Calls the landscape agentic graph.
- call_diagram_agentic: Calls the diagram agentic graph.
- call_togaf_agentic: Calls the togaf agentic graph.
- websearch: Search for real-time data to answer user's question
Classes:
- Router: TypedDict representing the worker to route to next.
Attributes:
- model: The LLM client for the supervisor model.
- super_builder: The StateGraph builder for constructing the graph.
- super_graph: The compiled EA4ALL Agentic Workflow Graph.
Note: This module depends on other modules and packages such as langchain_core, langgraph, shared, ea4all_apm, ea4all_vqa, and ea4all_gra.
"""
"""Changelog:
- lanchain_openapi: 0.2.9 (0.3.6 issue with max_tokens for HF models)
#2025-06-03
- Refactored code to fix problems with linter and type checking (Standard mode)
"""
from langgraph.types import Command
from langchain_core.messages import (
HumanMessage,
AIMessage
)
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.runnables import RunnableConfig
from langchain import hub
from langgraph.graph import (
START,
END,
StateGraph,
)
from langgraph.checkpoint.memory import MemorySaver
from typing_extensions import Literal, TypedDict
import uuid
from ea4all.src.shared.configuration import BaseConfiguration
from ea4all.src.shared.utils import get_llm_client
from ea4all.src.shared.state import State
from ea4all.src.tools.tools import websearch
from ea4all.src.ea4all_indexer.graph import indexer_graph
from ea4all.src.ea4all_apm.graph import apm_graph
from ea4all.src.ea4all_vqa.graph import diagram_graph
from ea4all.src.ea4all_gra.graph import togaf_graph
async def call_indexer_apm(state: State, config: RunnableConfig):
response = await indexer_graph.ainvoke(input={"docs":[]}, config=config)
def make_supervisor_node(model: BaseChatModel, members: list[str]):
options = ["FINISH"] + members
system_prompt = hub.pull("ea4all_super_graph").template
class Router(TypedDict):
"""Worker to route to next. If no workers needed, route to FINISH."""
next: Literal["FINISH", "portfolio_team", "diagram_team", "blueprint_team", "websearch_team"]
async def supervisor_node(state: State, config: RunnableConfig) -> Command[Literal["portfolio_team", "diagram_team", "blueprint_team", "websearch_team", '__end__']]:
"""An LLM-based router."""
messages = [
{"role": "system", "content": system_prompt},
] + [state["messages"][-1]]
response = await model.with_structured_output(Router).ainvoke(messages, config=config)
_goto = "__end__"
if isinstance(response, dict):
_goto = response["next"]
# Ensure _goto is one of the allowed Literal values
if _goto not in ["portfolio_team", "diagram_team", "blueprint_team", "websearch_team"]:
_goto = "__end__"
print(f"---Supervisor got a request--- Question: {state['messages'][-1].content} ==> Routing to {_goto}\n")
return Command(
#update={"next": _goto},
goto=_goto
)
return supervisor_node
async def call_landscape_agentic(state: State, config: RunnableConfig) -> Command[Literal['__end__']]: ##2025-02-21: NOT passing CHAT MEMORY to the APM_graph
response = await apm_graph.ainvoke({"question": state["messages"][-1].content}, config=config)
return Command(
update={
"messages": [
AIMessage(
content=str(response), name="landscape_agentic"
)
]
},
goto="__end__",
)
async def call_diagram_agentic(state: State, config: RunnableConfig) -> Command[Literal['__end__']]: #NOT passing CHAT MEMORY to the Diagram_graph
inputs = {
"messages": [{"role": "user", "content": state.get('messages')[-1].content}],
"question": state['messages'][-1].content, "image":""
} #user response
response = await diagram_graph.ainvoke(
input=inputs,
config=config
)
return Command(
update={
"messages": [
AIMessage(
content=response['messages'][-1].content, name="landscape_agentic"
)
]
},
goto="__end__",
)
async def call_togaf_agentic(state: State, config: RunnableConfig) -> Command[Literal["__end__"]]: #NOT passing CHAT MEMORY to the Togaf_graph
print(f"---TOGAF ROUTE team node ready to --- CALL_TOGAF_AGENTIC Routing to {state['next']} with User Question: {state['messages'][-1].content}")
inputs = {"messages": [{"role": "user", "content": state.get('messages')[-1].content}]} #user response
response = await togaf_graph.ainvoke(
input=inputs,
config=config
) #astream not loading the graph
return Command(
update={
"messages": [
AIMessage(
content=response["messages"][-1].content, name="togaf_route"
)
]
},
goto="__end__",
)
# Wrap-up websearch answer to user's question
async def call_generate_websearch(state:State, config: RunnableConfig) -> Command[Literal["__end__"]]:
from ea4all.src.ea4all_apm.state import OverallState
if config is not None:
source = config.get('metadata', {}).get('langgraph_node', 'unknown')
# Invoke GENERATOR node in the APMGraph
state_dict = {
"documents": state['messages'][-1].content,
"web_search": "Yes",
"question": state['messages'][-2].content,
"source": source
}
apm_state = OverallState(**state_dict)
generation = await apm_graph.nodes["generate"].ainvoke(apm_state, config)
return Command(
update={
"messages": [
AIMessage(
content=generation['generation'], name="generate_websearch"
)
]
},
goto="__end__",
)
async def blueprint_team(state: State) -> Command[Literal["togaf_route"]]:
print("---Blueprint team got a request--- Routing to TOGAF_ROUTE node")
return Command(update={**state}, goto="togaf_route")
async def diagram_team(state: State) -> Command[Literal["diagram_route"]]:
print("---Diagram team got a request--- Routing to DIAGRAM_ROUTE node")
return Command(update={**state}, goto="diagram_route")
async def super_graph_entry_point(state: State):
# Generate a unique thread ID
thread_config = RunnableConfig({"configurable": {"thread_id": str(uuid.uuid4())}})
# Initialize state if not provided
if state is None:
state = {
"messages": [
("system", "You are a helpful assistant"),
("human", "Start the workflow")
]
}
# Build and compile the graph
graph = build_super_graph()
# Async invocation
try:
# Use ainvoke for async execution
result = await graph.ainvoke(state, config=RunnableConfig(thread_config))
return result
except Exception as e:
print(f"Graph execution error: {e}")
raise
# Define & build the graph.
def build_super_graph():
model = get_llm_client(BaseConfiguration.supervisor_model, api_base_url="", streaming=BaseConfiguration.streaming)
teams_supervisor_node = make_supervisor_node(model, ["portfolio_team", "diagram_team", "blueprint_team","websearch_team"])
super_builder = StateGraph(State, config_schema=BaseConfiguration)
super_builder.add_node("apm_indexer", call_indexer_apm)
super_builder.add_node("supervisor", teams_supervisor_node)
super_builder.add_node("portfolio_team", call_landscape_agentic)
super_builder.add_node("websearch_team", websearch)
super_builder.add_node("diagram_team", diagram_team)
super_builder.add_node("blueprint_team", blueprint_team)
super_builder.add_node("generate_websearch", call_generate_websearch)
super_builder.add_node("diagram_route", call_diagram_agentic)
super_builder.add_node("togaf_route", call_togaf_agentic)
super_builder.add_edge(START, "apm_indexer")
super_builder.add_edge("apm_indexer", "supervisor")
super_builder.add_edge("websearch_team", "generate_websearch")
super_builder.add_edge("blueprint_team", "togaf_route")
super_builder.add_edge("diagram_team", "diagram_route")
super_builder.add_edge("portfolio_team", END)
super_builder.add_edge("generate_websearch", END)
super_builder.add_edge("togaf_route", END)
super_builder.add_edge("diagram_route", END)
#memory = MemorySaver() #With LangGraph API, inMemmory is handled directly by the platform
super_graph = super_builder.compile() #checkpointer=memory)
super_graph.name = "EA4ALL Agentic Workflow Graph"
return super_graph
# Export the graph for LangGraph Dev/Studio
super_graph = build_super_graph()
|