Spaces:
Sleeping
Sleeping
| """ | |
| Bring-your-own-key LLM factory. | |
| Users send their provider + API key with each chat request. The key is used to | |
| build a chat model for that one request and is never stored, logged or written | |
| to disk on the server. | |
| """ | |
| import os | |
| import re | |
| from dataclasses import dataclass | |
| from back_end.config import ALLOW_SERVER_KEY_FALLBACK | |
| PROVIDERS = { | |
| "google": {"label": "Google Gemini", "default_model": "gemini-2.5-flash"}, | |
| "openai": {"label": "OpenAI", "default_model": "gpt-5-mini"}, | |
| "anthropic": {"label": "Anthropic Claude", "default_model": "claude-opus-5-5"}, | |
| } | |
| # Loose shape checks: enough to reject garbage / injection attempts early, | |
| # without hard-coding every key format a provider might roll out. | |
| _KEY_PATTERN = re.compile(r"^[A-Za-z0-9_\-\.]{20,300}$") | |
| _MODEL_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_\-\.:/]{0,99}$") | |
| class LLMConfigError(ValueError): | |
| """Raised for invalid user-supplied provider/key/model values.""" | |
| class LLMConfig: | |
| provider: str | |
| api_key: str | |
| model: str | |
| def __repr__(self) -> str: # never let the key end up in logs / tracebacks | |
| return f"LLMConfig(provider={self.provider!r}, model={self.model!r}, api_key='***')" | |
| __str__ = __repr__ | |
| def resolve_llm_config(provider: str | None, api_key: str | None, model: str | None) -> LLMConfig: | |
| provider = (provider or "google").strip().lower() | |
| if provider not in PROVIDERS: | |
| raise LLMConfigError("Unsupported provider.") | |
| api_key = (api_key or "").strip() | |
| if not api_key: | |
| server_key = os.getenv("GOOGLE_API_KEY", "").strip() | |
| if provider == "google" and ALLOW_SERVER_KEY_FALLBACK and server_key: | |
| api_key = server_key | |
| else: | |
| raise LLMConfigError("Please add your API key in Settings (key icon, top right).") | |
| if not _KEY_PATTERN.match(api_key): | |
| raise LLMConfigError("That API key doesn't look valid.") | |
| model = (model or "").strip() or PROVIDERS[provider]["default_model"] | |
| if not _MODEL_PATTERN.match(model): | |
| raise LLMConfigError("Invalid model name.") | |
| return LLMConfig(provider=provider, api_key=api_key, model=model) | |
| def build_chat_model(cfg: LLMConfig, temperature: float = 0.0, max_tokens: int | None = None): | |
| """Create a LangChain chat model for the user's provider.""" | |
| if cfg.provider == "google": | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| kwargs = {"max_output_tokens": max_tokens} if max_tokens else {} | |
| return ChatGoogleGenerativeAI( | |
| model=cfg.model, google_api_key=cfg.api_key, temperature=temperature, max_retries=2, **kwargs | |
| ) | |
| if cfg.provider == "openai": | |
| from langchain_openai import ChatOpenAI | |
| kwargs = {"max_tokens": max_tokens} if max_tokens else {} | |
| # Reasoning models (o-series, gpt-5) reject custom temperatures | |
| if cfg.model.startswith(("gpt-4", "gpt-3.5")): | |
| kwargs["temperature"] = temperature | |
| return ChatOpenAI(model=cfg.model, api_key=cfg.api_key, max_retries=2, **kwargs) | |
| if cfg.provider == "anthropic": | |
| from langchain_anthropic import ChatAnthropic | |
| # Current Claude models reject non-default temperature (HTTP 400), so it is | |
| # never sent. 16k output keeps non-streaming calls under HTTP timeouts. | |
| return ChatAnthropic(model=cfg.model, api_key=cfg.api_key, max_tokens=16000, max_retries=2) | |
| raise LLMConfigError("Unsupported provider.") | |
| def with_structured_output(cfg: LLMConfig, llm, schema): | |
| """Structured output that works on every provider. Claude uses its native | |
| structured outputs: forcing a tool call (LangChain's default) is rejected | |
| by current Claude models.""" | |
| if cfg.provider == "anthropic": | |
| return llm.with_structured_output(schema, method="json_schema") | |
| return llm.with_structured_output(schema) | |
| def strip_thinking_blocks(messages: list) -> list: | |
| """Remove model reasoning blocks before re-sending a transcript under a different | |
| system prompt. Claude only accepts thinking blocks replayed in the exact | |
| conversation that produced them, and the supervisor doesn't need them anyway.""" | |
| cleaned = [] | |
| for m in messages: | |
| content = getattr(m, "content", None) | |
| if isinstance(content, list) and any( | |
| isinstance(b, dict) and b.get("type") in ("thinking", "redacted_thinking") for b in content | |
| ): | |
| kept = [b for b in content if not (isinstance(b, dict) and b.get("type") in ("thinking", "redacted_thinking"))] | |
| m = m.model_copy(update={"content": kept}) | |
| cleaned.append(m) | |
| return cleaned | |
| def friendly_llm_error(exc: Exception) -> str: | |
| """Map provider errors to a message that is safe to show (no key echo, no stack).""" | |
| text = str(exc).lower() | |
| if any(s in text for s in ("401", "403", "api key not valid", "invalid api key", "invalid x-api-key", "permission", "unauthorized", "authentication")): | |
| return "Your API key was rejected by the provider. Check the key and provider in Settings." | |
| if any(s in text for s in ("429", "quota", "rate limit", "resource_exhausted", "resource exhausted")): | |
| return "Your API key hit its rate limit or quota. Wait a bit or use a different key." | |
| if any(s in text for s in ("404", "not found", "does not exist", "model_not_found")): | |
| return "The provider doesn't recognise that model name. Check the model in Settings." | |
| return "The language model request failed. Please try again." | |