Spaces:
Sleeping
Sleeping
File size: 3,098 Bytes
5db30ac | 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 | 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
|