409 lines
17 KiB
Python
409 lines
17 KiB
Python
"""FastAPI Application Factory"""
|
|
|
|
import logging
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
from agentkit.core.agent_pool import AgentPool
|
|
from agentkit.llm.gateway import LLMGateway
|
|
from agentkit.llm.providers.anthropic import AnthropicProvider
|
|
from agentkit.llm.providers.openai import OpenAICompatibleProvider
|
|
from agentkit.mcp.manager import MCPManager
|
|
from agentkit.quality.gate import QualityGate
|
|
from agentkit.quality.output import OutputStandardizer
|
|
from agentkit.router.intent import IntentRouter
|
|
from agentkit.skills.base import Skill, SkillConfig
|
|
from agentkit.skills.registry import SkillRegistry
|
|
from agentkit.tools.registry import ToolRegistry
|
|
from agentkit.server.config import ServerConfig
|
|
from agentkit.server.routes import agents, tasks, skills, llm, health, metrics, ws, evolution, memory
|
|
from agentkit.server.middleware import APIKeyAuthMiddleware, RateLimitMiddleware
|
|
from agentkit.server.task_store import create_task_store
|
|
from agentkit.server.runner import BackgroundRunner
|
|
from agentkit.core.logging import setup_structured_logging
|
|
from agentkit.telemetry.setup import setup_telemetry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _build_llm_gateway(config: ServerConfig) -> LLMGateway:
|
|
"""Build LLMGateway from ServerConfig, registering all providers."""
|
|
gateway = LLMGateway(config=config.llm_config)
|
|
|
|
for name, pconf in config.llm_config.providers.items():
|
|
if not pconf.api_key:
|
|
continue # Skip providers without API keys
|
|
try:
|
|
if pconf.type == "anthropic":
|
|
provider = AnthropicProvider(
|
|
api_key=pconf.api_key,
|
|
model=list(pconf.models.keys())[0] if pconf.models else "claude-sonnet-4-20250514",
|
|
max_tokens=pconf.max_tokens,
|
|
base_url=pconf.base_url or "https://api.anthropic.com",
|
|
timeout=pconf.timeout,
|
|
)
|
|
elif pconf.type == "gemini":
|
|
provider = GeminiProvider(
|
|
api_key=pconf.api_key,
|
|
model=list(pconf.models.keys())[0] if pconf.models else "gemini-2.0-flash",
|
|
max_output_tokens=pconf.max_tokens,
|
|
base_url=pconf.base_url or "https://generativelanguage.googleapis.com",
|
|
timeout=pconf.timeout,
|
|
)
|
|
else:
|
|
provider = OpenAICompatibleProvider(
|
|
api_key=pconf.api_key,
|
|
base_url=pconf.base_url,
|
|
)
|
|
gateway.register_provider(name, provider)
|
|
except Exception as e:
|
|
import logging
|
|
logging.getLogger(__name__).warning(f"Failed to register LLM provider '{name}': {e}")
|
|
|
|
return gateway
|
|
|
|
|
|
def _build_skill_registry(config: ServerConfig) -> SkillRegistry:
|
|
"""Build SkillRegistry from ServerConfig, loading all skill configs."""
|
|
registry = SkillRegistry()
|
|
skill_configs = config.load_skill_configs()
|
|
for skill_config in skill_configs:
|
|
skill = Skill(config=skill_config)
|
|
registry.register(skill)
|
|
return registry
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
# Startup
|
|
task_store = app.state.task_store
|
|
await task_store.start_cleanup()
|
|
|
|
# Start config watcher if server_config is available
|
|
server_config = getattr(app.state, "server_config", None)
|
|
if server_config is not None and server_config._config_path:
|
|
server_config.on_change = lambda cfg: _on_config_change(app, cfg)
|
|
server_config.watch_config()
|
|
logger.info("Config hot-reload enabled")
|
|
|
|
# Start MCP servers if configured
|
|
mcp_manager = getattr(app.state, "mcp_manager", None)
|
|
if mcp_manager is not None:
|
|
await mcp_manager.start_all()
|
|
|
|
yield
|
|
|
|
# Shutdown
|
|
# Stop MCP servers
|
|
if mcp_manager is not None:
|
|
await mcp_manager.stop_all()
|
|
|
|
if server_config is not None:
|
|
server_config.stop_watching()
|
|
|
|
await task_store.stop_cleanup()
|
|
|
|
|
|
def _on_config_change(app: FastAPI, config: ServerConfig) -> None:
|
|
"""Handle config change by reloading affected components.
|
|
|
|
Implements graceful rolling update:
|
|
- New tasks use the new configuration
|
|
- In-progress tasks continue with their original configuration
|
|
- Config version is incremented for audit tracking
|
|
"""
|
|
# Increment config version for audit
|
|
current_version = getattr(app.state, "config_version", 0) + 1
|
|
app.state.config_version = current_version
|
|
logger.info(f"Config change detected (v{current_version}), reloading...")
|
|
|
|
# Rebuild LLMGateway if llm config changed
|
|
try:
|
|
new_gateway = _build_llm_gateway(config)
|
|
app.state.llm_gateway = new_gateway
|
|
# Also update the agent pool's gateway reference
|
|
if hasattr(app.state, "agent_pool") and app.state.agent_pool is not None:
|
|
app.state.agent_pool._llm_gateway = new_gateway
|
|
if hasattr(app.state, "intent_router") and app.state.intent_router is not None:
|
|
app.state.intent_router._llm_gateway = new_gateway
|
|
logger.info(f"LLM Gateway reloaded (config v{current_version})")
|
|
except Exception as e:
|
|
logger.error(f"Failed to reload LLM Gateway: {e}")
|
|
|
|
# Reload skills if skill paths changed
|
|
try:
|
|
new_skill_registry = _build_skill_registry(config)
|
|
app.state.skill_registry = new_skill_registry
|
|
if hasattr(app.state, "agent_pool") and app.state.agent_pool is not None:
|
|
app.state.agent_pool._skill_registry = new_skill_registry
|
|
logger.info(f"Skills reloaded (config v{current_version})")
|
|
except Exception as e:
|
|
logger.error(f"Failed to reload skills: {e}")
|
|
|
|
# Update config version on all agents
|
|
if hasattr(app.state, "agent_pool") and app.state.agent_pool is not None:
|
|
for agent in app.state.agent_pool._agents.values():
|
|
if hasattr(agent, "_config_version"):
|
|
agent._config_version = current_version
|
|
|
|
logger.info(f"Config reload complete (v{current_version})")
|
|
|
|
|
|
def create_app(
|
|
llm_gateway: LLMGateway | None = None,
|
|
skill_registry: SkillRegistry | None = None,
|
|
tool_registry: ToolRegistry | None = None,
|
|
api_key: str | None = None,
|
|
rate_limit: int | None = None,
|
|
server_config: ServerConfig | None = None,
|
|
) -> FastAPI:
|
|
"""Create and configure the FastAPI application
|
|
|
|
When called by uvicorn (factory=True), automatically loads ServerConfig
|
|
from AGENTKIT_CONFIG_PATH env var if server_config is not provided.
|
|
"""
|
|
# Auto-load config from env var if not provided (uvicorn factory mode)
|
|
if server_config is None:
|
|
config_path = os.environ.get("AGENTKIT_CONFIG_PATH")
|
|
if config_path and os.path.exists(config_path):
|
|
server_config = ServerConfig.from_yaml(config_path)
|
|
app = FastAPI(title="AgentKit Server", version="2.0.0", lifespan=lifespan)
|
|
|
|
# Initialize structured logging
|
|
setup_structured_logging()
|
|
|
|
# Initialize OpenTelemetry (no-op if not installed or not configured)
|
|
if server_config:
|
|
setup_telemetry(app, server_config.telemetry)
|
|
|
|
# Resolve effective API key and rate limit
|
|
effective_api_key = api_key
|
|
effective_rate_limit = rate_limit
|
|
if server_config:
|
|
if effective_api_key is None:
|
|
effective_api_key = server_config.api_key
|
|
if effective_rate_limit is None:
|
|
effective_rate_limit = server_config.rate_limit
|
|
|
|
# CORS 配置
|
|
cors_origins = ["*"]
|
|
if server_config:
|
|
cors_origins = server_config.cors_origins
|
|
if cors_origins == ["*"]:
|
|
import logging
|
|
logging.getLogger(__name__).warning(
|
|
"CORS allows all origins (allow_origins=['*']). "
|
|
"Set server.cors_origins in agentkit.yaml for production."
|
|
)
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=cors_origins,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# Auth middleware
|
|
app.add_middleware(APIKeyAuthMiddleware, api_key=effective_api_key)
|
|
|
|
# Rate limiting middleware
|
|
if effective_rate_limit is not None:
|
|
os.environ["AGENTKIT_RATE_LIMIT_PER_MINUTE"] = str(effective_rate_limit)
|
|
app.add_middleware(RateLimitMiddleware)
|
|
|
|
# Build LLM Gateway from config if not provided
|
|
if llm_gateway is None and server_config:
|
|
llm_gateway = _build_llm_gateway(server_config)
|
|
|
|
# Build Skill Registry from config if not provided
|
|
if skill_registry is None and server_config:
|
|
skill_registry = _build_skill_registry(server_config)
|
|
|
|
# Initialize shared state
|
|
app.state.llm_gateway = llm_gateway or LLMGateway()
|
|
app.state.skill_registry = skill_registry or SkillRegistry()
|
|
app.state.tool_registry = tool_registry or ToolRegistry()
|
|
# Initialize MCPManager if MCP servers are configured
|
|
if server_config and server_config.mcp_servers:
|
|
mcp_manager = MCPManager(
|
|
configs=server_config.mcp_servers,
|
|
tool_registry=app.state.tool_registry,
|
|
)
|
|
app.state.mcp_manager = mcp_manager
|
|
else:
|
|
app.state.mcp_manager = None
|
|
# Initialize compressor if compression is configured
|
|
from agentkit.core.compressor import create_compressor
|
|
compressor = create_compressor(server_config.compression) if server_config else None
|
|
app.state.compressor = compressor
|
|
# Register headroom_retrieve tool if HeadroomCompressor is active
|
|
if compressor is not None:
|
|
try:
|
|
from agentkit.core.headroom_compressor import HeadroomCompressor
|
|
if isinstance(compressor, HeadroomCompressor) and compressor.is_available():
|
|
from agentkit.tools.headroom_retrieve import HeadroomRetrieveTool
|
|
retrieve_tool = HeadroomRetrieveTool(compressor=compressor)
|
|
app.state.tool_registry.register(retrieve_tool)
|
|
logger.info("HeadroomRetrieveTool registered (CCR retrieval enabled)")
|
|
except ImportError:
|
|
pass
|
|
app.state.agent_pool = AgentPool(
|
|
llm_gateway=app.state.llm_gateway,
|
|
skill_registry=app.state.skill_registry,
|
|
tool_registry=app.state.tool_registry,
|
|
compressor=compressor,
|
|
)
|
|
app.state.intent_router = IntentRouter(llm_gateway=app.state.llm_gateway)
|
|
app.state.quality_gate = QualityGate()
|
|
app.state.output_standardizer = OutputStandardizer()
|
|
# Initialize task store from config
|
|
ts_config = server_config.task_store if server_config else {}
|
|
# Merge CLI overrides from AGENTKIT_TASK_STORE env var
|
|
ts_env = os.environ.get("AGENTKIT_TASK_STORE")
|
|
if ts_env:
|
|
import json as _json
|
|
try:
|
|
ts_config = {**ts_config, **_json.loads(ts_env)}
|
|
except Exception:
|
|
pass
|
|
task_store = create_task_store(
|
|
backend=ts_config.get("backend", "memory"),
|
|
redis_url=ts_config.get("redis_url", "redis://localhost:6379/0"),
|
|
ttl_seconds=ts_config.get("ttl_seconds", 3600),
|
|
max_records=ts_config.get("max_records", 10000),
|
|
)
|
|
app.state.task_store = task_store
|
|
app.state.runner = BackgroundRunner(task_store=app.state.task_store)
|
|
app.state.server_config = server_config
|
|
app.state.api_key = effective_api_key
|
|
|
|
# Initialize evolution store if configured
|
|
if server_config and hasattr(server_config, 'evolution') and server_config.evolution:
|
|
try:
|
|
from agentkit.evolution.evolution_store import create_evolution_store
|
|
evo_conf = server_config.evolution
|
|
app.state.evolution_store = create_evolution_store(
|
|
backend=evo_conf.get("backend", "memory"),
|
|
db_path=evo_conf.get("db_path", "~/.agentkit/evolution.db"),
|
|
)
|
|
except Exception as e:
|
|
import logging
|
|
logging.getLogger(__name__).warning(f"Failed to initialize evolution store: {e}")
|
|
app.state.evolution_store = None
|
|
else:
|
|
app.state.evolution_store = None
|
|
|
|
# Initialize memory components if configured
|
|
if server_config and hasattr(server_config, 'memory') and server_config.memory:
|
|
try:
|
|
from agentkit.memory.retriever import MemoryRetriever
|
|
from agentkit.memory.working import WorkingMemory
|
|
from agentkit.memory.semantic import SemanticMemory
|
|
from agentkit.memory.http_rag import HttpRAGService
|
|
|
|
working = None
|
|
episodic = None
|
|
semantic = None
|
|
|
|
if server_config.memory.get("working", {}).get("enabled"):
|
|
import redis.asyncio as aioredis
|
|
redis_url = server_config.memory["working"].get("redis_url", "redis://localhost:6379")
|
|
redis_client = aioredis.from_url(redis_url, decode_responses=True)
|
|
working = WorkingMemory(redis=redis_client)
|
|
|
|
if server_config.memory.get("semantic", {}).get("enabled"):
|
|
sem_conf = server_config.memory["semantic"]
|
|
rag_service = HttpRAGService(
|
|
base_url=sem_conf["base_url"],
|
|
api_key=sem_conf.get("api_key"),
|
|
knowledge_base_ids=sem_conf.get("knowledge_base_ids", []),
|
|
timeout=sem_conf.get("timeout", 30),
|
|
)
|
|
semantic = SemanticMemory(
|
|
rag_service=rag_service,
|
|
knowledge_base_ids=sem_conf.get("knowledge_base_ids", []),
|
|
search_mode=sem_conf.get("search_mode", "standard"),
|
|
use_rerank=sem_conf.get("use_rerank", True),
|
|
use_compression=sem_conf.get("use_compression", False),
|
|
kb_weights=sem_conf.get("kb_weights"),
|
|
)
|
|
|
|
if server_config.memory.get("episodic", {}).get("enabled"):
|
|
try:
|
|
from agentkit.memory.episodic import EpisodicMemory
|
|
from agentkit.memory.embedder import OpenAIEmbedder, EmbeddingCache
|
|
from agentkit.memory.models import EpisodeModel, create_episodic_session_factory
|
|
|
|
epi_conf = server_config.memory["episodic"]
|
|
embedder = None
|
|
if epi_conf.get("embedder_api_key") or os.environ.get("OPENAI_API_KEY"):
|
|
cache = EmbeddingCache(
|
|
max_size=epi_conf.get("cache_max_size", 1000),
|
|
ttl=epi_conf.get("cache_ttl", 3600),
|
|
)
|
|
embedder = OpenAIEmbedder(
|
|
api_key=epi_conf.get("embedder_api_key"),
|
|
model=epi_conf.get("embedder_model", "text-embedding-3-small"),
|
|
base_url=epi_conf.get("embedder_base_url"),
|
|
cache=cache,
|
|
)
|
|
# Resolve session_factory and model from database_url if configured
|
|
epi_session_factory = None
|
|
epi_model = None
|
|
database_url = epi_conf.get("database_url") or os.environ.get("DATABASE_URL")
|
|
if database_url:
|
|
try:
|
|
epi_session_factory = create_episodic_session_factory(database_url)
|
|
epi_model = EpisodeModel
|
|
except Exception as db_err:
|
|
import logging as _log
|
|
_log.getLogger(__name__).warning(
|
|
f"Failed to create episodic DB session: {db_err}"
|
|
)
|
|
|
|
episodic = EpisodicMemory(
|
|
session_factory=epi_session_factory,
|
|
episodic_model=epi_model,
|
|
embedder=embedder,
|
|
decay_rate=epi_conf.get("decay_rate", 0.01),
|
|
alpha=epi_conf.get("alpha", 0.7),
|
|
retrieve_limit=epi_conf.get("retrieve_limit", 200),
|
|
pgvector_enabled=epi_conf.get("pgvector_enabled", True),
|
|
table_name=epi_conf.get("table_name", "episodic_memories"),
|
|
)
|
|
except Exception as e:
|
|
import logging
|
|
logging.getLogger(__name__).warning(f"Failed to initialize episodic memory: {e}")
|
|
|
|
memory_retriever = MemoryRetriever(
|
|
working_memory=working,
|
|
episodic_memory=episodic,
|
|
semantic_memory=semantic,
|
|
)
|
|
app.state.memory_retriever = memory_retriever
|
|
|
|
# Auto-register retrieve_knowledge tool if semantic memory is configured
|
|
if memory_retriever:
|
|
retrieve_tool = memory_retriever.create_retrieve_tool()
|
|
if retrieve_tool:
|
|
app.state.retrieve_knowledge_tool = retrieve_tool
|
|
except Exception as e:
|
|
import logging
|
|
logging.getLogger(__name__).warning(f"Failed to initialize memory components: {e}")
|
|
app.state.memory_retriever = None
|
|
|
|
# Include routes
|
|
app.include_router(agents.router, prefix="/api/v1")
|
|
app.include_router(tasks.router, prefix="/api/v1")
|
|
app.include_router(skills.router, prefix="/api/v1")
|
|
app.include_router(llm.router, prefix="/api/v1")
|
|
app.include_router(health.router, prefix="/api/v1")
|
|
app.include_router(metrics.router, prefix="/api/v1")
|
|
app.include_router(ws.router, prefix="/api/v1")
|
|
app.include_router(evolution.router, prefix="/api/v1")
|
|
app.include_router(memory.router, prefix="/api/v1")
|
|
|
|
return app
|