Spaces:
Sleeping
Sleeping
File size: 2,634 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 80 81 82 83 | import asyncio
from typing import Dict, Optional
from tqdm.asyncio import tqdm
from langgraph.graph.state import CompiledStateGraph
from langgraph.errors import GraphRecursionError
from elbench.basic_education_runtime.agent import init_agent_map
from elbench.basic_education_runtime.config import CONFIG
from elbench.basic_education_runtime.directions import apply_agent_direction
from elbench.basic_education_runtime.entity import Prompt
from elbench.basic_education_runtime.model import init_model_map
from elbench.basic_education_runtime.utils import replace_prompt
from tenacity import RetryError
async def run(workers_num: int = CONFIG.globals.concurrency):
sem = asyncio.Semaphore(workers_num)
model_map = init_model_map()
agents = []
for task in CONFIG.tasks.variables:
agent_map, task = init_agent_map(model_map, task)
if CONFIG.tasks.start_prompt is not None:
if task is not None:
start_prompt = replace_prompt(CONFIG.tasks.start_prompt, task)
else:
start_prompt = CONFIG.tasks.start_prompt
else:
start_prompt = None
agent, _ = await apply_agent_direction(agent_map, task=task)
agents.append((agent, start_prompt))
async def arun(
agent: CompiledStateGraph, prompt: Optional[Prompt | Dict[str, str]]
):
if isinstance(prompt, Prompt):
n_prompt = prompt.model_dump()
elif prompt is None:
n_prompt = []
else:
n_prompt = prompt
async with sem:
try:
await agent.ainvoke(
{"messages": n_prompt},
{
"configurable": {"thread_id": "0"},
"recursion_limit": CONFIG.globals.recursion_limit,
},
stream_mode="values",
)
except GraphRecursionError:
print(
f"Recursion limit {CONFIG.globals.recursion_limit} reached for one task"
)
return
except RetryError as e:
exception = e.last_attempt.exception()
if exception is not None:
raise exception
else:
raise ValueError("Retry error occurred without exception")
tasks = []
for agent, prompt in agents:
tasks.append(arun(agent, prompt))
await tqdm.gather(*tasks)
conns = CONFIG.context.conns
for conn in conns:
await conn.close()
if __name__ == "__main__":
asyncio.run(run())
|