diff --git a/docker-compose.yaml b/docker-compose.yaml index af9261f..41cd688 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -28,9 +28,51 @@ services: # GitHub MCP (Remote) - GITHUB_PERSONAL_ACCESS_TOKEN=${GITHUB_PERSONAL_ACCESS_TOKEN} - GITHUB_MCP_URL=https://api.githubcopilot.com/mcp/ + # LiteLLM + - LITELLM_URL=http://litellm:4000 + - LITELLM_MASTER_KEY=sk-sre-agent-local depends_on: slack: condition: service_started + litellm: + condition: service_healthy + restart: unless-stopped + + # LiteLLM Proxy (cost tracking + model routing) + litellm: + image: ghcr.io/berriai/litellm:main-latest + command: ["--config", "/app/litellm_config.yaml"] + volumes: + - ./litellm_config.yaml:/app/litellm_config.yaml:ro + environment: + - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY} + - DATABASE_URL=postgresql://litellm:litellm@postgres:5432/litellm + ports: + - "4000:4000" + healthcheck: + test: ["CMD-SHELL", "curl -f http://localhost:4000/health/liveliness || exit 1"] + interval: 10s + timeout: 5s + retries: 5 + depends_on: + postgres: + condition: service_healthy + restart: unless-stopped + + # PostgreSQL (for LiteLLM cost tracking) + postgres: + image: postgres:16 + environment: + - POSTGRES_USER=litellm + - POSTGRES_PASSWORD=litellm + - POSTGRES_DB=litellm + volumes: + - postgres_data:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U litellm -d litellm"] + interval: 5s + timeout: 5s + retries: 5 restart: unless-stopped # Slack MCP Server (korotovsky/slack-mcp-server) @@ -44,3 +86,6 @@ services: restart: unless-stopped ports: - "13080:13080" + +volumes: + postgres_data: diff --git a/litellm_config.yaml b/litellm_config.yaml new file mode 100644 index 0000000..ca2bead --- /dev/null +++ b/litellm_config.yaml @@ -0,0 +1,9 @@ +model_list: + - model_name: claude-sonnet-4-5 + litellm_params: + model: anthropic/claude-sonnet-4-5-20250929 + api_key: os.environ/ANTHROPIC_API_KEY + +general_settings: + master_key: sk-sre-agent-local + database_url: "postgresql://litellm:litellm@postgres:5432/litellm" diff --git a/src/sre_agent/core/agent.py b/src/sre_agent/core/agent.py index caa703c..473a1cb 100644 --- a/src/sre_agent/core/agent.py +++ b/src/sre_agent/core/agent.py @@ -1,6 +1,8 @@ """SRE Agent using pydantic-ai.""" from pydantic_ai import Agent +from pydantic_ai.models.anthropic import AnthropicModel +from pydantic_ai.providers.anthropic import AnthropicProvider from sre_agent.core.models import ErrorDiagnosis from sre_agent.core.prompts import SYSTEM_PROMPT, build_diagnosis_prompt @@ -27,8 +29,16 @@ def create_sre_agent(config: AgentSettings) -> Agent[None, ErrorDiagnosis]: create_slack_mcp_toolset(config), ] - return Agent( + model = AnthropicModel( config.model, + provider=AnthropicProvider( + api_key=config.litellm.master_key, + base_url=f"{config.litellm.url}/anthropic/", + ), + ) + + return Agent( + model, system_prompt=SYSTEM_PROMPT, output_type=ErrorDiagnosis, toolsets=toolsets, diff --git a/src/sre_agent/core/settings.py b/src/sre_agent/core/settings.py index d2d2cfc..cc610a9 100644 --- a/src/sre_agent/core/settings.py +++ b/src/sre_agent/core/settings.py @@ -50,6 +50,19 @@ class SlackSettings(BaseSettings): mcp_url: str = Field(description="URL of Slack MCP server (SSE)") +class LiteLLMSettings(BaseSettings): + """LiteLLM proxy configuration.""" + + model_config = SettingsConfigDict( + env_prefix="LITELLM_", + env_file=ENV_FILE_PATH, + extra="ignore", + ) + + url: str = Field(default="http://litellm:4000", description="LiteLLM proxy URL") + master_key: str = Field(description="LiteLLM master key") + + class AgentSettings(BaseSettings): """Main agent configuration.""" @@ -67,6 +80,7 @@ class AgentSettings(BaseSettings): aws: AWSSettings github: GitHubSettings slack: SlackSettings + litellm: LiteLLMSettings def get_settings() -> AgentSettings: @@ -84,4 +98,5 @@ def get_settings() -> AgentSettings: aws=AWSSettings(), github=GitHubSettings(), # type: ignore[call-arg] slack=SlackSettings(), # type: ignore[call-arg] + litellm=LiteLLMSettings(), # type: ignore[call-arg] )