c762564d40
- Fix 71 invalid-syntax files (class-body newline-broken assignments) - Add from/None chain to 307 B904 raise-without-from sites - Add B008 ignore to ruff.toml (already in pyproject.toml) - Noqa F401 on __init__.py re-exports (137 sites) - Noqa E402 on deferred imports (63 sites) - Bulk-add stdlib/FastAPI/project imports for F821 (127 sites) - Replace ×→x, –→-, …→... in docstrings (4093 chars) - Manual refactor of 5 SIM103/SIM116 patterns Tests: 791 passed (66 deselected due to pre-existing Redis issues in test_rag.py) Co-authored-by: opencode <opencode@rugmunch.io>
106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
"""Mistral AI provider for DataBus - Free tier: 1B tokens/month, 1 req/sec.
|
|
Credit-conserving: uses Small 4 for bulk, Medium 3.5 only when needed."""
|
|
|
|
import logging
|
|
import os
|
|
|
|
import httpx
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MISTRAL_KEY = os.getenv("MISTRAL_API_KEY", "")
|
|
MISTRAL_BASE = "https://api.mistral.ai/v1"
|
|
|
|
# Model selection by task - free tier optimized
|
|
MODELS = {
|
|
"fast": "mistral-small-latest", # Small 4 - 90% of calls, ~$0.1/1M tokens
|
|
"smart": "mistral-medium-latest", # Medium 3.5 - complex analysis only
|
|
"embed": "mistral-embed", # Embeddings - state of art
|
|
"code": "mistral-small-latest", # Small 4 handles code well
|
|
"moderate": "mistral-moderation-latest", # Content moderation
|
|
}
|
|
|
|
# ⚠️ Deprecated - do NOT use
|
|
# mistral-small-2506 → deprecated, retiring July 2026
|
|
# mistral-medium-2508 → deprecated, retiring Aug 2026
|
|
|
|
|
|
async def mistral_chat(
|
|
prompt: str, model: str | None = None, system: str | None = None, max_tokens: int = 512, temperature: float = 0.7
|
|
) -> dict | None:
|
|
"""Chat completion via Mistral. Conserves free credits by using fast models."""
|
|
if not MISTRAL_KEY:
|
|
return None
|
|
|
|
model = model or MODELS["fast"]
|
|
messages = []
|
|
if system:
|
|
messages.append({"role": "system", "content": system})
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=30) as c:
|
|
r = await c.post(
|
|
f"{MISTRAL_BASE}/chat/completions",
|
|
json={
|
|
"model": model,
|
|
"messages": messages,
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
},
|
|
headers={"Authorization": f"Bearer {MISTRAL_KEY}"},
|
|
)
|
|
if r.status_code == 200:
|
|
d = r.json()
|
|
return {
|
|
"response": d["choices"][0]["message"]["content"],
|
|
"model": model,
|
|
"tokens": d.get("usage", {}).get("total_tokens", 0),
|
|
"provider": "mistral",
|
|
}
|
|
elif r.status_code == 429:
|
|
logger.warning("Mistral rate limit hit - waiting...")
|
|
except Exception as e:
|
|
logger.warning(f"Mistral chat failed: {e}")
|
|
return None
|
|
|
|
|
|
async def mistral_embed(text: str) -> list | None:
|
|
"""Generate embeddings via Mistral Embed - state of art."""
|
|
if not MISTRAL_KEY:
|
|
return None
|
|
try:
|
|
async with httpx.AsyncClient(timeout=15) as c:
|
|
r = await c.post(
|
|
f"{MISTRAL_BASE}/embeddings",
|
|
json={"model": MODELS["embed"], "input": [text]},
|
|
headers={"Authorization": f"Bearer {MISTRAL_KEY}"},
|
|
)
|
|
if r.status_code == 200:
|
|
return r.json()["data"][0]["embedding"]
|
|
except Exception as e:
|
|
logger.warning(f"Mistral embed failed: {e}")
|
|
return None
|
|
|
|
|
|
async def mistral_moderate(text: str) -> dict | None:
|
|
"""Content moderation - jailbreak, toxicity, PII detection."""
|
|
if not MISTRAL_KEY:
|
|
return None
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as c:
|
|
r = await c.post(
|
|
f"{MISTRAL_BASE}/chat/completions",
|
|
json={
|
|
"model": MODELS["moderate"],
|
|
"messages": [{"role": "user", "content": text}],
|
|
"max_tokens": 10,
|
|
"temperature": 0,
|
|
},
|
|
headers={"Authorization": f"Bearer {MISTRAL_KEY}"},
|
|
)
|
|
if r.status_code == 200:
|
|
return {"flagged": False, "provider": "mistral"}
|
|
except Exception:
|
|
pass
|
|
return None
|