ryan
bc3e3d2797
- Add transitional @jwt_access_required (accepts access tokens, and for now also refresh tokens) and use it on every protected route and inside admin_required. /api/user/refresh stays refresh-only. - Move pending OIDC state/PKCE verifiers from an in-process dict to an oidc_states table (10 min TTL, single-use) so logins survive a blue/green flip. Drop the unused redirect_after_login. - Unit tests for the decorator and the state store. Deploy this commit on its own, before the access-token frontend: the old backend answers access tokens with 422, so the new frontend only works once this backend is live on both colors. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
553 lines
20 KiB
Python
553 lines
20 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import time
|
|
import uuid
|
|
|
|
from quart import Blueprint, jsonify, make_response, request
|
|
from quart_jwt_extended import get_jwt_identity
|
|
|
|
import blueprints.users.models
|
|
from blueprints.users.decorators import jwt_access_required
|
|
from metrics import count_result, observe_response, time_response
|
|
from utils.image_process import image_content_block
|
|
from utils.strip_markdown import extract_agent_text
|
|
from utils.image_upload import ImageValidationError, process_image
|
|
from utils.s3_client import generate_presigned_url as s3_presigned_url
|
|
from utils.s3_client import get_image as s3_get_image
|
|
from utils.s3_client import upload_image as s3_upload_image
|
|
|
|
from .logic import (
|
|
add_message_to_conversation,
|
|
get_agent_for,
|
|
persist_agent_tool_messages,
|
|
tool_result_text,
|
|
)
|
|
from .memory import build_system_prompt_with_memories
|
|
from .models import (
|
|
Conversation,
|
|
PydConversation,
|
|
PydListConversation,
|
|
Speaker,
|
|
)
|
|
|
|
TTFT_TIMEOUT = float(os.getenv("TTFT_TIMEOUT_SECONDS", "30"))
|
|
|
|
_THINK_OPEN = "<think>"
|
|
_THINK_CLOSE = "</think>"
|
|
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
|
|
|
|
|
|
def _route_content(
|
|
buf: str, text: str, in_thinking: bool
|
|
) -> tuple[str, str, bool, str]:
|
|
"""Split a streaming content chunk into (thinking_text, content_text).
|
|
|
|
Handles tags split across chunk boundaries by keeping a small trailing
|
|
buffer. Returns the updated in_thinking state and any leftover buffer.
|
|
"""
|
|
combined = buf + text
|
|
thinking_parts: list[str] = []
|
|
content_parts: list[str] = []
|
|
|
|
while combined:
|
|
tag, hold = (
|
|
(_THINK_CLOSE, len(_THINK_CLOSE) - 1)
|
|
if in_thinking
|
|
else (_THINK_OPEN, len(_THINK_OPEN) - 1)
|
|
)
|
|
idx = combined.find(tag)
|
|
if idx == -1:
|
|
# Tag not yet complete — hold back enough bytes to catch a split.
|
|
if len(combined) <= hold:
|
|
return (
|
|
"".join(thinking_parts),
|
|
"".join(content_parts),
|
|
in_thinking,
|
|
combined,
|
|
)
|
|
safe, leftover = combined[:-hold], combined[-hold:]
|
|
(thinking_parts if in_thinking else content_parts).append(safe)
|
|
return (
|
|
"".join(thinking_parts),
|
|
"".join(content_parts),
|
|
in_thinking,
|
|
leftover,
|
|
)
|
|
(thinking_parts if in_thinking else content_parts).append(combined[:idx])
|
|
combined = combined[idx + len(tag) :]
|
|
in_thinking = not in_thinking
|
|
|
|
return "".join(thinking_parts), "".join(content_parts), in_thinking, ""
|
|
|
|
|
|
conversation_blueprint = Blueprint(
|
|
"conversation_api", __name__, url_prefix="/api/conversation"
|
|
)
|
|
|
|
# Detached persistence tasks. A strong reference keeps them from being garbage
|
|
# collected mid-flight when the SSE stream is cancelled by a client disconnect.
|
|
_persist_tasks: set[asyncio.Task] = set()
|
|
|
|
|
|
def _persist_in_background(coro) -> asyncio.Task:
|
|
task = asyncio.create_task(coro)
|
|
_persist_tasks.add(task)
|
|
task.add_done_callback(_persist_tasks.discard)
|
|
return task
|
|
|
|
|
|
async def _get_owned_conversation(conversation_id, user_id) -> Conversation | None:
|
|
"""Fetch a conversation only if it belongs to the given user."""
|
|
return await Conversation.get_or_none(id=conversation_id, user_id=user_id)
|
|
|
|
|
|
def _build_user_message(
|
|
query_text: str,
|
|
image_block: dict | None = None,
|
|
) -> dict:
|
|
"""Build the single new user message for this turn.
|
|
|
|
History is restored by the LangGraph checkpointer; only the current
|
|
message is passed to avoid duplicating prior turns in checkpoint state.
|
|
"""
|
|
if image_block:
|
|
content = [
|
|
{"type": "text", "text": query_text or "The user sent this image."},
|
|
image_block,
|
|
]
|
|
else:
|
|
content = query_text
|
|
return {"role": "user", "content": content}
|
|
|
|
|
|
@conversation_blueprint.post("/query")
|
|
@jwt_access_required
|
|
async def query():
|
|
current_user_uuid = get_jwt_identity()
|
|
user = await blueprints.users.models.User.get(id=current_user_uuid)
|
|
data = await request.get_json()
|
|
query = data.get("query")
|
|
conversation_id = data.get("conversation_id")
|
|
conversation = await _get_owned_conversation(conversation_id, current_user_uuid)
|
|
if conversation is None:
|
|
return jsonify({"error": "conversation not found"}), 404
|
|
await add_message_to_conversation(
|
|
conversation=conversation,
|
|
message=query,
|
|
speaker="user",
|
|
user=user,
|
|
)
|
|
|
|
with time_response(stream=False):
|
|
system_prompt = await build_system_prompt_with_memories(str(user.id))
|
|
payload = {"messages": [_build_user_message(query)]}
|
|
agent_config = {
|
|
"configurable": {
|
|
"user_id": str(user.id),
|
|
"is_admin": user.is_admin(),
|
|
"thread_id": str(conversation.id),
|
|
"system_prompt": system_prompt,
|
|
}
|
|
}
|
|
|
|
response = await get_agent_for(user).ainvoke(payload, config=agent_config)
|
|
message = extract_agent_text(response.get("messages", [])[-1].content)
|
|
await persist_agent_tool_messages(
|
|
conversation, user, response.get("messages", [])
|
|
)
|
|
await add_message_to_conversation(
|
|
conversation=conversation,
|
|
message=message,
|
|
speaker="simba",
|
|
user=user,
|
|
)
|
|
return jsonify({"response": message})
|
|
|
|
|
|
@conversation_blueprint.post("/upload-image")
|
|
@jwt_access_required
|
|
async def upload_image():
|
|
current_user_uuid = get_jwt_identity()
|
|
await blueprints.users.models.User.get(id=current_user_uuid)
|
|
|
|
files = await request.files
|
|
form = await request.form
|
|
file = files.get("file")
|
|
conversation_id = form.get("conversation_id")
|
|
|
|
if not file or not conversation_id:
|
|
return jsonify({"error": "file and conversation_id are required"}), 400
|
|
|
|
file_bytes = file.read()
|
|
content_type = file.content_type or "image/jpeg"
|
|
|
|
try:
|
|
processed_bytes, output_content_type = process_image(file_bytes, content_type)
|
|
except ImageValidationError as e:
|
|
return jsonify({"error": str(e)}), 400
|
|
|
|
ext = output_content_type.split("/")[-1]
|
|
if ext == "jpeg":
|
|
ext = "jpg"
|
|
key = f"conversations/{conversation_id}/{uuid.uuid4()}.{ext}"
|
|
|
|
await s3_upload_image(processed_bytes, key, output_content_type)
|
|
|
|
return jsonify({"image_key": key})
|
|
|
|
|
|
@conversation_blueprint.get("/image/<path:image_key>")
|
|
@jwt_access_required
|
|
async def serve_image(image_key: str):
|
|
url = await s3_presigned_url(image_key)
|
|
return jsonify({"url": url})
|
|
|
|
|
|
@conversation_blueprint.post("/stream-query")
|
|
@jwt_access_required
|
|
async def stream_query():
|
|
current_user_uuid = get_jwt_identity()
|
|
user = await blueprints.users.models.User.get(id=current_user_uuid)
|
|
data = await request.get_json()
|
|
query_text = data.get("query")
|
|
conversation_id = data.get("conversation_id")
|
|
image_key = data.get("image_key")
|
|
conversation = await _get_owned_conversation(conversation_id, current_user_uuid)
|
|
if conversation is None:
|
|
return jsonify({"error": "conversation not found"}), 404
|
|
await add_message_to_conversation(
|
|
conversation=conversation,
|
|
message=query_text or "",
|
|
speaker="user",
|
|
user=user,
|
|
image_key=image_key,
|
|
)
|
|
|
|
# If an image was uploaded, attach it to the agent request directly
|
|
image_block = None
|
|
if image_key:
|
|
try:
|
|
image_bytes, _ = await s3_get_image(image_key)
|
|
image_block = image_content_block(image_bytes)
|
|
except Exception as e:
|
|
logging.error(f"Failed to load image {image_key}: {e}")
|
|
query_text = f"{query_text or ''}\n\n[An image was sent but could not be loaded.]".strip()
|
|
|
|
system_prompt = await build_system_prompt_with_memories(str(user.id))
|
|
payload = {"messages": [_build_user_message(query_text or "", image_block)]}
|
|
agent_config = {
|
|
"configurable": {
|
|
"user_id": str(user.id),
|
|
"is_admin": user.is_admin(),
|
|
"thread_id": str(conversation.id),
|
|
"system_prompt": system_prompt,
|
|
}
|
|
}
|
|
# Capture the agent reference while the request/app context is still active.
|
|
# Quart exits the app context before iterating the SSE response body, so
|
|
# current_app cannot be accessed inside event_generator().
|
|
main_agent = get_agent_for(user)
|
|
|
|
async def event_generator():
|
|
final_message = None
|
|
streamed = ""
|
|
# Tag-routing state for <think>…</think> blocks.
|
|
in_thinking = False
|
|
tag_buf = ""
|
|
# Fallback heartbeat for models that emit empty chunks before thinking text.
|
|
last_thinking = 0.0
|
|
turn_has_content = False
|
|
|
|
# Timing / observability
|
|
t0 = time.monotonic()
|
|
llm_start_at: float | None = None
|
|
first_content_at: float | None = None
|
|
tool_timers: dict[str, float] = {} # run_id -> start time
|
|
tool_spans: list[dict] = []
|
|
n_tokens = 0
|
|
model_name: str | None = None
|
|
cancelled = False
|
|
|
|
async def persist(text: str) -> None:
|
|
await add_message_to_conversation(
|
|
conversation=conversation,
|
|
message=text,
|
|
speaker="simba",
|
|
user=user,
|
|
)
|
|
|
|
def _log_summary(*, error: str | None = None) -> None:
|
|
now = time.monotonic()
|
|
thinking_ms = (
|
|
round((first_content_at - llm_start_at) * 1000)
|
|
if llm_start_at and first_content_at
|
|
else None
|
|
)
|
|
ttft_ms = (
|
|
round((first_content_at - t0) * 1000) if first_content_at else None
|
|
)
|
|
logging.info(
|
|
json.dumps(
|
|
{
|
|
"event": "stream_request",
|
|
"conv": str(conversation.id),
|
|
"user": user.username,
|
|
"total_ms": round((now - t0) * 1000),
|
|
"thinking_ms": thinking_ms,
|
|
"ttft_ms": ttft_ms,
|
|
"tokens": n_tokens,
|
|
"tools": tool_spans,
|
|
"model": model_name,
|
|
"cancelled": cancelled,
|
|
**({"error": error} if error else {}),
|
|
}
|
|
)
|
|
)
|
|
|
|
try:
|
|
aiter = main_agent.astream_events(
|
|
payload, version="v2", config=agent_config
|
|
).__aiter__()
|
|
while True:
|
|
# Apply a short TTFT timeout only during the window between
|
|
# on_chat_model_start and the first content token. Before the
|
|
# LLM call starts (tool phase) and after the first token, let
|
|
# the httpx read timeout handle stalls instead.
|
|
if llm_start_at is not None and first_content_at is None:
|
|
ttft_wait = TTFT_TIMEOUT - (time.monotonic() - llm_start_at)
|
|
timeout = max(ttft_wait, 1.0)
|
|
else:
|
|
timeout = None
|
|
|
|
try:
|
|
if timeout is not None:
|
|
event = await asyncio.wait_for(
|
|
aiter.__anext__(), timeout=timeout
|
|
)
|
|
else:
|
|
event = await aiter.__anext__()
|
|
except StopAsyncIteration:
|
|
break
|
|
except asyncio.TimeoutError:
|
|
_log_summary(error=f"TTFT timeout after {TTFT_TIMEOUT:.0f}s")
|
|
yield f"data: {json.dumps({'type': 'error', 'message': 'Simba took too long to respond. Please try again.', 'stalled': True})}\n\n"
|
|
return
|
|
|
|
event_type = event.get("event")
|
|
|
|
if event_type == "on_tool_start":
|
|
tool_timers[event["run_id"]] = time.monotonic()
|
|
yield f"data: {json.dumps({'type': 'tool_start', 'tool': event['name']})}\n\n"
|
|
|
|
elif event_type == "on_tool_end":
|
|
start = tool_timers.pop(event["run_id"], None)
|
|
tool_spans.append(
|
|
{
|
|
"tool": event["name"],
|
|
"ms": round((time.monotonic() - start) * 1000)
|
|
if start
|
|
else None,
|
|
}
|
|
)
|
|
yield f"data: {json.dumps({'type': 'tool_end', 'tool': event['name']})}\n\n"
|
|
# Persist tool output so later turns keep the fetched context
|
|
output = event.get("data", {}).get("output")
|
|
if output is not None:
|
|
await add_message_to_conversation(
|
|
conversation=conversation,
|
|
message=tool_result_text(event["name"], output),
|
|
speaker="tool",
|
|
user=user,
|
|
)
|
|
|
|
elif event_type == "on_chat_model_start":
|
|
turn_has_content = False
|
|
if llm_start_at is None:
|
|
llm_start_at = time.monotonic()
|
|
if model_name is None:
|
|
model_name = event.get("metadata", {}).get(
|
|
"ls_model_name"
|
|
) or event.get("name")
|
|
|
|
elif event_type == "on_chat_model_stream":
|
|
chunk = event.get("data", {}).get("chunk")
|
|
raw = getattr(chunk, "content", None)
|
|
if isinstance(raw, str) and raw:
|
|
think_text, content_text, in_thinking, tag_buf = _route_content(
|
|
tag_buf, raw, in_thinking
|
|
)
|
|
if think_text:
|
|
yield f"data: {json.dumps({'type': 'thinking_delta', 'delta': think_text})}\n\n"
|
|
if content_text:
|
|
if first_content_at is None:
|
|
first_content_at = time.monotonic()
|
|
turn_has_content = True
|
|
n_tokens += 1
|
|
streamed += content_text
|
|
yield f"data: {json.dumps({'type': 'content', 'delta': content_text})}\n\n"
|
|
elif not turn_has_content:
|
|
# Empty chunk before any content — heartbeat for models
|
|
# that don't emit <think> tags at all.
|
|
now = time.monotonic()
|
|
if now - last_thinking >= 1.0:
|
|
last_thinking = now
|
|
yield f"data: {json.dumps({'type': 'thinking'})}\n\n"
|
|
|
|
elif event_type == "on_chain_end":
|
|
output = event.get("data", {}).get("output")
|
|
if isinstance(output, dict):
|
|
msgs = output.get("messages", [])
|
|
if msgs:
|
|
last_msg = msgs[-1]
|
|
content = getattr(last_msg, "content", None)
|
|
if isinstance(content, str) and content:
|
|
final_message = _THINK_RE.sub("", content).strip()
|
|
|
|
except asyncio.CancelledError:
|
|
cancelled = True
|
|
_log_summary()
|
|
text = final_message or streamed
|
|
if text:
|
|
_persist_in_background(persist(text))
|
|
raise
|
|
except Exception as e:
|
|
count_result(stream=True, result="error")
|
|
_log_summary(error=str(e))
|
|
yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
|
|
return
|
|
|
|
final_message = final_message or streamed
|
|
if final_message:
|
|
# Persist via a detached, shielded task so a disconnect at this
|
|
# point (after generation completes) still commits the answer.
|
|
task = _persist_in_background(persist(final_message))
|
|
await asyncio.shield(task)
|
|
yield f"data: {json.dumps({'type': 'response', 'message': final_message})}\n\n"
|
|
else:
|
|
yield f"data: {json.dumps({'type': 'error', 'message': 'No response generated'})}\n\n"
|
|
|
|
observe_response(stream=True, seconds=time.monotonic() - t0)
|
|
count_result(stream=True, result="success")
|
|
_log_summary()
|
|
yield "data: [DONE]\n\n"
|
|
|
|
return await make_response(
|
|
event_generator(),
|
|
200,
|
|
{
|
|
"Content-Type": "text/event-stream",
|
|
"Cache-Control": "no-cache",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
|
|
@conversation_blueprint.route("/<conversation_id>")
|
|
@jwt_access_required
|
|
async def get_conversation(conversation_id: str):
|
|
user_uuid = get_jwt_identity()
|
|
conversation = await _get_owned_conversation(conversation_id, user_uuid)
|
|
if conversation is None:
|
|
return jsonify({"error": "conversation not found"}), 404
|
|
await conversation.fetch_related("messages")
|
|
|
|
# Manually serialize the conversation with messages
|
|
messages = []
|
|
for msg in conversation.messages:
|
|
# Tool messages hold raw fetched context for the LLM; only surface
|
|
# their header line (e.g. "[web_search result]") in the UI.
|
|
text = msg.text.split("\n", 1)[0] if msg.speaker == Speaker.TOOL else msg.text
|
|
messages.append(
|
|
{
|
|
"id": str(msg.id),
|
|
"text": text,
|
|
"speaker": msg.speaker.value,
|
|
"created_at": msg.created_at.isoformat(),
|
|
"image_key": msg.image_key,
|
|
}
|
|
)
|
|
return jsonify(
|
|
{
|
|
"id": str(conversation.id),
|
|
"name": conversation.name,
|
|
"messages": messages,
|
|
"created_at": conversation.created_at.isoformat(),
|
|
"updated_at": conversation.updated_at.isoformat(),
|
|
}
|
|
)
|
|
|
|
|
|
@conversation_blueprint.post("/")
|
|
@jwt_access_required
|
|
async def create_conversation():
|
|
user_uuid = get_jwt_identity()
|
|
user = await blueprints.users.models.User.get(id=user_uuid)
|
|
conversation = await Conversation.create(
|
|
name="New Conversation",
|
|
user=user,
|
|
)
|
|
|
|
serialized_conversation = await PydConversation.from_tortoise_orm(conversation)
|
|
return jsonify(serialized_conversation.model_dump())
|
|
|
|
|
|
@conversation_blueprint.get("/")
|
|
@jwt_access_required
|
|
async def get_all_conversations():
|
|
user_uuid = get_jwt_identity()
|
|
user = await blueprints.users.models.User.get(id=user_uuid)
|
|
|
|
query = Conversation.filter(user=user)
|
|
|
|
search = request.args.get("search", "").strip()
|
|
if search:
|
|
query = query.filter(name__icontains=search)
|
|
|
|
query = query.order_by("-updated_at")
|
|
|
|
# `limit` is optional: when omitted we return the full list (backward
|
|
# compatible with any non-frontend callers). `offset` only applies with it.
|
|
limit = request.args.get("limit", type=int)
|
|
if limit is not None:
|
|
offset = request.args.get("offset", default=0, type=int)
|
|
query = query.offset(offset).limit(limit)
|
|
|
|
serialized_conversations = await PydListConversation.from_queryset(query)
|
|
|
|
return jsonify(serialized_conversations.model_dump())
|
|
|
|
|
|
@conversation_blueprint.patch("/<conversation_id>")
|
|
@jwt_access_required
|
|
async def rename_conversation(conversation_id: str):
|
|
user_uuid = get_jwt_identity()
|
|
body = await request.get_json()
|
|
name = (body or {}).get("name", "").strip()
|
|
if not name:
|
|
return jsonify({"error": "name is required"}), 400
|
|
|
|
conversation = await _get_owned_conversation(conversation_id, user_uuid)
|
|
if conversation is None:
|
|
return jsonify({"error": "conversation not found"}), 404
|
|
|
|
conversation.name = name[:255]
|
|
await conversation.save()
|
|
|
|
serialized_conversation = await PydConversation.from_tortoise_orm(conversation)
|
|
return jsonify(serialized_conversation.model_dump())
|
|
|
|
|
|
@conversation_blueprint.delete("/<conversation_id>")
|
|
@jwt_access_required
|
|
async def delete_conversation(conversation_id: str):
|
|
user_uuid = get_jwt_identity()
|
|
conversation = await _get_owned_conversation(conversation_id, user_uuid)
|
|
if conversation is None:
|
|
return jsonify({"error": "conversation not found"}), 404
|
|
|
|
await conversation.delete()
|
|
return jsonify({"success": True})
|