File size: 3,744 Bytes
22dcdfd |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
from dataclasses import dataclass
from langgraph.graph.state import CompiledStateGraph
from langgraph.pregel import Pregel
from agents.bg_task_agent.bg_task_agent import bg_task_agent
from agents.chatbot import chatbot
from agents.command_agent import command_agent
from agents.github_mcp_agent.github_mcp_agent import github_mcp_agent
from agents.interrupt_agent import interrupt_agent
from agents.knowledge_base_agent import kb_agent
from agents.langgraph_supervisor_agent import langgraph_supervisor_agent
from agents.langgraph_supervisor_hierarchy_agent import langgraph_supervisor_hierarchy_agent
from agents.portfolio_agent.portfolio_agent import portfolio_agent
from agents.lazy_agent import LazyLoadingAgent
from agents.rag_assistant import rag_assistant
from agents.research_assistant import research_assistant
from schema import AgentInfo
DEFAULT_AGENT = "portfolio-agent"
# Type alias to handle LangGraph's different agent patterns
# - @entrypoint functions return Pregel
# - StateGraph().compile() returns CompiledStateGraph
AgentGraph = CompiledStateGraph | Pregel # What get_agent() returns (always loaded)
AgentGraphLike = CompiledStateGraph | Pregel | LazyLoadingAgent # What can be stored in registry
@dataclass
class Agent:
description: str
graph_like: AgentGraphLike
agents: dict[str, Agent] = {
"chatbot": Agent(description="A simple chatbot.", graph_like=chatbot),
"research-assistant": Agent(
description="A research assistant with web search and calculator.",
graph_like=research_assistant,
),
"rag-assistant": Agent(
description="A RAG assistant with access to information in a database.",
graph_like=rag_assistant,
),
"portfolio-agent": Agent(
description="A portfolio assistant with access to information in a database.",
graph_like=portfolio_agent,
),
"command-agent": Agent(description="A command agent.", graph_like=command_agent),
"bg-task-agent": Agent(description="A background task agent.", graph_like=bg_task_agent),
"langgraph-supervisor-agent": Agent(
description="A langgraph supervisor agent", graph_like=langgraph_supervisor_agent
),
"langgraph-supervisor-hierarchy-agent": Agent(
description="A langgraph supervisor agent with a nested hierarchy of agents",
graph_like=langgraph_supervisor_hierarchy_agent,
),
"interrupt-agent": Agent(
description="An agent the uses interrupts.", graph_like=interrupt_agent
),
"knowledge-base-agent": Agent(
description="A retrieval-augmented generation agent using Amazon Bedrock Knowledge Base",
graph_like=kb_agent,
),
"github-mcp-agent": Agent(
description="A GitHub agent with MCP tools for repository management and development workflows.",
graph_like=github_mcp_agent,
),
}
async def load_agent(agent_id: str) -> None:
"""Load lazy agents if needed."""
graph_like = agents[agent_id].graph_like
if isinstance(graph_like, LazyLoadingAgent):
await graph_like.load()
def get_agent(agent_id: str) -> AgentGraph:
"""Get an agent graph, loading lazy agents if needed."""
agent_graph = agents[agent_id].graph_like
# If it's a lazy loading agent, ensure it's loaded and return its graph
if isinstance(agent_graph, LazyLoadingAgent):
if not agent_graph._loaded:
raise RuntimeError(f"Agent {agent_id} not loaded. Call load() first.")
return agent_graph.get_graph()
# Otherwise return the graph directly
return agent_graph
def get_all_agent_info() -> list[AgentInfo]:
return [
AgentInfo(key=agent_id, description=agent.description) for agent_id, agent in agents.items()
]
|