"""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()