64 lines
2.8 KiB
Python
64 lines
2.8 KiB
Python
import asyncio
|
|
from typing import Callable
|
|
from agno.agent import RunEvent, RunOutputEvent
|
|
from agno.team import Team, TeamRunEvent, TeamRunOutputEvent
|
|
from agno.tools.reasoning import ReasoningTools
|
|
from app.agents.prompts import *
|
|
from app.configs import AppConfig, AppModel
|
|
from app.api.tools import *
|
|
|
|
|
|
|
|
# Define a unified event type for team and agent events
|
|
AppEvent = TeamRunOutputEvent | RunOutputEvent
|
|
|
|
|
|
class AppTeam:
|
|
def __init__(self, configs: AppConfig, team_models: AppModel, team_leader: AppModel | None = None):
|
|
self.configs = configs
|
|
self.team_models = team_models
|
|
self.team_leader = team_leader or team_models
|
|
self.listeners: dict[str, Callable[[AppEvent], None]] = {}
|
|
|
|
def add_listener(self, event: RunEvent | TeamRunEvent, listener: Callable[[AppEvent], None]) -> None:
|
|
self.listeners[event.value] = listener
|
|
|
|
def run_team(self, query: str) -> str:
|
|
return asyncio.run(self.run_team_async(query))
|
|
|
|
async def run_team_async(self, query: str) -> str:
|
|
team = AppTeam.create_team_with(self.configs, self.team_models, self.team_leader)
|
|
result = "No output from team"
|
|
|
|
async for run_event in team.arun(query, stream=True, stream_intermediate_steps=True): # type: ignore
|
|
if run_event.event in self.listeners:
|
|
self.listeners[run_event.event](run_event)
|
|
if run_event.event in [TeamRunEvent.run_completed]:
|
|
if isinstance(run_event.content, str):
|
|
result = run_event.content
|
|
thinking = result.rfind("</think>")
|
|
if thinking != -1: result = result[thinking:]
|
|
|
|
return result
|
|
|
|
@classmethod
|
|
def create_team_with(cls, configs: AppConfig, team_model: AppModel, team_leader: AppModel | None = None) -> Team:
|
|
|
|
market_tool = MarketAPIsTool(currency=configs.api.currency)
|
|
market_tool.handler.set_retries(configs.api.retry_attempts, configs.api.retry_delay_seconds)
|
|
news_tool = NewsAPIsTool()
|
|
news_tool.handler.set_retries(configs.api.retry_attempts, configs.api.retry_delay_seconds)
|
|
social_tool = SocialAPIsTool()
|
|
social_tool.handler.set_retries(configs.api.retry_attempts, configs.api.retry_delay_seconds)
|
|
|
|
market_agent = team_model.get_agent(instructions=MARKET_INSTRUCTIONS, name="MarketAgent", tools=[market_tool])
|
|
news_agent = team_model.get_agent(instructions=NEWS_INSTRUCTIONS, name="NewsAgent", tools=[news_tool])
|
|
social_agent = team_model.get_agent(instructions=SOCIAL_INSTRUCTIONS, name="SocialAgent", tools=[social_tool])
|
|
|
|
team_leader = team_leader or team_model
|
|
return Team(
|
|
model=team_leader.get_model(COORDINATOR_INSTRUCTIONS),
|
|
name="CryptoAnalysisTeam",
|
|
tools=[ReasoningTools()],
|
|
members=[market_agent, news_agent, social_agent],
|
|
) |