ankn
c1c8f31539
Switch Gemini from deprecated google-generativeai to google-genai. Fix Mistral import path to mistralai.client.Mistral (v2.x layout). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
"""LLM provider abstraction. Each provider wraps one vendor SDK."""
|
|
import os
|
|
|
|
from rich.console import Console
|
|
|
|
_console = Console()
|
|
|
|
|
|
class Provider:
|
|
name: str
|
|
|
|
def chat(self, messages: list[dict], system: str, max_tokens: int = 4096) -> str:
|
|
raise NotImplementedError
|
|
|
|
|
|
class AnthropicProvider(Provider):
|
|
name = "Claude"
|
|
|
|
def __init__(self):
|
|
import anthropic
|
|
self._client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
|
|
|
|
def chat(self, messages: list[dict], system: str, max_tokens: int = 4096) -> str:
|
|
response = self._client.messages.create(
|
|
model="claude-sonnet-4-6",
|
|
max_tokens=max_tokens,
|
|
system=system,
|
|
messages=messages,
|
|
)
|
|
return response.content[0].text
|
|
|
|
|
|
class OpenAIProvider(Provider):
|
|
name = "ChatGPT"
|
|
|
|
def __init__(self):
|
|
from openai import OpenAI
|
|
self._client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
|
|
|
|
def chat(self, messages: list[dict], system: str, max_tokens: int = 4096) -> str:
|
|
oai_messages = [{"role": "system", "content": system}] + messages
|
|
response = self._client.chat.completions.create(
|
|
model="gpt-4o",
|
|
max_tokens=max_tokens,
|
|
messages=oai_messages,
|
|
)
|
|
return response.choices[0].message.content
|
|
|
|
|
|
class GeminiProvider(Provider):
|
|
name = "Gemini"
|
|
|
|
def __init__(self):
|
|
from google import genai
|
|
from google.genai import types as genai_types
|
|
self._client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
|
|
self._types = genai_types
|
|
|
|
def chat(self, messages: list[dict], system: str, max_tokens: int = 4096) -> str:
|
|
# Convert role "assistant" → "model" for Gemini
|
|
contents = [
|
|
{"role": "model" if m["role"] == "assistant" else "user", "parts": [{"text": m["content"]}]}
|
|
for m in messages
|
|
]
|
|
response = self._client.models.generate_content(
|
|
model="gemini-2.0-flash",
|
|
contents=contents,
|
|
config=self._types.GenerateContentConfig(
|
|
system_instruction=system,
|
|
max_output_tokens=max_tokens,
|
|
),
|
|
)
|
|
return response.text
|
|
|
|
|
|
class MistralProvider(Provider):
|
|
name = "Mistral"
|
|
|
|
def __init__(self):
|
|
from mistralai.client import Mistral
|
|
self._client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
|
|
|
|
def chat(self, messages: list[dict], system: str, max_tokens: int = 4096) -> str:
|
|
mistral_messages = [{"role": "system", "content": system}] + messages
|
|
response = self._client.chat.complete(
|
|
model="mistral-large-latest",
|
|
max_tokens=max_tokens,
|
|
messages=mistral_messages,
|
|
)
|
|
return response.choices[0].message.content
|
|
|
|
|
|
_REGISTRY = [
|
|
("ANTHROPIC_API_KEY", AnthropicProvider),
|
|
("OPENAI_API_KEY", OpenAIProvider),
|
|
("GOOGLE_API_KEY", GeminiProvider),
|
|
("MISTRAL_API_KEY", MistralProvider),
|
|
]
|
|
|
|
|
|
def get_configured() -> list[Provider]:
|
|
"""Return instantiated providers for every API key that is set. Skips on init error."""
|
|
result = []
|
|
for env_var, cls in _REGISTRY:
|
|
if os.environ.get(env_var):
|
|
try:
|
|
result.append(cls())
|
|
except Exception as e:
|
|
_console.print(f"[yellow][Warnung] {cls.__name__} konnte nicht initialisiert werden: {e}[/yellow]")
|
|
return result
|