|
|
"""Hierarchical orchestrator using middleware and sub-teams.""" |
|
|
|
|
|
import asyncio |
|
|
from collections.abc import AsyncGenerator |
|
|
|
|
|
import structlog |
|
|
|
|
|
from src.agents.judge_agent_llm import LLMSubIterationJudge |
|
|
from src.agents.magentic_agents import create_search_agent |
|
|
from src.middleware.sub_iteration import SubIterationMiddleware, SubIterationTeam |
|
|
from src.services.embeddings import get_embedding_service |
|
|
from src.state import init_magentic_state |
|
|
from src.utils.models import AgentEvent |
|
|
|
|
|
logger = structlog.get_logger() |
|
|
|
|
|
|
|
|
class ResearchTeam(SubIterationTeam): |
|
|
"""Adapts Magentic ChatAgent to SubIterationTeam protocol.""" |
|
|
|
|
|
def __init__(self) -> None: |
|
|
self.agent = create_search_agent() |
|
|
|
|
|
async def execute(self, task: str) -> str: |
|
|
response = await self.agent.run(task) |
|
|
if response.messages: |
|
|
for msg in reversed(response.messages): |
|
|
if msg.role == "assistant" and msg.text: |
|
|
return str(msg.text) |
|
|
return "No response from agent." |
|
|
|
|
|
|
|
|
class HierarchicalOrchestrator: |
|
|
"""Orchestrator that uses hierarchical teams and sub-iterations.""" |
|
|
|
|
|
def __init__(self) -> None: |
|
|
self.team = ResearchTeam() |
|
|
self.judge = LLMSubIterationJudge() |
|
|
self.middleware = SubIterationMiddleware(self.team, self.judge, max_iterations=5) |
|
|
|
|
|
async def run(self, query: str) -> AsyncGenerator[AgentEvent, None]: |
|
|
logger.info("Starting hierarchical orchestrator", query=query) |
|
|
|
|
|
try: |
|
|
service = get_embedding_service() |
|
|
init_magentic_state(service) |
|
|
except Exception as e: |
|
|
logger.warning( |
|
|
"Embedding service initialization failed, using default state", |
|
|
error=str(e), |
|
|
) |
|
|
init_magentic_state() |
|
|
|
|
|
yield AgentEvent(type="started", message=f"Starting research: {query}") |
|
|
|
|
|
queue: asyncio.Queue[AgentEvent | None] = asyncio.Queue() |
|
|
|
|
|
async def event_callback(event: AgentEvent) -> None: |
|
|
await queue.put(event) |
|
|
|
|
|
task_future = asyncio.create_task(self.middleware.run(query, event_callback)) |
|
|
|
|
|
while not task_future.done(): |
|
|
get_event = asyncio.create_task(queue.get()) |
|
|
done, _ = await asyncio.wait( |
|
|
{task_future, get_event}, return_when=asyncio.FIRST_COMPLETED |
|
|
) |
|
|
|
|
|
if get_event in done: |
|
|
event = get_event.result() |
|
|
if event: |
|
|
yield event |
|
|
else: |
|
|
get_event.cancel() |
|
|
|
|
|
|
|
|
while not queue.empty(): |
|
|
ev = queue.get_nowait() |
|
|
if ev: |
|
|
yield ev |
|
|
|
|
|
try: |
|
|
result, assessment = await task_future |
|
|
|
|
|
assessment_text = assessment.reasoning if assessment else "None" |
|
|
yield AgentEvent( |
|
|
type="complete", |
|
|
message=( |
|
|
f"Research complete.\n\nResult:\n{result}\n\nAssessment:\n{assessment_text}" |
|
|
), |
|
|
data={"assessment": assessment.model_dump() if assessment else None}, |
|
|
) |
|
|
except Exception as e: |
|
|
logger.error("Orchestrator failed", error=str(e)) |
|
|
yield AgentEvent(type="error", message=f"Orchestrator failed: {e}") |
|
|
|
