Replace custom provider abstraction with litellm
Drops the src/providers/ package (base class, AnthropicProvider, OpenAIProvider, factory) in favour of a single src/llm.py that calls litellm.completion() directly. litellm handles provider routing, authentication, and SDK differences for 100+ providers without any code we need to maintain. Changes: - Delete src/providers/ entirely - Add src/llm.py — one complete() function wrapping litellm - src/agent/core.py: EmailAgent takes model: str instead of LLMProvider - src/config.py: ai_provider + api key fields → single ai_model string in litellm format (e.g. "anthropic/claude-opus-4-6") - main.py: remove provider factory wiring; pass config.ai_model to agent - .env.example: simplify AI section, show litellm model string examples - pyproject.toml: replace anthropic + openai deps with litellm>=1.0.0 - uv.lock: regenerated https://claude.ai/code/session_01HaUFs7SaLD5SoiuGCY27Tw
This commit is contained in:
+7
-9
@@ -7,7 +7,7 @@ from datetime import datetime, timezone
|
||||
|
||||
from ..models.conversation import ConversationState, ChatMessage
|
||||
from ..models.registration import BookingDay, RegistrationData
|
||||
from ..providers.base import LLMProvider, LLMMessage
|
||||
from .. import llm
|
||||
from ..knowledge_base.loader import KnowledgeBase
|
||||
from ..storage.json_store import ConversationStore, normalize_email, _diff_registrations
|
||||
from ..notifications.notifier import AdminNotifier
|
||||
@@ -28,12 +28,12 @@ class EmailAgent:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: LLMProvider,
|
||||
model: str,
|
||||
kb: KnowledgeBase,
|
||||
store: ConversationStore,
|
||||
notifier: AdminNotifier,
|
||||
) -> None:
|
||||
self._provider = provider
|
||||
self._model = model
|
||||
self._kb = kb
|
||||
self._store = store
|
||||
self._notifier = notifier
|
||||
@@ -97,11 +97,10 @@ class EmailAgent:
|
||||
def _handle_registration(self, state: ConversationState) -> str:
|
||||
"""Drive the in-progress registration conversation."""
|
||||
system = build_system_prompt(self._kb, state)
|
||||
llm_messages = [LLMMessage(role=m.role, content=m.content) for m in state.messages]
|
||||
|
||||
try:
|
||||
response = self._provider.complete(system=system, messages=llm_messages)
|
||||
parsed = self._parse_llm_response(response.content)
|
||||
content = llm.complete(self._model, system, state.messages)
|
||||
parsed = self._parse_llm_response(content)
|
||||
except Exception:
|
||||
logger.exception("LLM call failed for %s", state.conversation_id)
|
||||
return self._fallback_message(state)
|
||||
@@ -140,11 +139,10 @@ class EmailAgent:
|
||||
def _handle_post_completion(self, state: ConversationState) -> str:
|
||||
"""Handle messages received after a registration is already complete."""
|
||||
system = build_system_prompt(self._kb, state)
|
||||
llm_messages = [LLMMessage(role=m.role, content=m.content) for m in state.messages]
|
||||
|
||||
try:
|
||||
response = self._provider.complete(system=system, messages=llm_messages)
|
||||
parsed = self._parse_llm_response(response.content)
|
||||
content = llm.complete(self._model, system, state.messages)
|
||||
parsed = self._parse_llm_response(content)
|
||||
except Exception:
|
||||
logger.exception("LLM call failed (post-completion) for %s", state.conversation_id)
|
||||
return self._fallback_message(state)
|
||||
|
||||
+4
-9
@@ -13,11 +13,9 @@ except ImportError:
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
# AI Provider
|
||||
ai_provider: str = "anthropic" # "anthropic" or "openai"
|
||||
ai_model: str = ""
|
||||
anthropic_api_key: str = ""
|
||||
openai_api_key: str = ""
|
||||
# AI model — litellm format, e.g. "anthropic/claude-opus-4-6" or "openai/gpt-4o".
|
||||
# The matching API key must be set as an env var (ANTHROPIC_API_KEY, OPENAI_API_KEY, …).
|
||||
ai_model: str = "anthropic/claude-opus-4-6"
|
||||
|
||||
# Email — IMAP (receiving)
|
||||
imap_host: str = ""
|
||||
@@ -48,10 +46,7 @@ class Config:
|
||||
@classmethod
|
||||
def from_env(cls) -> "Config":
|
||||
return cls(
|
||||
ai_provider=os.getenv("AI_PROVIDER", "anthropic"),
|
||||
ai_model=os.getenv("AI_MODEL", ""),
|
||||
anthropic_api_key=os.getenv("ANTHROPIC_API_KEY", ""),
|
||||
openai_api_key=os.getenv("OPENAI_API_KEY", ""),
|
||||
ai_model=os.getenv("AI_MODEL", "anthropic/claude-opus-4-6"),
|
||||
imap_host=os.getenv("IMAP_HOST", ""),
|
||||
imap_port=int(os.getenv("IMAP_PORT", "993")),
|
||||
imap_username=os.getenv("IMAP_USERNAME", ""),
|
||||
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
"""LLM completion via litellm — supports any provider with a single call."""
|
||||
|
||||
import litellm
|
||||
|
||||
|
||||
def complete(model: str, system: str, messages: list) -> str:
|
||||
"""Call any LLM and return the response text.
|
||||
|
||||
Args:
|
||||
model: litellm model string, e.g. "anthropic/claude-opus-4-6" or
|
||||
"openai/gpt-4o". The matching API key must be set as an
|
||||
environment variable (ANTHROPIC_API_KEY, OPENAI_API_KEY, …).
|
||||
system: System prompt text.
|
||||
messages: List of objects with .role and .content attributes.
|
||||
|
||||
Returns:
|
||||
The model's reply as a plain string.
|
||||
"""
|
||||
api_messages = [{"role": "system", "content": system}]
|
||||
api_messages += [{"role": m.role, "content": m.content} for m in messages]
|
||||
response = litellm.completion(model=model, messages=api_messages, max_tokens=2048)
|
||||
return response.choices[0].message.content
|
||||
@@ -1,34 +0,0 @@
|
||||
"""LLM provider registry."""
|
||||
|
||||
from .base import LLMProvider, LLMMessage, LLMResponse
|
||||
from .anthropic_provider import AnthropicProvider
|
||||
from .openai_provider import OpenAIProvider
|
||||
|
||||
__all__ = [
|
||||
"LLMProvider",
|
||||
"LLMMessage",
|
||||
"LLMResponse",
|
||||
"AnthropicProvider",
|
||||
"OpenAIProvider",
|
||||
"create_provider",
|
||||
]
|
||||
|
||||
|
||||
def create_provider(provider: str, api_key: str, model: str = "") -> LLMProvider:
|
||||
"""Instantiate the correct LLMProvider by name.
|
||||
|
||||
Args:
|
||||
provider: "anthropic" or "openai"
|
||||
api_key: API key for the chosen provider.
|
||||
model: Optional model name override.
|
||||
|
||||
Returns:
|
||||
Configured LLMProvider instance.
|
||||
"""
|
||||
if provider == "anthropic":
|
||||
return AnthropicProvider(api_key=api_key, model=model)
|
||||
if provider == "openai":
|
||||
return OpenAIProvider(api_key=api_key, model=model)
|
||||
raise ValueError(
|
||||
f"Unknown AI provider: '{provider}'. Supported values: 'anthropic', 'openai'."
|
||||
)
|
||||
@@ -1,37 +0,0 @@
|
||||
"""Anthropic (Claude) LLM provider."""
|
||||
|
||||
from .base import LLMProvider, LLMMessage, LLMResponse
|
||||
|
||||
|
||||
class AnthropicProvider(LLMProvider):
|
||||
DEFAULT_MODEL = "claude-opus-4-6"
|
||||
|
||||
def __init__(self, api_key: str, model: str = "") -> None:
|
||||
try:
|
||||
import anthropic
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Install the 'anthropic' package to use the Anthropic provider: "
|
||||
"pip install anthropic"
|
||||
) from exc
|
||||
|
||||
self._client = anthropic.Anthropic(api_key=api_key)
|
||||
self._model = model or self.DEFAULT_MODEL
|
||||
|
||||
def complete(self, system: str, messages: list) -> LLMResponse:
|
||||
api_messages = [
|
||||
{"role": m.role, "content": m.content}
|
||||
for m in messages
|
||||
if m.role in ("user", "assistant")
|
||||
]
|
||||
response = self._client.messages.create(
|
||||
model=self._model,
|
||||
max_tokens=2048,
|
||||
system=system,
|
||||
messages=api_messages,
|
||||
)
|
||||
return LLMResponse(content=response.content[0].text)
|
||||
|
||||
@property
|
||||
def model_name(self) -> str:
|
||||
return self._model
|
||||
@@ -1,36 +0,0 @@
|
||||
"""Abstract base class for LLM providers."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMMessage:
|
||||
role: str # "user" or "assistant"
|
||||
content: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMResponse:
|
||||
content: str
|
||||
|
||||
|
||||
class LLMProvider(ABC):
|
||||
"""Uniform interface for any LLM backend."""
|
||||
|
||||
@abstractmethod
|
||||
def complete(self, system: str, messages: list) -> LLMResponse:
|
||||
"""Generate a completion.
|
||||
|
||||
Args:
|
||||
system: System prompt text.
|
||||
messages: List of LLMMessage objects (user/assistant turns).
|
||||
|
||||
Returns:
|
||||
LLMResponse with the model's text output.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def model_name(self) -> str:
|
||||
"""Human-readable model identifier."""
|
||||
@@ -1,37 +0,0 @@
|
||||
"""OpenAI (GPT) LLM provider."""
|
||||
|
||||
from .base import LLMProvider, LLMMessage, LLMResponse
|
||||
|
||||
|
||||
class OpenAIProvider(LLMProvider):
|
||||
DEFAULT_MODEL = "gpt-4o"
|
||||
|
||||
def __init__(self, api_key: str, model: str = "") -> None:
|
||||
try:
|
||||
from openai import OpenAI
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Install the 'openai' package to use the OpenAI provider: "
|
||||
"pip install openai"
|
||||
) from exc
|
||||
|
||||
self._client = OpenAI(api_key=api_key)
|
||||
self._model = model or self.DEFAULT_MODEL
|
||||
|
||||
def complete(self, system: str, messages: list) -> LLMResponse:
|
||||
api_messages = [{"role": "system", "content": system}]
|
||||
api_messages.extend(
|
||||
{"role": m.role, "content": m.content}
|
||||
for m in messages
|
||||
if m.role in ("user", "assistant")
|
||||
)
|
||||
response = self._client.chat.completions.create(
|
||||
model=self._model,
|
||||
messages=api_messages,
|
||||
max_tokens=2048,
|
||||
)
|
||||
return LLMResponse(content=response.choices[0].message.content)
|
||||
|
||||
@property
|
||||
def model_name(self) -> str:
|
||||
return self._model
|
||||
Reference in New Issue
Block a user