Spaces:
Sleeping
Sleeping
| import asyncio | |
| from typing import Dict, Any, Tuple, Optional | |
| from uuid import uuid4 | |
| from langgraph.graph.state import CompiledStateGraph | |
| from langgraph.graph import StateGraph, START, END | |
| from langgraph.prebuilt.chat_agent_executor import AgentState | |
| from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver | |
| import aiosqlite | |
| from elbench.basic_education_runtime.entity import AgentConfig | |
| from elbench.basic_education_runtime.router import * # noqa: F403 | |
| from elbench.basic_education_runtime.config import CONFIG | |
| def add_node_to_graph(graph: StateGraph, node_id: str, node_instance: Any) -> None: | |
| if node_id not in graph.nodes: | |
| graph.add_node(node_id, node_instance) | |
| else: | |
| pass | |
| async def apply_agent_direction_from_dict( | |
| agent_map: Dict[str, Tuple[CompiledStateGraph, AgentConfig]], | |
| memory_id: Optional[str] = None, | |
| task: Optional[Dict[str, str]] = None, | |
| ) -> Tuple[CompiledStateGraph, str]: | |
| if memory_id is None: | |
| memory_id = str(uuid4()) | |
| directions = CONFIG.directions | |
| graph = StateGraph(AgentState) | |
| for node_id, node_instance in agent_map.items(): | |
| graph.add_node(node_id, node_instance[0]) | |
| for direction in directions: | |
| start_node, end_node = (i.strip() for i in direction.split("->")) | |
| if start_node == "START": | |
| start_node = START | |
| if end_node.startswith("router:"): | |
| function_call = end_node[len("router:") :] | |
| route, path_map = eval(function_call) | |
| end_node = end_node.replace(":", "_") | |
| graph.add_conditional_edges(start_node, route, path_map) | |
| continue | |
| # add_node_to_graph(graph, end_node, route) | |
| elif end_node == "END": | |
| end_node = END | |
| else: | |
| agent_config = agent_map[end_node][1] | |
| pregel_instance = agent_map[end_node][0] | |
| if not pregel_instance or not agent_config: | |
| raise ValueError(f"Invalid configuration for {end_node}.") | |
| graph.add_edge(start_node, end_node) | |
| CONFIG.globals.memory.path.mkdir(parents=True, exist_ok=True) | |
| path = CONFIG.globals.memory.path / f"{memory_id}.db" | |
| # conn = sqlite3.connect(f"{memory_id}.db", check_same_thread=False) | |
| if task is not None: | |
| async with aiosqlite.connect(path) as conn: | |
| sql = "create table task (key TEXT, value TEXT)" | |
| await conn.execute(sql) | |
| sql = "insert into task (key, value) values (?, ?)" | |
| for key, value in task.items(): | |
| await conn.execute(sql, (key, value)) | |
| await conn.commit() | |
| # await conn.close() | |
| conn = aiosqlite.connect(path, check_same_thread=False) | |
| if not hasattr(conn, "is_alive"): | |
| conn.is_alive = lambda: bool( # type: ignore[attr-defined] | |
| getattr(conn, "_running", False) and getattr(conn, "_connection", None) is not None | |
| ) | |
| CONFIG.context.conns.append(conn) | |
| memory = AsyncSqliteSaver(conn) | |
| return graph.compile(checkpointer=memory), memory_id | |
| apply_agent_direction = apply_agent_direction_from_dict | |