Commit 9f1fcd5a authored by Kantz's avatar Kantz
Browse files

Limit hinzugefügt

parent cbc052c2
...@@ -2,6 +2,8 @@ MATHPIX_APP_ID="" ...@@ -2,6 +2,8 @@ MATHPIX_APP_ID=""
MATHPIX_APP_KEY="" MATHPIX_APP_KEY=""
POSTGRES_URL="" POSTGRES_URL=""
DAILY_LLM_CALL_LIMIT="100"
DAILY_LLM_TOKEN_LIMIT="50000"
FRONTEND_URL="http://localhost:5173" FRONTEND_URL="http://localhost:5173"
...@@ -28,4 +30,4 @@ OLLAMA_TIMEOUT="60" ...@@ -28,4 +30,4 @@ OLLAMA_TIMEOUT="60"
MISTRAL_CHAT_MODEL="mistral-large-3-675b-instruct-2512" MISTRAL_CHAT_MODEL="mistral-large-3-675b-instruct-2512"
MISTRAL_API_KEY="" MISTRAL_API_KEY=""
MISTRAL_CHAT_TIMEOUT="60" MISTRAL_CHAT_TIMEOUT="60"
MISTRAL_CHAT_TEMPERATURE="0.2" MISTRAL_CHAT_TEMPERATURE="0.2"
\ No newline at end of file
...@@ -9,6 +9,7 @@ import psycopg ...@@ -9,6 +9,7 @@ import psycopg
from fastapi import APIRouter from fastapi import APIRouter
import app.config as config import app.config as config
from app.deterministic_services import llm_quota
router = APIRouter() router = APIRouter()
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
...@@ -98,12 +99,36 @@ def _check_postgres() -> dict: ...@@ -98,12 +99,36 @@ def _check_postgres() -> dict:
return {"status": "error", "detail": str(exc)} return {"status": "error", "detail": str(exc)}
def _check_llm_quota() -> dict:
try:
pg_url = config.get_postgres_url()
except ValueError:
return {"status": "missing_config"}
try:
today_usage = llm_quota.get_today_usage(pg_url)
limits = config.get_llm_quota_settings()
return {
"status": "ok",
"usage_date": today_usage.usage_date.isoformat(),
"call_count": today_usage.call_count,
"token_count": today_usage.token_count,
"limits": {
"daily_call_limit": limits.daily_call_limit,
"daily_token_limit": limits.daily_token_limit,
},
}
except Exception as exc:
return {"status": "error", "detail": str(exc)}
@router.get("/api/health") @router.get("/api/health")
def health() -> Dict[str, Any]: def health() -> Dict[str, Any]:
services = { services = {
"ollama": _check_ollama(), "ollama": _check_ollama(),
"openai": _check_openai(), "openai": _check_openai(),
"postgres": _check_postgres(), "postgres": _check_postgres(),
"llm_quota": _check_llm_quota(),
} }
required_statuses = {"ok"} required_statuses = {"ok"}
overall = "ok" if all( overall = "ok" if all(
......
...@@ -92,6 +92,12 @@ class MathpixSettings: ...@@ -92,6 +92,12 @@ class MathpixSettings:
app_key: str app_key: str
@dataclass(frozen=True)
class LLMQuotaSettings:
daily_call_limit: int | None
daily_token_limit: int | None
def _read_float(value: str | None) -> float | None: def _read_float(value: str | None) -> float | None:
if value is None or value == "": if value is None or value == "":
return None return None
...@@ -101,6 +107,15 @@ def _read_float(value: str | None) -> float | None: ...@@ -101,6 +107,15 @@ def _read_float(value: str | None) -> float | None:
return None return None
def _read_int(value: str | None) -> int | None:
if value is None or value == "":
return None
try:
return int(value)
except ValueError:
return None
def get_mathpix_settings() -> MathpixSettings: def get_mathpix_settings() -> MathpixSettings:
return MathpixSettings( return MathpixSettings(
app_id=os.getenv("MATHPIX_APP_ID"), app_id=os.getenv("MATHPIX_APP_ID"),
...@@ -172,3 +187,10 @@ def get_postgres_url() -> str: ...@@ -172,3 +187,10 @@ def get_postgres_url() -> str:
if not pg_url: if not pg_url:
raise ValueError("Missing POSTGRES_URL") raise ValueError("Missing POSTGRES_URL")
return pg_url return pg_url
def get_llm_quota_settings() -> LLMQuotaSettings:
return LLMQuotaSettings(
daily_call_limit=_read_int(os.getenv("DAILY_LLM_CALL_LIMIT")),
daily_token_limit=_read_int(os.getenv("DAILY_LLM_TOKEN_LIMIT")),
)
import inspect import inspect
import json import json
from datetime import date
from typing import Any, Callable from typing import Any, Callable
import ollama import ollama
...@@ -7,6 +8,7 @@ from openai import OpenAI ...@@ -7,6 +8,7 @@ from openai import OpenAI
from mistralai.client import Mistral from mistralai.client import Mistral
from app import config from app import config
from app.deterministic_services import llm_quota
def _filter_kwargs(func, kwargs: dict) -> dict: def _filter_kwargs(func, kwargs: dict) -> dict:
...@@ -23,6 +25,39 @@ def _extract_message(response) -> object: ...@@ -23,6 +25,39 @@ def _extract_message(response) -> object:
return response.message return response.message
return {} return {}
def _extract_total_tokens(response: object) -> int:
usage = None
if isinstance(response, dict):
usage = response.get("usage")
elif hasattr(response, "usage"):
usage = getattr(response, "usage")
if usage is None:
return 0
if isinstance(usage, dict):
total = usage.get("total_tokens")
if total is not None:
return int(total or 0)
prompt = usage.get("prompt_tokens") or usage.get("input_tokens") or 0
completion = usage.get("completion_tokens") or usage.get("output_tokens") or 0
return int(prompt) + int(completion)
total = getattr(usage, "total_tokens", None)
if total is not None:
return int(total or 0)
prompt = getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", None) or 0
completion = getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", None) or 0
return int(prompt) + int(completion)
def _record_call(result: dict, tokens: int | None = None) -> dict:
pg_url = config.get_postgres_url()
token_count = _extract_total_tokens(result.get("raw")) if tokens is None else tokens
llm_quota.record_usage(pg_url, date.today(), calls=1, tokens=token_count)
return result
def _chat_openai(messages: list[dict]) -> dict: def _chat_openai(messages: list[dict]) -> dict:
settings = config.get_openai_chat_settings() settings = config.get_openai_chat_settings()
if not settings: if not settings:
...@@ -77,21 +112,72 @@ def chat( ...@@ -77,21 +112,72 @@ def chat(
use_ollama: bool = False, use_ollama: bool = False,
use_mistral: bool = False, use_mistral: bool = False,
) -> dict: ) -> dict:
quota_settings = config.get_llm_quota_settings()
pg_url = config.get_postgres_url()
if use_mistral and tools: if use_mistral and tools:
raise RuntimeError("Mistral chat is currently only implemented for calls without tools.") raise RuntimeError("Mistral chat is currently only implemented for calls without tools.")
if use_mistral: if use_mistral:
mistral_result = _chat_mistral(messages) mistral_settings = config.get_mistral_chat_settings()
if not mistral_settings:
return {}
llm_quota.ensure_within_limits(
pg_url,
date.today(),
call_limit=quota_settings.daily_call_limit,
token_limit=quota_settings.daily_token_limit,
add_calls=1,
add_tokens=0,
)
try:
mistral_result = _chat_mistral(messages)
except Exception:
_record_call({"raw": None}, tokens=0)
raise
if mistral_result: if mistral_result:
return mistral_result return _record_call(mistral_result)
_record_call({"raw": None}, tokens=0)
return mistral_result
if not tools and not use_ollama: if not tools and not use_ollama:
openai_result = _chat_openai(messages) openai_settings = config.get_openai_chat_settings()
if openai_result: if openai_settings:
llm_quota.ensure_within_limits(
pg_url,
date.today(),
call_limit=quota_settings.daily_call_limit,
token_limit=quota_settings.daily_token_limit,
add_calls=1,
add_tokens=0,
)
try:
openai_result = _chat_openai(messages)
except Exception:
_record_call({"raw": None}, tokens=0)
raise
if openai_result:
return _record_call(openai_result)
_record_call({"raw": None}, tokens=0)
return openai_result return openai_result
mistral_result = _chat_mistral(messages) mistral_settings = config.get_mistral_chat_settings()
if mistral_result: if mistral_settings:
llm_quota.ensure_within_limits(
pg_url,
date.today(),
call_limit=quota_settings.daily_call_limit,
token_limit=quota_settings.daily_token_limit,
add_calls=1,
add_tokens=0,
)
try:
mistral_result = _chat_mistral(messages)
except Exception:
_record_call({"raw": None}, tokens=0)
raise
if mistral_result:
return _record_call(mistral_result)
_record_call({"raw": None}, tokens=0)
return mistral_result return mistral_result
settings = config.get_ollama_settings() settings = config.get_ollama_settings()
......
from __future__ import annotations
from dataclasses import dataclass
from datetime import date
from typing import Any
import psycopg
from psycopg.rows import dict_row
DATABASE_CREATION_SQL = """
CREATE TABLE IF NOT EXISTS llm_daily_usage (
usage_date DATE PRIMARY KEY,
call_count BIGINT NOT NULL DEFAULT 0,
token_count BIGINT NOT NULL DEFAULT 0,
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
"""
@dataclass(frozen=True)
class LLMDailyUsage:
usage_date: date
call_count: int
token_count: int
updated_at: str | None = None
def init_db(pg_url: str) -> None:
with psycopg.connect(pg_url) as conn:
with conn.cursor() as cur:
cur.execute(DATABASE_CREATION_SQL)
conn.commit()
def _ensure_row(conn: psycopg.Connection[Any], usage_date: date) -> None:
with conn.cursor() as cur:
cur.execute(
"""
INSERT INTO llm_daily_usage (usage_date)
VALUES (%s)
ON CONFLICT (usage_date) DO NOTHING
""",
(usage_date,),
)
def get_daily_usage(pg_url: str, usage_date: date) -> LLMDailyUsage:
init_db(pg_url)
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
cur.execute(
"""
SELECT usage_date, call_count, token_count, updated_at
FROM llm_daily_usage
WHERE usage_date = %s
""",
(usage_date,),
)
row = cur.fetchone()
if not row:
return LLMDailyUsage(usage_date=usage_date, call_count=0, token_count=0)
return LLMDailyUsage(
usage_date=row["usage_date"],
call_count=int(row["call_count"]),
token_count=int(row["token_count"]),
updated_at=row.get("updated_at").isoformat() if row.get("updated_at") else None,
)
def ensure_within_limits(
pg_url: str,
usage_date: date,
*,
call_limit: int | None,
token_limit: int | None,
add_calls: int = 1,
add_tokens: int = 0,
) -> None:
init_db(pg_url)
with psycopg.connect(pg_url) as conn:
with conn.transaction():
_ensure_row(conn, usage_date)
with conn.cursor() as cur:
cur.execute(
"""
SELECT call_count, token_count
FROM llm_daily_usage
WHERE usage_date = %s
FOR UPDATE
""",
(usage_date,),
)
row = cur.fetchone() or (0, 0)
next_calls = int(row[0]) + add_calls
next_tokens = int(row[1]) + add_tokens
if call_limit is not None and next_calls > call_limit:
raise QuotaExceededError("daily LLM call limit exceeded")
if token_limit is not None and next_tokens > token_limit:
raise QuotaExceededError("daily LLM token limit exceeded")
def record_usage(
pg_url: str,
usage_date: date,
*,
calls: int = 1,
tokens: int = 0,
) -> None:
init_db(pg_url)
with psycopg.connect(pg_url) as conn:
with conn.transaction():
_ensure_row(conn, usage_date)
with conn.cursor() as cur:
cur.execute(
"""
UPDATE llm_daily_usage
SET call_count = call_count + %s,
token_count = token_count + %s,
updated_at = NOW()
WHERE usage_date = %s
""",
(calls, tokens, usage_date),
)
def get_today_usage(pg_url: str) -> LLMDailyUsage:
return get_daily_usage(pg_url, date.today())
class QuotaExceededError(RuntimeError):
pass
...@@ -5,7 +5,8 @@ from fastapi import FastAPI ...@@ -5,7 +5,8 @@ from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from app.api import canvas, chat, context, health, orchestrator, tasks from app.api import canvas, chat, context, health, orchestrator, tasks
from app.config import get_frontend_url from app.config import get_frontend_url
from app.deterministic_services import embedding_provider from app.deterministic_services import embedding_provider, llm_quota
from app import config
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
...@@ -23,6 +24,15 @@ async def lifespan(_app: FastAPI): ...@@ -23,6 +24,15 @@ async def lifespan(_app: FastAPI):
yield yield
return return
try:
llm_quota.init_db(config.get_postgres_url())
logger.info("LLM quota table ensured")
except Exception:
logger.exception("LLM quota init failed")
health.set_readiness_failed("LLM quota init failed")
yield
return
try: try:
warmup_timing = embedding_provider.warmup_embedder() warmup_timing = embedding_provider.warmup_embedder()
logger.info("Embedding warmup finished: %s", warmup_timing) logger.info("Embedding warmup finished: %s", warmup_timing)
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment