大模型速度增加提示
This commit is contained in:
1 parent
4e764350a4
commit
3ea33e1dee
52 files changed
+4673
-176
No files matched your search
@@ -27,6 +27,7 @@ from app.gateway.routers import (
|
||||
auth,
|
||||
channels,
|
||||
feedback,
|
||||
llm_metrics,
|
||||
mcp,
|
||||
memory,
|
||||
models,
|
||||
@@ -35,6 +36,7 @@ from app.gateway.routers import (
|
||||
scheduled_tasks,
|
||||
skills,
|
||||
suggestions,
|
||||
system_settings,
|
||||
thread_runs,
|
||||
threads,
|
||||
uploads,
|
||||
@@ -495,6 +497,12 @@ This gateway provides custom endpoints for models, MCP configuration, skills, an
|
||||
# Scheduled task API is mounted at /api/scheduled-tasks
|
||||
app.include_router(scheduled_tasks.router)
|
||||
|
||||
# System settings API is mounted at /api/system-settings
|
||||
app.include_router(system_settings.router)
|
||||
|
||||
# Admin LLM metrics audit + export
|
||||
app.include_router(llm_metrics.router)
|
||||
|
||||
# Channels API is mounted at /api/channels
|
||||
app.include_router(channels.router)
|
||||
|
||||
|
||||
@@ -62,6 +62,15 @@ class LocalAuthProvider(AuthProvider):
|
||||
"""Get user by ID."""
|
||||
return await self._repo.get_user_by_id(user_id)
|
||||
|
||||
async def find_user_by_email_local_part(self, local_part: str) -> User | None:
|
||||
"""Find a user whose email starts with ``local_part@`` regardless of domain.
|
||||
|
||||
Used by the passwordless username-login flow to keep a single
|
||||
username mapped to the same account across legacy / canonical
|
||||
email-domain defaults.
|
||||
"""
|
||||
return await self._repo.find_user_by_email_local_part(local_part)
|
||||
|
||||
async def create_user(self, email: str, password: str | None = None, system_role: str = "user", needs_setup: bool = False) -> User:
|
||||
"""Create a new local user.
|
||||
|
||||
|
||||
@@ -61,6 +61,26 @@ class UserRepository(ABC):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def find_user_by_email_local_part(self, local_part: str) -> User | None:
|
||||
"""Find a user by the local part of their email (the bit before ``@``).
|
||||
|
||||
Used by the passwordless username-login flow so that a single
|
||||
username keeps mapping to the same account even when the synthetic
|
||||
email domain (e.g. ``cm.com`` vs the legacy ``local.deerflow``)
|
||||
differs between registration and current login.
|
||||
|
||||
Returns the oldest matching user when multiple rows share the same
|
||||
local part, so collisions resolve deterministically.
|
||||
|
||||
Args:
|
||||
local_part: Substring to match before the ``@``. Case-insensitive.
|
||||
|
||||
Returns:
|
||||
User if found, None otherwise.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def update_user(self, user: User) -> User:
|
||||
"""Update an existing user.
|
||||
|
||||
@@ -88,6 +88,27 @@ class SQLiteUserRepository(UserRepository):
|
||||
row = result.scalar_one_or_none()
|
||||
return self._row_to_user(row) if row is not None else None
|
||||
|
||||
async def find_user_by_email_local_part(self, local_part: str) -> User | None:
|
||||
local = (local_part or "").strip().lower()
|
||||
if not local:
|
||||
return None
|
||||
# Escape SQL LIKE wildcards so usernames like ``foo_bar`` or ``a%b`` do
|
||||
# not match unrelated rows. ``\`` is the chosen escape char below.
|
||||
escaped = local.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||||
pattern = f"{escaped}@%"
|
||||
# Case-insensitive prefix match on the email field, oldest first so a
|
||||
# subsequent registration cannot shadow an existing account.
|
||||
stmt = (
|
||||
select(UserRow)
|
||||
.where(func.lower(UserRow.email).like(pattern, escape="\\"))
|
||||
.order_by(UserRow.created_at.asc())
|
||||
.limit(1)
|
||||
)
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
return self._row_to_user(row) if row is not None else None
|
||||
|
||||
async def update_user(self, user: User) -> User:
|
||||
async with self._sf() as session:
|
||||
row = await session.get(UserRow, str(user.id))
|
||||
|
||||
@@ -90,6 +90,10 @@ async def langgraph_runtime(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
|
||||
app.state.scheduled_task_store = make_scheduled_task_store(sf)
|
||||
|
||||
from deerflow.persistence.llm_metrics import make_llm_metrics_store
|
||||
|
||||
app.state.llm_metrics_store = make_llm_metrics_store(sf)
|
||||
|
||||
from deerflow.persistence.agents import make_agent_store
|
||||
|
||||
app.state.agent_store = make_agent_store(sf)
|
||||
@@ -158,6 +162,11 @@ def get_scheduled_task_store(request: Request) -> ScheduledTaskStore:
|
||||
return val
|
||||
|
||||
|
||||
def get_llm_metrics_store(request: Request):
|
||||
"""Return the LLM metrics store, or None when not configured."""
|
||||
return getattr(request.app.state, "llm_metrics_store", None)
|
||||
|
||||
|
||||
def get_agent_store(request: Request) -> AgentStore:
|
||||
"""Return the custom-agent store."""
|
||||
val = getattr(request.app.state, "agent_store", None)
|
||||
|
||||
@@ -161,25 +161,40 @@ class MessageResponse(BaseModel):
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _normalize_username_local_part(username: str) -> str | None:
|
||||
"""Return the canonical local-part for a username, or None if user typed an email.
|
||||
|
||||
Used both to build the synthetic email and to look up legacy accounts that
|
||||
were registered with a different default domain. Returning ``None`` for
|
||||
raw-email input tells callers to skip the domain-agnostic fallback —
|
||||
if you typed ``alice@old.com``, you meant that specific row.
|
||||
"""
|
||||
raw = username.strip()
|
||||
if not raw:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="username cannot be empty")
|
||||
if "@" in raw:
|
||||
return None
|
||||
normalized = "".join(ch if (ch.isalnum() or ch in "._-") else "-" for ch in raw.lower()).strip("._-")
|
||||
if not normalized:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="username is invalid")
|
||||
return normalized
|
||||
|
||||
|
||||
def _username_to_email(username: str) -> str:
|
||||
"""Map a frontend username to a valid email used by user storage.
|
||||
|
||||
- If username already looks like an email (`contains @`), use it directly.
|
||||
- Otherwise generate `<normalized-username>@<domain>`.
|
||||
"""
|
||||
raw = username.strip()
|
||||
if not raw:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="username cannot be empty")
|
||||
local = _normalize_username_local_part(username)
|
||||
if local is None:
|
||||
return username.strip().lower()
|
||||
|
||||
if "@" in raw:
|
||||
return raw.lower()
|
||||
|
||||
normalized = "".join(ch if (ch.isalnum() or ch in "._-") else "-" for ch in raw.lower()).strip("._-")
|
||||
if not normalized:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="username is invalid")
|
||||
|
||||
domain = os.getenv("DEER_FLOW_USERNAME_LOGIN_DOMAIN", "local.deerflow").strip().lower() or "local.deerflow"
|
||||
return f"{normalized}@{domain}"
|
||||
# Default mail domain for passwordless / username-only registrations.
|
||||
# Kept overridable via env so individual deployments can swap it out, but the
|
||||
# built-in default targets the cm.com tenant.
|
||||
domain = os.getenv("DEER_FLOW_USERNAME_LOGIN_DOMAIN", "cm.com").strip().lower() or "cm.com"
|
||||
return f"{local}@{domain}"
|
||||
|
||||
|
||||
def _set_session_cookie(response: Response, token: str, request: Request) -> None:
|
||||
@@ -365,9 +380,21 @@ async def login_by_username(request: Request, response: Response, body: Username
|
||||
"""
|
||||
provider = get_local_provider()
|
||||
email = _username_to_email(body.username)
|
||||
local_part = _normalize_username_local_part(body.username)
|
||||
|
||||
created = False
|
||||
user = await provider.get_user_by_email(email)
|
||||
# Domain-agnostic fallback: if no row exists at the canonical
|
||||
# ``<username>@<current-default-domain>``, look for an existing user with
|
||||
# the same local part under any domain. This is the safety net that keeps
|
||||
# usernames globally unique when the default email domain changes (e.g.
|
||||
# ``local.deerflow`` → ``cm.com``) — without it, switching the default
|
||||
# would let the same username log into two different accounts.
|
||||
if user is None and local_part is not None:
|
||||
legacy = await provider.find_user_by_email_local_part(local_part)
|
||||
if legacy is not None:
|
||||
user = legacy
|
||||
logger.info("Matched username %r to legacy account %s", body.username, legacy.email)
|
||||
if user is None:
|
||||
try:
|
||||
role = "admin" if body.username == "admin" else "user"
|
||||
@@ -377,6 +404,8 @@ async def login_by_username(request: Request, response: Response, body: Username
|
||||
except ValueError:
|
||||
# Concurrent first login may create the same user in parallel.
|
||||
user = await provider.get_user_by_email(email)
|
||||
if user is None and local_part is not None:
|
||||
user = await provider.find_user_by_email_local_part(local_part)
|
||||
if user is None:
|
||||
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to create or load user")
|
||||
elif body.username == "admin" and user.system_role != "admin":
|
||||
|
||||
@@ -0,0 +1,323 @@
|
||||
"""Admin-only LLM call metrics: list + Excel export.
|
||||
|
||||
Lets administrators audit every provider invocation made by the runtime —
|
||||
including failures — and download the full table as ``.xlsx``. The recording
|
||||
is unconditional on the backend, so this view always reflects every call.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
from urllib.parse import quote
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Query, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.deps import get_llm_metrics_store, get_local_provider, get_optional_user_from_request
|
||||
from deerflow.persistence.llm_metrics.base import LlmMetricsQuery, LlmMetricsStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/api/admin/llm-metrics", tags=["llm-metrics"])
|
||||
|
||||
|
||||
class LlmMetricResponse(BaseModel):
|
||||
id: str
|
||||
created_at: str | None = None
|
||||
user_id: str | None = None
|
||||
user_name: str | None = Field(
|
||||
default=None,
|
||||
description="Local part (before @) of the user's email, resolved from user_id",
|
||||
)
|
||||
thread_id: str | None = None
|
||||
run_id: str | None = None
|
||||
agent_name: str | None = None
|
||||
model_name: str | None = None
|
||||
duration_ms: int = 0
|
||||
input_tokens: int | None = None
|
||||
output_tokens: int | None = None
|
||||
total_tokens: int | None = None
|
||||
tokens_per_sec: float | None = None
|
||||
status: str = "success"
|
||||
error_type: str | None = None
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
class LlmMetricsListResponse(BaseModel):
|
||||
items: list[LlmMetricResponse]
|
||||
total: int = Field(..., description="Total rows matching the filter, ignoring pagination")
|
||||
limit: int
|
||||
offset: int
|
||||
|
||||
|
||||
async def _require_admin(request: Request) -> None:
|
||||
user = await get_optional_user_from_request(request)
|
||||
if user is None:
|
||||
return
|
||||
if getattr(user, "system_role", None) != "admin":
|
||||
raise HTTPException(status_code=403, detail="LLM metrics admin view is admin-only")
|
||||
|
||||
|
||||
def _require_store(request: Request) -> LlmMetricsStore:
|
||||
store = get_llm_metrics_store(request)
|
||||
if store is None:
|
||||
raise HTTPException(status_code=503, detail="LLM metrics store is not configured")
|
||||
return store
|
||||
|
||||
|
||||
def _parse_dt(value: str | None) -> datetime | None:
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
# Accept both ``Z`` suffix and ``+00:00``.
|
||||
return datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=422, detail=f"Invalid datetime: {value!r} ({exc})") from exc
|
||||
|
||||
|
||||
def _email_local_part(email: str | None) -> str | None:
|
||||
if not email:
|
||||
return None
|
||||
at = email.find("@")
|
||||
return email[:at] if at > 0 else email
|
||||
|
||||
|
||||
def _derive_tokens_per_sec(row: dict[str, Any]) -> None:
|
||||
"""Fill in ``tokens_per_sec`` when the row was recorded before we knew how
|
||||
to extract ``output_tokens`` from response_metadata.token_usage.
|
||||
|
||||
Uses ``output_tokens / elapsed`` when available, otherwise falls back to
|
||||
``total_tokens / elapsed`` as an approximation. Newer rows already have
|
||||
the field set during recording, so this only kicks in for legacy data.
|
||||
"""
|
||||
if not isinstance(row, dict):
|
||||
return
|
||||
if row.get("tokens_per_sec") not in (None, 0):
|
||||
return
|
||||
duration_ms = row.get("duration_ms") or 0
|
||||
if duration_ms <= 0:
|
||||
return
|
||||
numerator = row.get("output_tokens") or row.get("total_tokens")
|
||||
if not numerator:
|
||||
return
|
||||
row["tokens_per_sec"] = round((numerator / duration_ms) * 1000, 3)
|
||||
|
||||
|
||||
async def _resolve_user_name_to_id(name: str | None) -> str | None:
|
||||
"""Look up a UUID for the given username (email local-part)."""
|
||||
if not name:
|
||||
return None
|
||||
try:
|
||||
provider = get_local_provider()
|
||||
except Exception:
|
||||
return None
|
||||
try:
|
||||
user = await provider.find_user_by_email_local_part(name)
|
||||
except Exception:
|
||||
return None
|
||||
return str(user.id) if user is not None else None
|
||||
|
||||
|
||||
async def _attach_usernames(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Resolve each row's ``user_id`` UUID to a friendly ``user_name``.
|
||||
|
||||
Batches lookups so one query per distinct user_id, not per row. Tolerates
|
||||
missing accounts (deleted users) by leaving ``user_name`` blank.
|
||||
"""
|
||||
if not rows:
|
||||
return rows
|
||||
user_ids = {row["user_id"] for row in rows if isinstance(row, dict) and row.get("user_id")}
|
||||
if not user_ids:
|
||||
return rows
|
||||
try:
|
||||
provider = get_local_provider()
|
||||
except Exception:
|
||||
provider = None
|
||||
email_by_id: dict[str, str] = {}
|
||||
if provider is not None:
|
||||
for uid in user_ids:
|
||||
try:
|
||||
user = await provider.get_user(uid)
|
||||
except Exception:
|
||||
user = None
|
||||
if user is not None and user.email:
|
||||
email_by_id[uid] = user.email
|
||||
for row in rows:
|
||||
if isinstance(row, dict) and row.get("user_id"):
|
||||
email = email_by_id.get(row["user_id"])
|
||||
row["user_name"] = _email_local_part(email) or row["user_id"]
|
||||
return rows
|
||||
|
||||
|
||||
def _build_query(
|
||||
*,
|
||||
user_id: str | None,
|
||||
model_name: str | None,
|
||||
status: str | None,
|
||||
since: str | None,
|
||||
until: str | None,
|
||||
limit: int,
|
||||
offset: int,
|
||||
) -> LlmMetricsQuery:
|
||||
return LlmMetricsQuery(
|
||||
user_id=user_id or None,
|
||||
model_name=model_name or None,
|
||||
status=status or None,
|
||||
since=_parse_dt(since),
|
||||
until=_parse_dt(until),
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
|
||||
@router.get("", response_model=LlmMetricsListResponse)
|
||||
async def list_llm_metrics(
|
||||
request: Request,
|
||||
user_id: str | None = Query(default=None),
|
||||
user_name: str | None = Query(default=None, description="Username (email local part) to filter by"),
|
||||
model_name: str | None = Query(default=None),
|
||||
status: str | None = Query(default=None, pattern="^(success|error)$"),
|
||||
since: str | None = Query(default=None, description="ISO 8601 datetime, inclusive lower bound"),
|
||||
until: str | None = Query(default=None, description="ISO 8601 datetime, inclusive upper bound"),
|
||||
limit: int = Query(default=50, ge=1, le=500),
|
||||
offset: int = Query(default=0, ge=0),
|
||||
) -> LlmMetricsListResponse:
|
||||
await _require_admin(request)
|
||||
store = _require_store(request)
|
||||
if not user_id and user_name:
|
||||
resolved = await _resolve_user_name_to_id(user_name)
|
||||
if resolved is None:
|
||||
# Username unknown → no rows; short-circuit instead of a global scan.
|
||||
return LlmMetricsListResponse(items=[], total=0, limit=limit, offset=offset)
|
||||
user_id = resolved
|
||||
query = _build_query(
|
||||
user_id=user_id,
|
||||
model_name=model_name,
|
||||
status=status,
|
||||
since=since,
|
||||
until=until,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
rows = await store.list(query)
|
||||
total = await store.count(query)
|
||||
await _attach_usernames(rows)
|
||||
for row in rows:
|
||||
_derive_tokens_per_sec(row)
|
||||
return LlmMetricsListResponse(
|
||||
items=[LlmMetricResponse(**row) for row in rows],
|
||||
total=total,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
|
||||
def _format_cell(value: Any) -> Any:
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, datetime):
|
||||
return value.replace(tzinfo=None) if value.tzinfo else value
|
||||
return value
|
||||
|
||||
|
||||
@router.get("/export")
|
||||
async def export_llm_metrics(
|
||||
request: Request,
|
||||
user_id: str | None = Query(default=None),
|
||||
user_name: str | None = Query(default=None),
|
||||
model_name: str | None = Query(default=None),
|
||||
status: str | None = Query(default=None, pattern="^(success|error)$"),
|
||||
since: str | None = Query(default=None),
|
||||
until: str | None = Query(default=None),
|
||||
):
|
||||
"""Stream the full filtered metric set as an ``.xlsx`` file."""
|
||||
await _require_admin(request)
|
||||
store = _require_store(request)
|
||||
if not user_id and user_name:
|
||||
resolved = await _resolve_user_name_to_id(user_name)
|
||||
if resolved is None:
|
||||
# No matching user — return an empty xlsx so the user still gets a
|
||||
# file (matching the list view's empty state).
|
||||
user_id = "__no_such_user__"
|
||||
else:
|
||||
user_id = resolved
|
||||
|
||||
try:
|
||||
from openpyxl import Workbook
|
||||
except ImportError as exc:
|
||||
raise HTTPException(status_code=500, detail=f"openpyxl is not installed: {exc}") from exc
|
||||
|
||||
query = _build_query(
|
||||
user_id=user_id,
|
||||
model_name=model_name,
|
||||
status=status,
|
||||
since=since,
|
||||
until=until,
|
||||
limit=1, # unused: iter_all paginates internally
|
||||
offset=0,
|
||||
)
|
||||
|
||||
workbook = Workbook(write_only=True)
|
||||
sheet = workbook.create_sheet("llm_call_metrics")
|
||||
columns = [
|
||||
"id",
|
||||
"created_at",
|
||||
"status",
|
||||
"user_name",
|
||||
"user_id",
|
||||
"thread_id",
|
||||
"run_id",
|
||||
"agent_name",
|
||||
"model_name",
|
||||
"duration_ms",
|
||||
"input_tokens",
|
||||
"output_tokens",
|
||||
"total_tokens",
|
||||
"tokens_per_sec",
|
||||
"error_type",
|
||||
"error_message",
|
||||
]
|
||||
sheet.append(columns)
|
||||
|
||||
# Pre-resolve all distinct user_ids in chunks so the export doesn't pay one
|
||||
# round trip per row.
|
||||
buffered: list[dict[str, Any]] = []
|
||||
FLUSH_AT = 200
|
||||
|
||||
async def flush() -> int:
|
||||
if not buffered:
|
||||
return 0
|
||||
await _attach_usernames(buffered)
|
||||
for row in buffered:
|
||||
_derive_tokens_per_sec(row)
|
||||
sheet.append([_format_cell(row.get(col)) for col in columns])
|
||||
n = len(buffered)
|
||||
buffered.clear()
|
||||
return n
|
||||
|
||||
count = 0
|
||||
async for row in store.iter_all(query):
|
||||
buffered.append(row)
|
||||
if len(buffered) >= FLUSH_AT:
|
||||
count += await flush()
|
||||
count += await flush()
|
||||
|
||||
buffer = BytesIO()
|
||||
workbook.save(buffer)
|
||||
buffer.seek(0)
|
||||
logger.info("Exported %d llm_call_metrics rows", count)
|
||||
|
||||
timestamp = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
|
||||
filename = f"llm_call_metrics-{timestamp}.xlsx"
|
||||
|
||||
headers = {
|
||||
"Content-Disposition": f"attachment; filename=\"{filename}\"; filename*=UTF-8''{quote(filename)}",
|
||||
}
|
||||
return StreamingResponse(
|
||||
buffer,
|
||||
media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
headers=headers,
|
||||
)
|
||||
@@ -3,15 +3,31 @@ import logging
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.deps import get_optional_user_from_request
|
||||
from deerflow.config.extensions_config import ExtensionsConfig, get_extensions_config, reload_extensions_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/api", tags=["mcp"])
|
||||
|
||||
|
||||
async def _require_admin(request: Request) -> None:
|
||||
"""Reject the request unless the caller is an authenticated admin user.
|
||||
|
||||
Allows the call through silently when auth is disabled (no user attached),
|
||||
so existing single-user / no-auth deployments keep working.
|
||||
"""
|
||||
user = await get_optional_user_from_request(request)
|
||||
if user is None:
|
||||
# No authenticated user — auth might be disabled or middleware bypassed.
|
||||
# Defer to the AuthMiddleware decision and let the call through.
|
||||
return
|
||||
if getattr(user, "system_role", None) != "admin":
|
||||
raise HTTPException(status_code=403, detail="MCP configuration requires admin privileges")
|
||||
|
||||
|
||||
class McpOAuthConfigResponse(BaseModel):
|
||||
"""OAuth configuration for an MCP server."""
|
||||
|
||||
@@ -43,6 +59,10 @@ class McpServerConfigResponse(BaseModel):
|
||||
headers: dict[str, str] = Field(default_factory=dict, description="HTTP headers to send (for sse or http type)")
|
||||
oauth: McpOAuthConfigResponse | None = Field(default=None, description="OAuth configuration for MCP HTTP/SSE servers")
|
||||
description: str = Field(default="", description="Human-readable description of what this MCP server provides")
|
||||
disabled_tools: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Tool names (without server-name prefix) to hide from the agent",
|
||||
)
|
||||
|
||||
|
||||
class McpConfigResponse(BaseModel):
|
||||
@@ -101,7 +121,7 @@ async def get_mcp_configuration() -> McpConfigResponse:
|
||||
summary="Update MCP Configuration",
|
||||
description="Update Model Context Protocol (MCP) server configurations and save to file.",
|
||||
)
|
||||
async def update_mcp_configuration(request: McpConfigUpdateRequest) -> McpConfigResponse:
|
||||
async def update_mcp_configuration(request: McpConfigUpdateRequest, http_request: Request) -> McpConfigResponse:
|
||||
"""Update the MCP configuration.
|
||||
|
||||
This will:
|
||||
@@ -133,6 +153,22 @@ async def update_mcp_configuration(request: McpConfigUpdateRequest) -> McpConfig
|
||||
}
|
||||
```
|
||||
"""
|
||||
await _require_admin(http_request)
|
||||
|
||||
# Light validation: each server must have the fields required by its transport.
|
||||
for server_name, server in request.mcp_servers.items():
|
||||
if not server_name.strip():
|
||||
raise HTTPException(status_code=400, detail="MCP server name cannot be empty")
|
||||
transport = (server.type or "stdio").lower()
|
||||
if transport == "stdio":
|
||||
if not (server.command or "").strip():
|
||||
raise HTTPException(status_code=400, detail=f"MCP server '{server_name}' uses stdio transport but command is empty")
|
||||
elif transport in ("http", "sse"):
|
||||
if not (server.url or "").strip():
|
||||
raise HTTPException(status_code=400, detail=f"MCP server '{server_name}' uses {transport} transport but url is empty")
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail=f"MCP server '{server_name}' has unsupported transport: {transport}")
|
||||
|
||||
try:
|
||||
# Get the current config path (or determine where to save it)
|
||||
config_path = ExtensionsConfig.resolve_config_path()
|
||||
|
||||
@@ -14,7 +14,7 @@ from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import FileResponse, Response, StreamingResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.deps import get_current_user
|
||||
from app.gateway.deps import get_current_user, get_local_provider
|
||||
from deerflow.runtime.scheduler import get_scheduled_task_service
|
||||
|
||||
router = APIRouter(prefix="/api/scheduled-tasks", tags=["scheduled-tasks"])
|
||||
@@ -61,6 +61,50 @@ async def _user_id(request: Request) -> str:
|
||||
return user_id
|
||||
|
||||
|
||||
async def _enrich_with_creator(tasks: list[dict[str, Any]] | dict[str, Any] | None) -> Any:
|
||||
"""Attach ``creator_email`` / ``creator_user_id`` to one or many task dicts.
|
||||
|
||||
Looks up the auth user once per distinct ``user_id`` in the batch so we do
|
||||
not pay N+1 round trips when rendering a long task list. Tolerates missing
|
||||
user records (deleted accounts) by leaving the email blank.
|
||||
"""
|
||||
if tasks is None:
|
||||
return tasks
|
||||
items = tasks if isinstance(tasks, list) else [tasks]
|
||||
|
||||
user_ids = {str(t.get("user_id")) for t in items if isinstance(t, dict) and t.get("user_id")}
|
||||
if not user_ids:
|
||||
return tasks
|
||||
|
||||
try:
|
||||
provider = get_local_provider()
|
||||
except Exception:
|
||||
# Auth provider not initialised (e.g. memory-mode test harness): keep
|
||||
# the bare user_id and let the frontend render the raw id.
|
||||
provider = None
|
||||
|
||||
email_by_id: dict[str, str] = {}
|
||||
if provider is not None:
|
||||
for uid in user_ids:
|
||||
try:
|
||||
user = await provider.get_user(uid)
|
||||
except Exception:
|
||||
user = None
|
||||
if user is not None:
|
||||
email_by_id[uid] = user.email
|
||||
|
||||
for item in items:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
uid = item.get("user_id")
|
||||
if uid is None:
|
||||
continue
|
||||
uid_str = str(uid)
|
||||
item["creator_user_id"] = uid_str
|
||||
item["creator_email"] = email_by_id.get(uid_str, "")
|
||||
return tasks
|
||||
|
||||
|
||||
def _safe_filename(value: str, default: str = "scheduled-task-result") -> str:
|
||||
name = re.sub(r'[\\/:*?"<>|\r\n]+', "-", value).strip(" .")
|
||||
return name or default
|
||||
@@ -133,7 +177,9 @@ async def get_scheduler_thread(request: Request) -> dict:
|
||||
async def list_scheduled_tasks(request: Request) -> list[dict]:
|
||||
service = get_scheduled_task_service()
|
||||
user_id = await _user_id(request)
|
||||
return await service.list_tasks(user_id)
|
||||
tasks = await service.list_tasks(user_id)
|
||||
await _enrich_with_creator(tasks)
|
||||
return tasks
|
||||
|
||||
|
||||
@router.get("/{task_id}/runs")
|
||||
@@ -158,7 +204,10 @@ async def get_scheduled_task_run(run_id: str, request: Request) -> dict:
|
||||
async def delete_scheduled_task_run(run_id: str, request: Request) -> dict:
|
||||
service = get_scheduled_task_service()
|
||||
user_id = await _user_id(request)
|
||||
deleted = await service.delete_run(user_id, run_id)
|
||||
try:
|
||||
deleted = await service.delete_run(user_id, run_id)
|
||||
except PermissionError as exc:
|
||||
raise HTTPException(status_code=403, detail=str(exc))
|
||||
if not deleted:
|
||||
raise HTTPException(status_code=404, detail=f"Run {run_id} not found")
|
||||
return {"success": True}
|
||||
@@ -298,7 +347,7 @@ async def update_delivery_profile(body: ScheduledTaskDeliveryProfileUpdateReques
|
||||
async def create_scheduled_task(body: ScheduledTaskCreateRequest, request: Request) -> dict:
|
||||
service = get_scheduled_task_service()
|
||||
user_id = await _user_id(request)
|
||||
return await service.create_task(
|
||||
task = await service.create_task(
|
||||
user_id=user_id,
|
||||
name=body.name,
|
||||
prompt=body.prompt,
|
||||
@@ -307,6 +356,8 @@ async def create_scheduled_task(body: ScheduledTaskCreateRequest, request: Reque
|
||||
execution_agent_name=body.execution_agent_name,
|
||||
execution_context=body.execution_context,
|
||||
)
|
||||
await _enrich_with_creator(task)
|
||||
return task
|
||||
|
||||
|
||||
@router.patch("/{task_id}")
|
||||
@@ -316,6 +367,7 @@ async def update_scheduled_task(task_id: str, body: ScheduledTaskUpdateRequest,
|
||||
task = await service.update_task(user_id, task_id, body.model_dump(exclude_unset=True))
|
||||
if task is None:
|
||||
raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
|
||||
await _enrich_with_creator(task)
|
||||
return task
|
||||
|
||||
|
||||
@@ -366,7 +418,17 @@ async def trigger_scheduled_task(task_id: str, request: Request) -> dict:
|
||||
async def list_published_scheduled_tasks(request: Request, limit: int = 50, offset: int = 0) -> list[dict]:
|
||||
service = get_scheduled_task_service()
|
||||
await _user_id(request)
|
||||
return await service.list_published(limit=max(1, min(limit, 200)), offset=max(0, offset))
|
||||
tasks = await service.list_published(limit=max(1, min(limit, 200)), offset=max(0, offset))
|
||||
# published list uses owner_user_id; resolve it into a creator_email here too
|
||||
# so the discovery view can show the author's account.
|
||||
for task in tasks:
|
||||
if isinstance(task, dict) and task.get("owner_user_id"):
|
||||
task["user_id"] = task["owner_user_id"]
|
||||
await _enrich_with_creator(tasks)
|
||||
for task in tasks:
|
||||
if isinstance(task, dict):
|
||||
task.pop("user_id", None)
|
||||
return tasks
|
||||
|
||||
|
||||
@router.post("/{task_id}/publish")
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
"""Admin-configurable system-wide runtime settings."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.deps import get_optional_user_from_request
|
||||
from deerflow.config.system_settings import (
|
||||
PromptPrefixSettings,
|
||||
SystemSettings,
|
||||
load_system_settings,
|
||||
save_system_settings,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/api/system-settings", tags=["system-settings"])
|
||||
|
||||
|
||||
class PromptPrefixResponse(BaseModel):
|
||||
enabled: bool = Field(..., description="Whether the toggle is exposed to users")
|
||||
switch_label: str = Field(..., description="Toggle label shown in the chat input")
|
||||
prompt_prefix: str = Field(..., description="Prefix text prepended to the first user message")
|
||||
|
||||
|
||||
class PromptPrefixUpdateRequest(BaseModel):
|
||||
enabled: bool | None = None
|
||||
switch_label: str | None = None
|
||||
prompt_prefix: str | None = None
|
||||
|
||||
|
||||
def _to_response(settings: PromptPrefixSettings) -> PromptPrefixResponse:
|
||||
return PromptPrefixResponse(
|
||||
enabled=settings.enabled,
|
||||
switch_label=settings.switch_label,
|
||||
prompt_prefix=settings.prompt_prefix,
|
||||
)
|
||||
|
||||
|
||||
async def _require_admin(request: Request) -> None:
|
||||
user = await get_optional_user_from_request(request)
|
||||
if user is None:
|
||||
return
|
||||
if getattr(user, "system_role", None) != "admin":
|
||||
raise HTTPException(status_code=403, detail="System settings can only be modified by an admin")
|
||||
|
||||
|
||||
@router.get("/prompt-prefix", response_model=PromptPrefixResponse)
|
||||
async def get_prompt_prefix() -> PromptPrefixResponse:
|
||||
"""Return the current prompt-prefix configuration (readable by any user)."""
|
||||
return _to_response(load_system_settings().prompt_prefix)
|
||||
|
||||
|
||||
@router.put("/prompt-prefix", response_model=PromptPrefixResponse)
|
||||
async def update_prompt_prefix(body: PromptPrefixUpdateRequest, request: Request) -> PromptPrefixResponse:
|
||||
"""Update the prompt-prefix configuration. Admin only."""
|
||||
await _require_admin(request)
|
||||
|
||||
current = load_system_settings()
|
||||
fields = body.model_dump(exclude_unset=True)
|
||||
updated_prefix = current.prompt_prefix.model_copy(update=fields)
|
||||
saved = save_system_settings(SystemSettings(prompt_prefix=updated_prefix))
|
||||
return _to_response(saved.prompt_prefix)
|
||||
@@ -19,6 +19,7 @@ from langchain_core.messages import HumanMessage
|
||||
|
||||
from app.gateway.deps import get_run_context, get_run_manager, get_stream_bridge
|
||||
from app.gateway.utils import sanitize_log_param
|
||||
from deerflow.config.system_settings import load_system_settings
|
||||
from deerflow.runtime import (
|
||||
END_SENTINEL,
|
||||
HEARTBEAT_SENTINEL,
|
||||
@@ -98,6 +99,45 @@ def normalize_input(raw_input: dict[str, Any] | None) -> dict[str, Any]:
|
||||
_DEFAULT_ASSISTANT_ID = "lead_agent"
|
||||
|
||||
|
||||
def _prepend_prompt_prefix(graph_input: dict[str, Any], prefix: str) -> dict[str, Any]:
|
||||
"""Return ``graph_input`` with the prefix prepended to its first HumanMessage.
|
||||
|
||||
Leaves the input unchanged when there are no messages or the first message
|
||||
is not a human turn. Idempotent in the sense that it only touches the first
|
||||
human message, so re-running this with the same prefix would double-prepend
|
||||
— callers must gate it on first-turn detection.
|
||||
"""
|
||||
if not prefix:
|
||||
return graph_input
|
||||
messages = graph_input.get("messages")
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return graph_input
|
||||
new_messages = list(messages)
|
||||
for idx, msg in enumerate(new_messages):
|
||||
if isinstance(msg, HumanMessage):
|
||||
existing = msg.content
|
||||
if isinstance(existing, str):
|
||||
new_messages[idx] = HumanMessage(content=f"{prefix}\n\n{existing}".strip("\n"))
|
||||
elif isinstance(existing, list):
|
||||
# Structured content (text blocks + media). Prepend a new text block.
|
||||
new_messages[idx] = HumanMessage(content=[{"type": "text", "text": prefix}, *existing])
|
||||
else:
|
||||
new_messages[idx] = HumanMessage(content=f"{prefix}\n\n{existing!s}")
|
||||
break
|
||||
return {**graph_input, "messages": new_messages}
|
||||
|
||||
|
||||
def _prompt_prefix_opted_in(body: Any) -> bool:
|
||||
"""Detect whether the caller asked for the prompt prefix to be applied."""
|
||||
context = getattr(body, "context", None)
|
||||
if isinstance(context, Mapping) and context.get("prompt_prefix_enabled"):
|
||||
return True
|
||||
metadata = getattr(body, "metadata", None)
|
||||
if isinstance(metadata, Mapping) and metadata.get("prompt_prefix_enabled"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# Whitelist of run-context keys that the langgraph-compat layer forwards from
|
||||
# ``body.context`` into the run config. ``config["context"]`` exists in
|
||||
# LangGraph >=0.6, but these values must be written to both ``configurable``
|
||||
@@ -340,6 +380,7 @@ async def start_run(
|
||||
# Upsert thread metadata so the thread appears in /threads/search,
|
||||
# even for threads that were never explicitly created via POST /threads
|
||||
# (e.g. stateless runs).
|
||||
thread_existed = False
|
||||
try:
|
||||
existing = await run_ctx.thread_store.get(thread_id)
|
||||
if existing is None:
|
||||
@@ -349,12 +390,24 @@ async def start_run(
|
||||
metadata=body.metadata,
|
||||
)
|
||||
else:
|
||||
thread_existed = True
|
||||
await run_ctx.thread_store.update_status(thread_id, "running")
|
||||
except Exception:
|
||||
logger.warning("Failed to upsert thread_meta for %s (non-fatal)", sanitize_log_param(thread_id))
|
||||
|
||||
agent_factory = resolve_agent_factory(body.assistant_id)
|
||||
graph_input = normalize_input(body.input)
|
||||
|
||||
# Server-side safety net for the admin-configurable prompt prefix: when the
|
||||
# caller opted in via context.prompt_prefix_enabled AND this is the very
|
||||
# first turn of the thread, prepend the configured prefix to the human
|
||||
# message. We deliberately skip already-running threads so the prefix is
|
||||
# never inserted twice — the frontend toggle is a per-message hint, but
|
||||
# this check is what makes the feature safe against client bugs / replays.
|
||||
if not thread_existed and _prompt_prefix_opted_in(body):
|
||||
prefix_settings = load_system_settings().prompt_prefix
|
||||
if prefix_settings.enabled and prefix_settings.prompt_prefix.strip():
|
||||
graph_input = _prepend_prompt_prefix(graph_input, prefix_settings.prompt_prefix.strip())
|
||||
config = build_run_config(thread_id, body.config, body.metadata, assistant_id=body.assistant_id)
|
||||
|
||||
# Merge DeerFlow-specific context overrides into both ``configurable`` and ``context``.
|
||||
|
||||
+335
@@ -0,0 +1,335 @@
|
||||
"""Middleware that records per-LLM-call metrics for the admin audit page.
|
||||
|
||||
Captures **every** provider invocation regardless of the frontend speed-meter
|
||||
toggle: the goal is a complete admin-side audit log of LLM usage, including
|
||||
failures. The user-facing speed indicator (controlled by
|
||||
``LocalSettings.display.streamSpeedMeterEnabled``) is a separate, purely
|
||||
cosmetic feature — turning it off does not stop this recording.
|
||||
|
||||
Each call ends with one row in ``llm_call_metrics``:
|
||||
- Successes record duration, token counts, and ``tokens_per_sec``.
|
||||
- Errors record the duration spent before the failure, the exception class as
|
||||
``error_type`` and a one-line ``error_message``.
|
||||
|
||||
Writes go through a best-effort store: a DB outage logs a warning and the
|
||||
agent runtime keeps streaming. We intentionally place this middleware
|
||||
*inside* :class:`LLMErrorHandlingMiddleware` (i.e. nested deeper in the
|
||||
``wrap_model_call`` chain) so we observe the **real** provider call timings —
|
||||
not the retry-wrapped successful response.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any, override
|
||||
|
||||
from langchain.agents import AgentState
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain.agents.middleware.types import (
|
||||
ModelCallResult,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
)
|
||||
from langchain_core.messages import AIMessage
|
||||
from langgraph.errors import GraphBubbleUp
|
||||
|
||||
from deerflow.persistence.engine import get_session_factory
|
||||
from deerflow.persistence.llm_metrics import LlmMetricRecord, make_llm_metrics_store
|
||||
from deerflow.persistence.llm_metrics.base import LlmMetricsStore
|
||||
from deerflow.runtime.user_context import get_effective_user_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def _safe_get_thread_id(request: ModelRequest) -> str | None:
|
||||
state = getattr(request, "state", None)
|
||||
if isinstance(state, dict):
|
||||
thread_id = state.get("thread_id")
|
||||
if isinstance(thread_id, str) and thread_id:
|
||||
return thread_id
|
||||
runtime = getattr(request, "runtime", None)
|
||||
if runtime is None:
|
||||
return None
|
||||
config = getattr(runtime, "config", None)
|
||||
configurable = (config or {}).get("configurable", {}) if isinstance(config, dict) else {}
|
||||
thread_id = configurable.get("thread_id") if isinstance(configurable, dict) else None
|
||||
if isinstance(thread_id, str) and thread_id:
|
||||
return thread_id
|
||||
return None
|
||||
|
||||
|
||||
def _safe_get_run_id(request: ModelRequest) -> str | None:
|
||||
runtime = getattr(request, "runtime", None)
|
||||
if runtime is None:
|
||||
return None
|
||||
context = getattr(runtime, "context", None)
|
||||
if isinstance(context, dict):
|
||||
for key in ("run_id", "langgraph_run_id"):
|
||||
value = context.get(key)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
config = getattr(runtime, "config", None)
|
||||
if isinstance(config, dict):
|
||||
metadata = config.get("metadata") if isinstance(config.get("metadata"), dict) else {}
|
||||
if isinstance(metadata, dict):
|
||||
for key in ("run_id", "langgraph_run_id"):
|
||||
value = metadata.get(key)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _safe_get_model_name(request: ModelRequest) -> str | None:
|
||||
model = getattr(request, "model", None)
|
||||
for attr in ("model_name", "model", "name"):
|
||||
value = getattr(model, attr, None)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _safe_get_agent_name(request: ModelRequest) -> str | None:
|
||||
runtime = getattr(request, "runtime", None)
|
||||
if runtime is None:
|
||||
return None
|
||||
context = getattr(runtime, "context", None)
|
||||
if isinstance(context, dict):
|
||||
for key in ("agent_name", "agent_id"):
|
||||
value = context.get(key)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _coerce_int(value: Any) -> int | None:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
if isinstance(value, float):
|
||||
return int(value)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _extract_usage(response: Any) -> dict[str, int]:
|
||||
"""Pull token counts out of an AIMessage, falling back to provider-native shapes.
|
||||
|
||||
The handler in ``wrap_model_call`` actually hands us a
|
||||
:class:`langchain.agents.middleware.types.ModelResponse` dataclass whose
|
||||
``.result`` is a ``list[BaseMessage]``. Earlier versions of this helper
|
||||
only matched ``AIMessage`` and ``list``, which meant we silently dropped
|
||||
every successful call and stored None for input/output/total. To stay
|
||||
robust we now also handle ``ModelResponse`` (and ``ExtendedModelResponse``,
|
||||
which wraps a ModelResponse), and we walk the messages list for the first
|
||||
``AIMessage``.
|
||||
|
||||
LangChain normalises usage into ``message.usage_metadata`` with the keys
|
||||
``input_tokens`` / ``output_tokens`` / ``total_tokens``. Several
|
||||
OpenAI-compatible providers (DeepSeek, vLLM, some Azure deployments) leave
|
||||
that dict partial — typically only ``total_tokens`` is populated — and the
|
||||
real breakdown sits under ``response_metadata.token_usage`` with the
|
||||
provider-native key names ``prompt_tokens`` / ``completion_tokens``. We
|
||||
merge both so the admin audit view always has the most complete numbers
|
||||
available.
|
||||
"""
|
||||
message: AIMessage | None = None
|
||||
|
||||
def _find_ai(candidate: Any) -> AIMessage | None:
|
||||
if isinstance(candidate, AIMessage):
|
||||
return candidate
|
||||
if isinstance(candidate, list):
|
||||
for item in candidate:
|
||||
found = _find_ai(item)
|
||||
if found is not None:
|
||||
return found
|
||||
return None
|
||||
|
||||
message = _find_ai(response)
|
||||
if message is None:
|
||||
# ModelResponse / ExtendedModelResponse — defined as dataclasses with a
|
||||
# ``result`` field holding ``list[BaseMessage]``. Recurse via attribute
|
||||
# access so we don't tightly couple to the dataclass class itself.
|
||||
for attr in ("result", "model_response"):
|
||||
inner = getattr(response, attr, None)
|
||||
if inner is None:
|
||||
continue
|
||||
message = _find_ai(inner)
|
||||
if message is None and attr == "model_response":
|
||||
inner_result = getattr(inner, "result", None)
|
||||
message = _find_ai(inner_result)
|
||||
if message is not None:
|
||||
break
|
||||
if message is None:
|
||||
return {}
|
||||
|
||||
out: dict[str, int] = {}
|
||||
|
||||
usage = getattr(message, "usage_metadata", None) or {}
|
||||
if isinstance(usage, dict):
|
||||
for key in ("input_tokens", "output_tokens", "total_tokens"):
|
||||
coerced = _coerce_int(usage.get(key))
|
||||
if coerced is not None:
|
||||
out[key] = coerced
|
||||
|
||||
response_metadata = getattr(message, "response_metadata", None) or {}
|
||||
token_usage = response_metadata.get("token_usage") if isinstance(response_metadata, dict) else None
|
||||
if isinstance(token_usage, dict):
|
||||
# Map both the LangChain-canonical names and OpenAI-native names; later
|
||||
# writes win so a more-precise provider-native value can override a
|
||||
# zero placeholder from usage_metadata.
|
||||
candidates = {
|
||||
"input_tokens": ("input_tokens", "prompt_tokens"),
|
||||
"output_tokens": ("output_tokens", "completion_tokens"),
|
||||
"total_tokens": ("total_tokens",),
|
||||
}
|
||||
for canonical, keys in candidates.items():
|
||||
for key in keys:
|
||||
coerced = _coerce_int(token_usage.get(key))
|
||||
if coerced is not None and coerced > 0:
|
||||
out[canonical] = coerced
|
||||
break
|
||||
|
||||
# Derive a missing or implausibly-zero field if the other two are present.
|
||||
# Some providers report ``input_tokens=0`` / ``output_tokens=0`` next to a
|
||||
# real ``total_tokens`` — treat those zeros as missing so we don't silently
|
||||
# store nonsense breakdowns.
|
||||
def _missing_or_zero(key: str) -> bool:
|
||||
return out.get(key) in (None, 0)
|
||||
|
||||
if "total_tokens" not in out and out.get("input_tokens") and out.get("output_tokens"):
|
||||
out["total_tokens"] = out["input_tokens"] + out["output_tokens"]
|
||||
if _missing_or_zero("output_tokens") and out.get("total_tokens") and out.get("input_tokens"):
|
||||
derived = out["total_tokens"] - out["input_tokens"]
|
||||
if derived > 0:
|
||||
out["output_tokens"] = derived
|
||||
if _missing_or_zero("input_tokens") and out.get("total_tokens") and out.get("output_tokens"):
|
||||
derived = out["total_tokens"] - out["output_tokens"]
|
||||
if derived > 0:
|
||||
out["input_tokens"] = derived
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def _classify_exception(exc: BaseException) -> tuple[str, str]:
|
||||
"""Return (error_type, error_message). Message is trimmed to 500 chars."""
|
||||
name = exc.__class__.__name__
|
||||
raw = (str(exc) or name).strip()
|
||||
return name, (raw[:500] if len(raw) > 500 else raw)
|
||||
|
||||
|
||||
class LlmMetricsMiddleware(AgentMiddleware[AgentState]):
|
||||
"""Wrap each model call to time it and persist a metric row."""
|
||||
|
||||
def __init__(self, *, store: LlmMetricsStore | None = None, **kwargs: Any) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self._store_override = store
|
||||
|
||||
def _store(self) -> LlmMetricsStore | None:
|
||||
if self._store_override is not None:
|
||||
return self._store_override
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
except Exception:
|
||||
return None
|
||||
return make_llm_metrics_store(sf)
|
||||
|
||||
def _build_base_record(self, request: ModelRequest) -> LlmMetricRecord:
|
||||
return LlmMetricRecord(
|
||||
id=uuid.uuid4().hex,
|
||||
created_at=_now(),
|
||||
user_id=get_effective_user_id(),
|
||||
thread_id=_safe_get_thread_id(request),
|
||||
run_id=_safe_get_run_id(request),
|
||||
agent_name=_safe_get_agent_name(request),
|
||||
model_name=_safe_get_model_name(request),
|
||||
)
|
||||
|
||||
def _finalize_success(self, record: LlmMetricRecord, response: ModelResponse, started: float) -> None:
|
||||
elapsed = max(0.0, time.monotonic() - started)
|
||||
record.duration_ms = int(elapsed * 1000)
|
||||
usage = _extract_usage(response)
|
||||
record.input_tokens = usage.get("input_tokens")
|
||||
record.output_tokens = usage.get("output_tokens")
|
||||
record.total_tokens = usage.get("total_tokens") or (
|
||||
(record.input_tokens or 0) + (record.output_tokens or 0) or None
|
||||
)
|
||||
if record.output_tokens is not None and elapsed > 0:
|
||||
record.tokens_per_sec = round(record.output_tokens / elapsed, 3)
|
||||
record.status = "success"
|
||||
|
||||
def _finalize_error(self, record: LlmMetricRecord, exc: BaseException, started: float) -> None:
|
||||
record.duration_ms = int(max(0.0, time.monotonic() - started) * 1000)
|
||||
record.status = "error"
|
||||
record.error_type, record.error_message = _classify_exception(exc)
|
||||
|
||||
def _flush(self, record: LlmMetricRecord) -> None:
|
||||
store = self._store()
|
||||
if store is None:
|
||||
return
|
||||
try:
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
loop = None
|
||||
if loop is not None and loop.is_running():
|
||||
# Fire-and-forget — never let metric IO block the request path.
|
||||
loop.create_task(store.record(record))
|
||||
else:
|
||||
asyncio.run(store.record(record))
|
||||
except Exception:
|
||||
logger.warning("LlmMetricsMiddleware: failed to flush metric %s", record.id, exc_info=True)
|
||||
|
||||
@override
|
||||
def wrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], ModelResponse],
|
||||
) -> ModelCallResult:
|
||||
record = self._build_base_record(request)
|
||||
started = time.monotonic()
|
||||
try:
|
||||
response = handler(request)
|
||||
except GraphBubbleUp:
|
||||
# LangGraph control flow — not a real LLM failure.
|
||||
raise
|
||||
except Exception as exc:
|
||||
self._finalize_error(record, exc, started)
|
||||
self._flush(record)
|
||||
raise
|
||||
self._finalize_success(record, response, started)
|
||||
self._flush(record)
|
||||
return response
|
||||
|
||||
@override
|
||||
async def awrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
||||
) -> ModelCallResult:
|
||||
record = self._build_base_record(request)
|
||||
started = time.monotonic()
|
||||
try:
|
||||
response = await handler(request)
|
||||
except GraphBubbleUp:
|
||||
raise
|
||||
except Exception as exc:
|
||||
self._finalize_error(record, exc, started)
|
||||
self._flush(record)
|
||||
raise
|
||||
self._finalize_success(record, response, started)
|
||||
self._flush(record)
|
||||
return response
|
||||
+8
@@ -96,6 +96,14 @@ def _build_runtime_middlewares(
|
||||
|
||||
middlewares.append(LLMErrorHandlingMiddleware(app_config=app_config))
|
||||
|
||||
# LLM call metrics — always on, even when the user-facing speed indicator
|
||||
# is disabled. We append *after* the retry middleware so the metrics
|
||||
# middleware sits deeper in the wrap chain and observes individual provider
|
||||
# attempts (including the failed ones).
|
||||
from deerflow.agents.middlewares.llm_metrics_middleware import LlmMetricsMiddleware
|
||||
|
||||
middlewares.append(LlmMetricsMiddleware())
|
||||
|
||||
# Guardrail middleware (if configured)
|
||||
guardrails_config = app_config.guardrails
|
||||
if guardrails_config.enabled and guardrails_config.provider:
|
||||
|
||||
+6
-1
@@ -45,7 +45,12 @@ class McpServerConfig(BaseModel):
|
||||
headers: dict[str, str] = Field(default_factory=dict, description="HTTP headers to send (for sse or http type)")
|
||||
oauth: McpOAuthConfig | None = Field(default=None, description="OAuth configuration (for sse or http type)")
|
||||
description: str = Field(default="", description="Human-readable description of what this MCP server provides")
|
||||
model_config = ConfigDict(extra="allow")
|
||||
disabled_tools: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="List of tool names (without server-name prefix) to hide from the agent",
|
||||
alias="disabledTools",
|
||||
)
|
||||
model_config = ConfigDict(extra="allow", populate_by_name=True)
|
||||
|
||||
|
||||
class SkillStateConfig(BaseModel):
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
"""System-wide runtime settings editable from the admin UI.
|
||||
|
||||
Backed by a single JSON file under the DeerFlow runtime home (default
|
||||
``.deer-flow/system_settings.json``). Unlike :mod:`extensions_config` this
|
||||
file is **not** part of the project source tree — it is per-deployment
|
||||
state, and it stays writable at runtime.
|
||||
|
||||
The only feature stored here at the moment is the **prompt prefix toggle**
|
||||
shown in the chat input: an administrator can configure a label and a
|
||||
prefix string, and any user can flip the toggle on/off; when enabled, the
|
||||
prefix is prepended to the first human message of a new thread.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from deerflow.config.runtime_paths import runtime_home
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SETTINGS_FILE = "system_settings.json"
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
class PromptPrefixSettings(BaseModel):
|
||||
"""Admin-configurable first-message prompt prefix."""
|
||||
|
||||
enabled: bool = Field(default=False, description="Whether the prefix toggle is offered to users at all")
|
||||
switch_label: str = Field(default="加入立场前缀", description="Toggle label shown in the chat input")
|
||||
prompt_prefix: str = Field(default="", description="Text prepended to the first user message when enabled")
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
|
||||
class SystemSettings(BaseModel):
|
||||
"""Top-level system settings document."""
|
||||
|
||||
prompt_prefix: PromptPrefixSettings = Field(default_factory=PromptPrefixSettings)
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
|
||||
def _settings_path() -> Path:
|
||||
home = runtime_home()
|
||||
home.mkdir(parents=True, exist_ok=True)
|
||||
return home / _SETTINGS_FILE
|
||||
|
||||
|
||||
def load_system_settings() -> SystemSettings:
|
||||
"""Read the on-disk system settings, returning defaults when missing."""
|
||||
path = _settings_path()
|
||||
if not path.is_file():
|
||||
return SystemSettings()
|
||||
try:
|
||||
raw = json.loads(path.read_text(encoding="utf-8"))
|
||||
if not isinstance(raw, dict):
|
||||
logger.warning("System settings file %s is not a JSON object; ignoring", path)
|
||||
return SystemSettings()
|
||||
return SystemSettings.model_validate(raw)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to read system settings from %s: %s", path, exc)
|
||||
return SystemSettings()
|
||||
|
||||
|
||||
def save_system_settings(settings: SystemSettings) -> SystemSettings:
|
||||
"""Persist the settings document atomically, returning the saved value."""
|
||||
path = _settings_path()
|
||||
payload = settings.model_dump(mode="json")
|
||||
with _lock:
|
||||
tmp = path.with_suffix(path.suffix + ".tmp")
|
||||
tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
tmp.replace(path)
|
||||
return settings
|
||||
@@ -121,6 +121,25 @@ async def get_mcp_tools() -> list[BaseTool]:
|
||||
|
||||
# Get all tools from all servers
|
||||
tools = await client.get_tools()
|
||||
|
||||
# Filter out tools listed in each server's `disabled_tools`. Tool names from
|
||||
# MultiServerMCPClient are prefixed with the server name (because we pass
|
||||
# `tool_name_prefix=True`), so we compare against `<server>_<tool>` here.
|
||||
disabled_full_names: set[str] = set()
|
||||
for server_name, server_cfg in extensions_config.get_enabled_mcp_servers().items():
|
||||
for tool_name in server_cfg.disabled_tools or []:
|
||||
name = (tool_name or "").strip()
|
||||
if not name:
|
||||
continue
|
||||
disabled_full_names.add(f"{server_name}_{name}")
|
||||
disabled_full_names.add(name)
|
||||
if disabled_full_names:
|
||||
kept = [t for t in tools if t.name not in disabled_full_names]
|
||||
removed = len(tools) - len(kept)
|
||||
if removed:
|
||||
logger.info(f"Filtered out {removed} disabled MCP tool(s)")
|
||||
tools = kept
|
||||
|
||||
logger.info(f"Successfully loaded {len(tools)} tool(s) from MCP servers")
|
||||
|
||||
# Patch tools to support sync invocation, as deerflow client streams synchronously
|
||||
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
"""Persistence layer for per-LLM-call metrics.
|
||||
|
||||
Each completed LLM invocation (success or failure) writes a single row so
|
||||
administrators can audit usage and speed. Designed to be cheap to write —
|
||||
the middleware fires once per provider call and never blocks streaming.
|
||||
"""
|
||||
|
||||
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsStore
|
||||
from deerflow.persistence.llm_metrics.model import LlmCallMetricRow
|
||||
from deerflow.persistence.llm_metrics.sql import SqlLlmMetricsStore
|
||||
|
||||
__all__ = [
|
||||
"LlmCallMetricRow",
|
||||
"LlmMetricRecord",
|
||||
"LlmMetricsStore",
|
||||
"SqlLlmMetricsStore",
|
||||
"make_llm_metrics_store",
|
||||
]
|
||||
|
||||
|
||||
def make_llm_metrics_store(session_factory) -> LlmMetricsStore | None:
|
||||
"""Return a SQL-backed metrics store, or ``None`` when no DB is configured."""
|
||||
if session_factory is None:
|
||||
return None
|
||||
return SqlLlmMetricsStore(session_factory)
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
"""Abstract repository contract for LLM call metrics."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class LlmMetricRecord:
|
||||
"""In-memory shape of a metric row before persistence."""
|
||||
|
||||
id: str
|
||||
created_at: datetime
|
||||
user_id: str | None = None
|
||||
thread_id: str | None = None
|
||||
run_id: str | None = None
|
||||
agent_name: str | None = None
|
||||
model_name: str | None = None
|
||||
duration_ms: int = 0
|
||||
input_tokens: int | None = None
|
||||
output_tokens: int | None = None
|
||||
total_tokens: int | None = None
|
||||
tokens_per_sec: float | None = None
|
||||
status: str = "success"
|
||||
error_type: str | None = None
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class LlmMetricsQuery:
|
||||
"""Filter set for the admin list / export endpoints."""
|
||||
|
||||
user_id: str | None = None
|
||||
model_name: str | None = None
|
||||
status: str | None = None
|
||||
since: datetime | None = None
|
||||
until: datetime | None = None
|
||||
limit: int = 100
|
||||
offset: int = 0
|
||||
|
||||
|
||||
class LlmMetricsStore(ABC):
|
||||
"""Persistence contract for the metrics middleware + admin API."""
|
||||
|
||||
@abstractmethod
|
||||
async def record(self, metric: LlmMetricRecord) -> None:
|
||||
"""Persist one metric row. Must never raise on best-effort writes."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def list(self, query: LlmMetricsQuery) -> list[dict[str, Any]]:
|
||||
"""Return the matching rows (newest first), paginated by ``query``."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def count(self, query: LlmMetricsQuery) -> int:
|
||||
"""Return the unpaginated row count for ``query``."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def iter_all(self, query: LlmMetricsQuery):
|
||||
"""Yield rows matching ``query`` without pagination — used by export."""
|
||||
raise NotImplementedError
|
||||
+62
@@ -0,0 +1,62 @@
|
||||
"""ORM model for per-LLM-call metrics."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from sqlalchemy import DateTime, Float, Integer, String, Text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from deerflow.persistence.base import Base
|
||||
|
||||
|
||||
class LlmCallMetricRow(Base):
|
||||
"""One row per completed (or failed) LLM provider invocation."""
|
||||
|
||||
__tablename__ = "llm_call_metrics"
|
||||
|
||||
id: Mapped[str] = mapped_column(String(64), primary_key=True)
|
||||
# Wall-clock time the LLM call ended. Indexed because admin views always
|
||||
# sort/filter by recency.
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True)
|
||||
|
||||
user_id: Mapped[str | None] = mapped_column(String(64), index=True)
|
||||
thread_id: Mapped[str | None] = mapped_column(String(64), index=True)
|
||||
run_id: Mapped[str | None] = mapped_column(String(64), index=True)
|
||||
agent_name: Mapped[str | None] = mapped_column(String(128), index=True)
|
||||
model_name: Mapped[str | None] = mapped_column(String(128), index=True)
|
||||
|
||||
# Wall-clock duration of the provider invocation in milliseconds. Always set.
|
||||
duration_ms: Mapped[int] = mapped_column(Integer, default=0)
|
||||
|
||||
input_tokens: Mapped[int | None] = mapped_column(Integer)
|
||||
output_tokens: Mapped[int | None] = mapped_column(Integer)
|
||||
total_tokens: Mapped[int | None] = mapped_column(Integer)
|
||||
|
||||
# Output tokens / seconds; null when there is no completion (errors).
|
||||
tokens_per_sec: Mapped[float | None] = mapped_column(Float)
|
||||
|
||||
# "success" or "error". Errors store the classification and original message
|
||||
# so admins can scan for transient vs auth/quota issues.
|
||||
status: Mapped[str] = mapped_column(String(16), default="success", index=True)
|
||||
error_type: Mapped[str | None] = mapped_column(String(64))
|
||||
error_message: Mapped[str | None] = mapped_column(Text)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"id": self.id,
|
||||
"created_at": self.created_at.isoformat() if isinstance(self.created_at, datetime) else self.created_at,
|
||||
"user_id": self.user_id,
|
||||
"thread_id": self.thread_id,
|
||||
"run_id": self.run_id,
|
||||
"agent_name": self.agent_name,
|
||||
"model_name": self.model_name,
|
||||
"duration_ms": self.duration_ms,
|
||||
"input_tokens": self.input_tokens,
|
||||
"output_tokens": self.output_tokens,
|
||||
"total_tokens": self.total_tokens,
|
||||
"tokens_per_sec": self.tokens_per_sec,
|
||||
"status": self.status,
|
||||
"error_type": self.error_type,
|
||||
"error_message": self.error_message,
|
||||
}
|
||||
+109
@@ -0,0 +1,109 @@
|
||||
"""SQLAlchemy-backed LLM metrics repository."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import desc, func, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsQuery, LlmMetricsStore
|
||||
from deerflow.persistence.llm_metrics.model import LlmCallMetricRow
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _row_to_dict(row: LlmCallMetricRow) -> dict[str, Any]:
|
||||
return row.to_dict()
|
||||
|
||||
|
||||
def _normalize_created_at(value: datetime) -> datetime:
|
||||
return value if value.tzinfo else value.replace(tzinfo=UTC)
|
||||
|
||||
|
||||
class SqlLlmMetricsStore(LlmMetricsStore):
|
||||
"""SQLite/Postgres-backed implementation; writes are intentionally fire-and-forget."""
|
||||
|
||||
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
||||
self._sf = session_factory
|
||||
|
||||
async def record(self, metric: LlmMetricRecord) -> None:
|
||||
row = LlmCallMetricRow(
|
||||
id=metric.id,
|
||||
created_at=_normalize_created_at(metric.created_at),
|
||||
user_id=metric.user_id,
|
||||
thread_id=metric.thread_id,
|
||||
run_id=metric.run_id,
|
||||
agent_name=metric.agent_name,
|
||||
model_name=metric.model_name,
|
||||
duration_ms=metric.duration_ms,
|
||||
input_tokens=metric.input_tokens,
|
||||
output_tokens=metric.output_tokens,
|
||||
total_tokens=metric.total_tokens,
|
||||
tokens_per_sec=metric.tokens_per_sec,
|
||||
status=metric.status,
|
||||
error_type=metric.error_type,
|
||||
error_message=metric.error_message,
|
||||
)
|
||||
try:
|
||||
async with self._sf() as session:
|
||||
session.add(row)
|
||||
await session.commit()
|
||||
except Exception:
|
||||
# Metrics are best-effort. Never let a persistence failure surface
|
||||
# to the agent runtime — log and move on.
|
||||
logger.warning("Failed to persist llm_call_metrics row %s", metric.id, exc_info=True)
|
||||
|
||||
def _apply_filters(self, stmt, query: LlmMetricsQuery):
|
||||
if query.user_id:
|
||||
stmt = stmt.where(LlmCallMetricRow.user_id == query.user_id)
|
||||
if query.model_name:
|
||||
stmt = stmt.where(LlmCallMetricRow.model_name == query.model_name)
|
||||
if query.status:
|
||||
stmt = stmt.where(LlmCallMetricRow.status == query.status)
|
||||
if query.since is not None:
|
||||
stmt = stmt.where(LlmCallMetricRow.created_at >= _normalize_created_at(query.since))
|
||||
if query.until is not None:
|
||||
stmt = stmt.where(LlmCallMetricRow.created_at <= _normalize_created_at(query.until))
|
||||
return stmt
|
||||
|
||||
async def list(self, query: LlmMetricsQuery) -> list[dict[str, Any]]:
|
||||
stmt = select(LlmCallMetricRow)
|
||||
stmt = self._apply_filters(stmt, query)
|
||||
stmt = stmt.order_by(desc(LlmCallMetricRow.created_at)).limit(max(1, min(query.limit, 1000))).offset(max(0, query.offset))
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(stmt)
|
||||
return [_row_to_dict(row) for row in result.scalars()]
|
||||
|
||||
async def count(self, query: LlmMetricsQuery) -> int:
|
||||
stmt = select(func.count()).select_from(LlmCallMetricRow)
|
||||
stmt = self._apply_filters(stmt, query)
|
||||
async with self._sf() as session:
|
||||
return (await session.scalar(stmt)) or 0
|
||||
|
||||
async def iter_all(self, query: LlmMetricsQuery) -> AsyncIterator[dict[str, Any]]:
|
||||
# Page through the result so the Excel export stays bounded in memory
|
||||
# for large date windows.
|
||||
page_size = 500
|
||||
offset = 0
|
||||
while True:
|
||||
page_query = LlmMetricsQuery(
|
||||
user_id=query.user_id,
|
||||
model_name=query.model_name,
|
||||
status=query.status,
|
||||
since=query.since,
|
||||
until=query.until,
|
||||
limit=page_size,
|
||||
offset=offset,
|
||||
)
|
||||
page = await self.list(page_query)
|
||||
if not page:
|
||||
return
|
||||
for item in page:
|
||||
yield item
|
||||
if len(page) < page_size:
|
||||
return
|
||||
offset += page_size
|
||||
+59
@@ -0,0 +1,59 @@
|
||||
"""Create llm_call_metrics table.
|
||||
|
||||
Revision ID: 20260514_02
|
||||
Revises: 20260514_01
|
||||
Create Date: 2026-05-14
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy import inspect
|
||||
|
||||
|
||||
revision = "20260514_02"
|
||||
down_revision = "20260514_01"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _has_table(table_name: str) -> bool:
|
||||
inspector = inspect(op.get_bind())
|
||||
return table_name in inspector.get_table_names()
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if _has_table("llm_call_metrics"):
|
||||
return
|
||||
|
||||
op.create_table(
|
||||
"llm_call_metrics",
|
||||
sa.Column("id", sa.String(64), primary_key=True),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("user_id", sa.String(64), nullable=True),
|
||||
sa.Column("thread_id", sa.String(64), nullable=True),
|
||||
sa.Column("run_id", sa.String(64), nullable=True),
|
||||
sa.Column("agent_name", sa.String(128), nullable=True),
|
||||
sa.Column("model_name", sa.String(128), nullable=True),
|
||||
sa.Column("duration_ms", sa.Integer(), nullable=False, server_default="0"),
|
||||
sa.Column("input_tokens", sa.Integer(), nullable=True),
|
||||
sa.Column("output_tokens", sa.Integer(), nullable=True),
|
||||
sa.Column("total_tokens", sa.Integer(), nullable=True),
|
||||
sa.Column("tokens_per_sec", sa.Float(), nullable=True),
|
||||
sa.Column("status", sa.String(16), nullable=False, server_default="success"),
|
||||
sa.Column("error_type", sa.String(64), nullable=True),
|
||||
sa.Column("error_message", sa.Text(), nullable=True),
|
||||
)
|
||||
op.create_index("ix_llm_call_metrics_created_at", "llm_call_metrics", ["created_at"])
|
||||
op.create_index("ix_llm_call_metrics_user_id", "llm_call_metrics", ["user_id"])
|
||||
op.create_index("ix_llm_call_metrics_model_name", "llm_call_metrics", ["model_name"])
|
||||
op.create_index("ix_llm_call_metrics_status", "llm_call_metrics", ["status"])
|
||||
op.create_index("ix_llm_call_metrics_thread_id", "llm_call_metrics", ["thread_id"])
|
||||
op.create_index("ix_llm_call_metrics_run_id", "llm_call_metrics", ["run_id"])
|
||||
op.create_index("ix_llm_call_metrics_agent_name", "llm_call_metrics", ["agent_name"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if _has_table("llm_call_metrics"):
|
||||
op.drop_table("llm_call_metrics")
|
||||
+4
@@ -17,7 +17,9 @@ there is no matching entity directory.
|
||||
|
||||
from deerflow.persistence.agents.model import AgentRow
|
||||
from deerflow.persistence.feedback.model import FeedbackRow
|
||||
from deerflow.persistence.llm_metrics.model import LlmCallMetricRow
|
||||
from deerflow.persistence.models.run_event import RunEventRow
|
||||
from deerflow.persistence.recommended_questions.model import RecommendedQuestionRow
|
||||
from deerflow.persistence.run.model import RunRow
|
||||
from deerflow.persistence.scheduled_tasks.model import (
|
||||
ScheduledTaskDeliveryProfileRow,
|
||||
@@ -32,6 +34,8 @@ from deerflow.persistence.user.model import UserRow
|
||||
__all__ = [
|
||||
"FeedbackRow",
|
||||
"AgentRow",
|
||||
"LlmCallMetricRow",
|
||||
"RecommendedQuestionRow",
|
||||
"RunEventRow",
|
||||
"RunRow",
|
||||
"ScheduledTaskDeliveryProfileRow",
|
||||
|
||||
+10
-2
@@ -247,7 +247,11 @@ class ScheduledTaskService:
|
||||
accessible = await self._accessible_task(user_id, task_id)
|
||||
if accessible is None:
|
||||
return []
|
||||
return await self.store.list_runs(task_id, accessible["user_id"], limit=limit)
|
||||
runs = await self.store.list_runs(task_id, accessible["user_id"], limit=limit)
|
||||
can_edit = accessible.get("role") == "owner"
|
||||
for run in runs:
|
||||
run["can_edit"] = can_edit
|
||||
return runs
|
||||
|
||||
async def get_run(self, user_id: str, run_id: str) -> dict[str, Any] | None:
|
||||
run = await self.store.get_run(run_id, user_id)
|
||||
@@ -316,7 +320,7 @@ class ScheduledTaskService:
|
||||
except ValueError as exc:
|
||||
raise PermissionError("Refusing to edit file outside the run archive") from exc
|
||||
|
||||
archive_path.write_text(new_content, encoding="utf-8")
|
||||
archive_path.write_text(new_content, encoding="utf-8", newline="")
|
||||
new_size = archive_path.stat().st_size
|
||||
|
||||
# Reflect the new size in the run row so subsequent list/preview calls
|
||||
@@ -342,6 +346,10 @@ class ScheduledTaskService:
|
||||
shutil.rmtree(archive_dir)
|
||||
except Exception:
|
||||
logger.warning("Failed to delete scheduled task run archive %s", archive_dir, exc_info=True)
|
||||
return True
|
||||
candidate = await self.store.get_run_any(run_id)
|
||||
if candidate is not None and await self.store.get_subscription(candidate["task_id"], user_id) is not None:
|
||||
raise PermissionError("Only the run owner may delete scheduled task runs")
|
||||
return deleted
|
||||
|
||||
async def get_delivery_profile(self, user_id: str) -> dict[str, Any]:
|
||||
|
||||
@@ -0,0 +1,321 @@
|
||||
"""Tests for the LLM call metrics middleware + repository.
|
||||
|
||||
Covers:
|
||||
- Successful calls emit a row with duration, token counts and tokens_per_sec.
|
||||
- Failed calls record the exception class as ``error_type`` plus a trimmed
|
||||
``error_message``.
|
||||
- The middleware never blocks on a backend exception (best-effort writes).
|
||||
- The SQL store applies filters and pagination for the admin list endpoint.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from langchain_core.messages import AIMessage
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
|
||||
from deerflow.agents.middlewares.llm_metrics_middleware import (
|
||||
LlmMetricsMiddleware,
|
||||
_extract_usage,
|
||||
)
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.llm_metrics import SqlLlmMetricsStore
|
||||
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsQuery
|
||||
|
||||
|
||||
class _RecordingStore:
|
||||
"""Minimal in-memory store used to assert what the middleware persists."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.records: list[LlmMetricRecord] = []
|
||||
|
||||
async def record(self, metric: LlmMetricRecord) -> None:
|
||||
self.records.append(metric)
|
||||
|
||||
async def list(self, query: LlmMetricsQuery):
|
||||
return [r.__dict__ for r in self.records]
|
||||
|
||||
async def count(self, query: LlmMetricsQuery) -> int:
|
||||
return len(self.records)
|
||||
|
||||
async def iter_all(self, query: LlmMetricsQuery):
|
||||
for record in self.records:
|
||||
yield record.__dict__
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def sql_store(tmp_path):
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'metrics.db'}")
|
||||
async with engine.begin() as conn:
|
||||
# Importing the package registers the LlmCallMetricRow with Base.metadata.
|
||||
import deerflow.persistence.llm_metrics.model # noqa: F401
|
||||
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
store = SqlLlmMetricsStore(sf)
|
||||
try:
|
||||
yield store
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
def _make_request():
|
||||
return type(
|
||||
"FakeRequest",
|
||||
(),
|
||||
{
|
||||
"model": type("FakeModel", (), {"model_name": "gpt-test"}),
|
||||
"runtime": type(
|
||||
"FakeRuntime",
|
||||
(),
|
||||
{
|
||||
"config": {"configurable": {"thread_id": "t1"}},
|
||||
"context": {"agent_name": "lead_agent", "run_id": "r1"},
|
||||
},
|
||||
),
|
||||
"state": {"thread_id": "t1"},
|
||||
},
|
||||
)()
|
||||
|
||||
|
||||
def test_middleware_records_success_with_speed():
|
||||
import time as _time
|
||||
|
||||
store = _RecordingStore()
|
||||
mw = LlmMetricsMiddleware(store=store)
|
||||
response = AIMessage(
|
||||
content="hi",
|
||||
usage_metadata={"input_tokens": 5, "output_tokens": 20, "total_tokens": 25},
|
||||
)
|
||||
|
||||
def slow_handler(_req):
|
||||
# Force a measurable elapsed window so tokens_per_sec is populated.
|
||||
# 50ms is comfortably above Windows monotonic-clock granularity.
|
||||
_time.sleep(0.05)
|
||||
return response
|
||||
|
||||
result = mw.wrap_model_call(_make_request(), slow_handler)
|
||||
|
||||
assert result is response
|
||||
# awrap_model_call uses a fire-and-forget task; here the sync path runs
|
||||
# asyncio.run synchronously, so the record landed already.
|
||||
assert len(store.records) == 1
|
||||
record = store.records[0]
|
||||
assert record.status == "success"
|
||||
assert record.model_name == "gpt-test"
|
||||
assert record.thread_id == "t1"
|
||||
assert record.run_id == "r1"
|
||||
assert record.input_tokens == 5
|
||||
assert record.output_tokens == 20
|
||||
assert record.total_tokens == 25
|
||||
# tokens_per_sec may be very high in tests because the handler is
|
||||
# ~instant; we only assert that the field is populated.
|
||||
assert record.tokens_per_sec is not None
|
||||
assert record.duration_ms >= 0
|
||||
|
||||
|
||||
def test_middleware_records_error_with_classification():
|
||||
store = _RecordingStore()
|
||||
mw = LlmMetricsMiddleware(store=store)
|
||||
|
||||
class BoomError(RuntimeError):
|
||||
pass
|
||||
|
||||
def explode(_req):
|
||||
raise BoomError("provider exploded")
|
||||
|
||||
with pytest.raises(BoomError):
|
||||
mw.wrap_model_call(_make_request(), explode)
|
||||
|
||||
assert len(store.records) == 1
|
||||
record = store.records[0]
|
||||
assert record.status == "error"
|
||||
assert record.error_type == "BoomError"
|
||||
assert record.error_message == "provider exploded"
|
||||
assert record.tokens_per_sec is None
|
||||
|
||||
|
||||
def test_middleware_swallows_store_errors():
|
||||
"""A persistence outage must not break the agent runtime."""
|
||||
|
||||
class FaultyStore:
|
||||
async def record(self, metric):
|
||||
raise RuntimeError("disk gone")
|
||||
|
||||
async def list(self, q):
|
||||
return []
|
||||
|
||||
async def count(self, q):
|
||||
return 0
|
||||
|
||||
async def iter_all(self, q):
|
||||
if False:
|
||||
yield {}
|
||||
|
||||
mw = LlmMetricsMiddleware(store=FaultyStore())
|
||||
# The wrap returns normally even though _flush would raise; the call below
|
||||
# must not raise.
|
||||
response = AIMessage(
|
||||
content="ok",
|
||||
usage_metadata={"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
)
|
||||
out = mw.wrap_model_call(_make_request(), lambda req: response)
|
||||
assert out is response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sql_store_round_trip(sql_store):
|
||||
base = datetime.now(UTC)
|
||||
for i in range(3):
|
||||
await sql_store.record(
|
||||
LlmMetricRecord(
|
||||
id=uuid.uuid4().hex,
|
||||
created_at=base - timedelta(minutes=i),
|
||||
user_id="alice",
|
||||
model_name="gpt-test" if i % 2 == 0 else "gpt-other",
|
||||
duration_ms=100 * (i + 1),
|
||||
output_tokens=10 * (i + 1),
|
||||
tokens_per_sec=10.0 * (i + 1),
|
||||
status="success",
|
||||
)
|
||||
)
|
||||
rows = await sql_store.list(LlmMetricsQuery(limit=10))
|
||||
assert len(rows) == 3
|
||||
# Newest first.
|
||||
assert rows[0]["duration_ms"] == 100
|
||||
|
||||
filtered = await sql_store.list(LlmMetricsQuery(model_name="gpt-test", limit=10))
|
||||
assert all(r["model_name"] == "gpt-test" for r in filtered)
|
||||
assert len(filtered) == 2
|
||||
|
||||
assert await sql_store.count(LlmMetricsQuery(model_name="gpt-test")) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sql_store_filters_errors_only(sql_store):
|
||||
base = datetime.now(UTC)
|
||||
await sql_store.record(
|
||||
LlmMetricRecord(id=uuid.uuid4().hex, created_at=base, status="success", duration_ms=10, output_tokens=1)
|
||||
)
|
||||
await sql_store.record(
|
||||
LlmMetricRecord(
|
||||
id=uuid.uuid4().hex,
|
||||
created_at=base,
|
||||
status="error",
|
||||
duration_ms=5,
|
||||
error_type="APITimeoutError",
|
||||
error_message="timed out",
|
||||
)
|
||||
)
|
||||
only_errors = await sql_store.list(LlmMetricsQuery(status="error", limit=10))
|
||||
assert len(only_errors) == 1
|
||||
assert only_errors[0]["error_type"] == "APITimeoutError"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sql_store_iter_all_paginates(sql_store):
|
||||
base = datetime.now(UTC)
|
||||
for i in range(1200):
|
||||
await sql_store.record(
|
||||
LlmMetricRecord(
|
||||
id=uuid.uuid4().hex,
|
||||
created_at=base - timedelta(seconds=i),
|
||||
duration_ms=1,
|
||||
status="success",
|
||||
)
|
||||
)
|
||||
seen = 0
|
||||
async for _ in sql_store.iter_all(LlmMetricsQuery()):
|
||||
seen += 1
|
||||
assert seen == 1200
|
||||
|
||||
|
||||
def test_extract_usage_prefers_usage_metadata():
|
||||
response = AIMessage(
|
||||
content="x",
|
||||
usage_metadata={"input_tokens": 3, "output_tokens": 7, "total_tokens": 10},
|
||||
)
|
||||
out = _extract_usage(response)
|
||||
assert out == {"input_tokens": 3, "output_tokens": 7, "total_tokens": 10}
|
||||
|
||||
|
||||
def test_extract_usage_falls_back_to_response_metadata_token_usage():
|
||||
"""DeepSeek / vLLM-style: only total in usage_metadata, breakdown elsewhere."""
|
||||
response = AIMessage(
|
||||
content="x",
|
||||
usage_metadata={"input_tokens": 0, "output_tokens": 0, "total_tokens": 11093},
|
||||
response_metadata={
|
||||
"token_usage": {
|
||||
"prompt_tokens": 9000,
|
||||
"completion_tokens": 2093,
|
||||
"total_tokens": 11093,
|
||||
},
|
||||
},
|
||||
)
|
||||
out = _extract_usage(response)
|
||||
assert out["input_tokens"] == 9000
|
||||
assert out["output_tokens"] == 2093
|
||||
assert out["total_tokens"] == 11093
|
||||
|
||||
|
||||
def test_extract_usage_unwraps_model_response():
|
||||
"""The handler hands us a ModelResponse dataclass, not a bare AIMessage."""
|
||||
from langchain.agents.middleware.types import ModelResponse
|
||||
|
||||
inner = AIMessage(
|
||||
content="hi",
|
||||
usage_metadata={"input_tokens": 7, "output_tokens": 9, "total_tokens": 16},
|
||||
)
|
||||
response = ModelResponse(result=[inner])
|
||||
out = _extract_usage(response)
|
||||
assert out == {"input_tokens": 7, "output_tokens": 9, "total_tokens": 16}
|
||||
|
||||
|
||||
def test_extract_usage_unwraps_extended_model_response():
|
||||
"""ExtendedModelResponse wraps a ModelResponse in a ``model_response`` field."""
|
||||
from langchain.agents.middleware.types import ExtendedModelResponse, ModelResponse
|
||||
|
||||
inner = AIMessage(
|
||||
content="hi",
|
||||
usage_metadata={"input_tokens": 1, "output_tokens": 2, "total_tokens": 3},
|
||||
)
|
||||
extended = ExtendedModelResponse(model_response=ModelResponse(result=[inner]))
|
||||
out = _extract_usage(extended)
|
||||
assert out["output_tokens"] == 2
|
||||
|
||||
|
||||
def test_extract_usage_derives_missing_field_when_two_present():
|
||||
response = AIMessage(
|
||||
content="x",
|
||||
usage_metadata={"input_tokens": 9000, "output_tokens": 0, "total_tokens": 11093},
|
||||
)
|
||||
out = _extract_usage(response)
|
||||
# output should be derived as total - input = 2093.
|
||||
assert out["output_tokens"] == 2093
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_middleware_records_success():
|
||||
store = _RecordingStore()
|
||||
mw = LlmMetricsMiddleware(store=store)
|
||||
|
||||
async def handler(_req):
|
||||
return AIMessage(
|
||||
content="hi",
|
||||
usage_metadata={"input_tokens": 1, "output_tokens": 4, "total_tokens": 5},
|
||||
)
|
||||
|
||||
response = await mw.awrap_model_call(_make_request(), handler)
|
||||
assert isinstance(response, AIMessage)
|
||||
|
||||
# awrap_model_call schedules the write as an asyncio.task; give the loop a
|
||||
# chance to drain.
|
||||
await asyncio.sleep(0)
|
||||
assert len(store.records) == 1
|
||||
assert store.records[0].status == "success"
|
||||
assert store.records[0].output_tokens == 4
|
||||
@@ -0,0 +1,108 @@
|
||||
"""Tests for the admin-configurable prompt-prefix toggle."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
from app.gateway.services import _prepend_prompt_prefix, _prompt_prefix_opted_in
|
||||
from deerflow.config import system_settings as ss
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tmp_runtime_home(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Path:
|
||||
"""Redirect ``runtime_home()`` to a per-test tmp dir."""
|
||||
|
||||
def fake_home() -> Path:
|
||||
return tmp_path
|
||||
|
||||
monkeypatch.setattr(ss, "runtime_home", fake_home)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def test_load_returns_defaults_when_missing(tmp_runtime_home: Path) -> None:
|
||||
settings = ss.load_system_settings()
|
||||
assert settings.prompt_prefix.enabled is False
|
||||
assert settings.prompt_prefix.prompt_prefix == ""
|
||||
|
||||
|
||||
def test_save_then_load_roundtrips(tmp_runtime_home: Path) -> None:
|
||||
settings = ss.SystemSettings(
|
||||
prompt_prefix=ss.PromptPrefixSettings(
|
||||
enabled=True,
|
||||
switch_label="加入立场前缀",
|
||||
prompt_prefix="你需要站在中方角度",
|
||||
)
|
||||
)
|
||||
ss.save_system_settings(settings)
|
||||
|
||||
file_path = tmp_runtime_home / "system_settings.json"
|
||||
assert file_path.is_file()
|
||||
saved = json.loads(file_path.read_text(encoding="utf-8"))
|
||||
assert saved["prompt_prefix"]["prompt_prefix"] == "你需要站在中方角度"
|
||||
|
||||
reloaded = ss.load_system_settings()
|
||||
assert reloaded.prompt_prefix.enabled is True
|
||||
assert reloaded.prompt_prefix.switch_label == "加入立场前缀"
|
||||
assert reloaded.prompt_prefix.prompt_prefix == "你需要站在中方角度"
|
||||
|
||||
|
||||
def test_prepend_prefix_on_string_content() -> None:
|
||||
graph_input = {"messages": [HumanMessage(content="评价一下中美贸易")]}
|
||||
out = _prepend_prompt_prefix(graph_input, "你需要站在中方角度")
|
||||
msg = out["messages"][0]
|
||||
assert isinstance(msg, HumanMessage)
|
||||
assert msg.content == "你需要站在中方角度\n\n评价一下中美贸易"
|
||||
|
||||
|
||||
def test_prepend_prefix_on_block_content() -> None:
|
||||
graph_input = {
|
||||
"messages": [
|
||||
HumanMessage(content=[{"type": "text", "text": "评价一下中美贸易"}])
|
||||
]
|
||||
}
|
||||
out = _prepend_prompt_prefix(graph_input, "你需要站在中方角度")
|
||||
blocks = out["messages"][0].content
|
||||
assert isinstance(blocks, list)
|
||||
assert blocks[0] == {"type": "text", "text": "你需要站在中方角度"}
|
||||
assert blocks[1] == {"type": "text", "text": "评价一下中美贸易"}
|
||||
|
||||
|
||||
def test_prepend_prefix_only_touches_first_human() -> None:
|
||||
graph_input = {
|
||||
"messages": [
|
||||
AIMessage(content="hi"),
|
||||
HumanMessage(content="q1"),
|
||||
HumanMessage(content="q2"),
|
||||
]
|
||||
}
|
||||
out = _prepend_prompt_prefix(graph_input, "PREFIX")
|
||||
# first human is index 1; index 2 must remain untouched
|
||||
assert out["messages"][1].content == "PREFIX\n\nq1"
|
||||
assert out["messages"][2].content == "q2"
|
||||
|
||||
|
||||
def test_prepend_prefix_no_messages_is_noop() -> None:
|
||||
assert _prepend_prompt_prefix({"messages": []}, "X") == {"messages": []}
|
||||
assert _prepend_prompt_prefix({}, "X") == {}
|
||||
|
||||
|
||||
def test_opted_in_via_context() -> None:
|
||||
body = SimpleNamespace(context={"prompt_prefix_enabled": True}, metadata=None)
|
||||
assert _prompt_prefix_opted_in(body) is True
|
||||
|
||||
|
||||
def test_opted_in_via_metadata() -> None:
|
||||
body = SimpleNamespace(context=None, metadata={"prompt_prefix_enabled": True})
|
||||
assert _prompt_prefix_opted_in(body) is True
|
||||
|
||||
|
||||
def test_not_opted_in_by_default() -> None:
|
||||
body = SimpleNamespace(context=None, metadata=None)
|
||||
assert _prompt_prefix_opted_in(body) is False
|
||||
body2 = SimpleNamespace(context={"mode": "ultra"}, metadata={})
|
||||
assert _prompt_prefix_opted_in(body2) is False
|
||||
@@ -147,13 +147,23 @@ async def test_subscriber_can_read_runs_and_get_run(store):
|
||||
)
|
||||
assert run is not None
|
||||
|
||||
alice_runs = await service.list_runs("alice", "t1", limit=10)
|
||||
assert alice_runs[0]["can_edit"] is True
|
||||
|
||||
# Bob (subscriber) sees runs on the task.
|
||||
bob_runs = await service.list_runs("bob", "t1", limit=10)
|
||||
assert [r["id"] for r in bob_runs] == ["run-1"]
|
||||
assert bob_runs[0]["can_edit"] is False
|
||||
|
||||
# Bob can fetch the run by id.
|
||||
fetched = await service.get_run("bob", "run-1")
|
||||
assert fetched is not None and fetched["id"] == "run-1"
|
||||
assert fetched["can_edit"] is False
|
||||
|
||||
# Bob can read subscribed records, but cannot delete Alice's record.
|
||||
with pytest.raises(PermissionError):
|
||||
await service.delete_run("bob", "run-1")
|
||||
assert await store.get_run("run-1", "alice") is not None
|
||||
|
||||
# Carol (no subscription) cannot.
|
||||
assert await service.list_runs("carol", "t1", limit=10) == []
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
"""Tests for the username-based login domain handling.
|
||||
|
||||
Verifies that:
|
||||
- The default mail domain is ``cm.com`` (with env-var override still honoured).
|
||||
- ``_normalize_username_local_part`` distinguishes raw-email vs username input.
|
||||
- The login flow falls back to a legacy-domain lookup so the same username
|
||||
keeps mapping to the same account after the default domain changes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from datetime import UTC, datetime
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from app.gateway.auth.models import User
|
||||
from app.gateway.auth.repositories.base import UserRepository
|
||||
from app.gateway.routers.auth import (
|
||||
_normalize_username_local_part,
|
||||
_username_to_email,
|
||||
)
|
||||
|
||||
|
||||
class FakeUserRepository(UserRepository):
|
||||
"""In-memory repo just enough to exercise the local-part lookup."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.rows: dict[str, User] = {}
|
||||
self.create_calls: list[User] = []
|
||||
|
||||
async def create_user(self, user: User) -> User:
|
||||
for existing in self.rows.values():
|
||||
if existing.email == user.email:
|
||||
raise ValueError(f"Email already registered: {user.email}")
|
||||
u = User(
|
||||
id=user.id or uuid4(),
|
||||
email=user.email,
|
||||
password_hash=user.password_hash,
|
||||
system_role=user.system_role,
|
||||
created_at=user.created_at or datetime.now(UTC),
|
||||
oauth_provider=user.oauth_provider,
|
||||
oauth_id=user.oauth_id,
|
||||
needs_setup=user.needs_setup,
|
||||
token_version=user.token_version,
|
||||
)
|
||||
self.rows[str(u.id)] = u
|
||||
self.create_calls.append(u)
|
||||
return u
|
||||
|
||||
async def get_user_by_id(self, user_id: str) -> User | None:
|
||||
return self.rows.get(user_id)
|
||||
|
||||
async def get_user_by_email(self, email: str) -> User | None:
|
||||
for user in self.rows.values():
|
||||
if user.email == email:
|
||||
return user
|
||||
return None
|
||||
|
||||
async def find_user_by_email_local_part(self, local_part: str) -> User | None:
|
||||
prefix = f"{(local_part or '').strip().lower()}@"
|
||||
candidates = [u for u in self.rows.values() if u.email.lower().startswith(prefix)]
|
||||
candidates.sort(key=lambda u: u.created_at)
|
||||
return candidates[0] if candidates else None
|
||||
|
||||
async def update_user(self, user: User) -> User:
|
||||
self.rows[str(user.id)] = user
|
||||
return user
|
||||
|
||||
async def count_users(self) -> int:
|
||||
return len(self.rows)
|
||||
|
||||
async def count_admin_users(self) -> int:
|
||||
return sum(1 for u in self.rows.values() if u.system_role == "admin")
|
||||
|
||||
async def get_user_by_oauth(self, provider: str, oauth_id: str) -> User | None:
|
||||
return None
|
||||
|
||||
|
||||
def _seed(repo: FakeUserRepository, email: str, *, role: str = "user") -> User:
|
||||
user = User(
|
||||
id=uuid4(),
|
||||
email=email,
|
||||
password_hash=None,
|
||||
system_role=role, # type: ignore[arg-type]
|
||||
created_at=datetime.now(UTC),
|
||||
needs_setup=False,
|
||||
)
|
||||
repo.rows[str(user.id)] = user
|
||||
return user
|
||||
|
||||
|
||||
def test_normalize_returns_none_for_explicit_email() -> None:
|
||||
assert _normalize_username_local_part("alice@old.com") is None
|
||||
|
||||
|
||||
def test_normalize_lowercases_and_keeps_safe_chars() -> None:
|
||||
assert _normalize_username_local_part("Alice_Smith.1") == "alice_smith.1"
|
||||
|
||||
|
||||
def test_normalize_strips_unsafe_chars() -> None:
|
||||
assert _normalize_username_local_part("hi there!") == "hi-there"
|
||||
|
||||
|
||||
def test_username_to_email_uses_cm_com_by_default(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("DEER_FLOW_USERNAME_LOGIN_DOMAIN", raising=False)
|
||||
assert _username_to_email("alice") == "alice@cm.com"
|
||||
|
||||
|
||||
def test_username_to_email_honours_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("DEER_FLOW_USERNAME_LOGIN_DOMAIN", "example.org")
|
||||
assert _username_to_email("alice") == "alice@example.org"
|
||||
|
||||
|
||||
def test_username_to_email_passes_explicit_email_through() -> None:
|
||||
assert _username_to_email("bob@legacy.io") == "bob@legacy.io"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_part_lookup_finds_legacy_domain_user() -> None:
|
||||
"""Same username under legacy domain should match the new lookup."""
|
||||
repo = FakeUserRepository()
|
||||
_seed(repo, "alice@local.deerflow")
|
||||
legacy = await repo.find_user_by_email_local_part("alice")
|
||||
assert legacy is not None
|
||||
assert legacy.email == "alice@local.deerflow"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_part_lookup_returns_oldest_when_collision() -> None:
|
||||
"""When the same local part exists under multiple domains, prefer the
|
||||
oldest row so re-logins land on the original account."""
|
||||
repo = FakeUserRepository()
|
||||
older = _seed(repo, "alice@local.deerflow")
|
||||
older.created_at = datetime(2024, 1, 1, tzinfo=UTC)
|
||||
newer = _seed(repo, "alice@cm.com")
|
||||
newer.created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
match = await repo.find_user_by_email_local_part("alice")
|
||||
assert match is not None
|
||||
assert match.email == "alice@local.deerflow"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_part_lookup_misses_unknown_username() -> None:
|
||||
repo = FakeUserRepository()
|
||||
_seed(repo, "alice@cm.com")
|
||||
assert await repo.find_user_by_email_local_part("bob") is None
|
||||
Reference in new issue
Block a user