Spaces:
Sleeping
Sleeping
| 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()) | |