Refactor settings store to rename proModel to proThinking and update related logic; enhance CSS for energy efficiency and reduced motion preferences; improve i18n translations for better clarity and consistency; modify proBlock utility functions for clearer instruction handling; streamline Vite configuration by removing unnecessary Univer.js dependencies.
This commit is contained in:
“ydy0615”
@@ -0,0 +1,389 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Request, Security
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
from geoip import get_ip_location_text
|
||||
from llm import stream_ollama_events
|
||||
from models import UserPreferences
|
||||
|
||||
logger = logging.getLogger("api.pro")
|
||||
|
||||
PRO_COMPLETION_TIMEOUT = float(os.getenv("PRO_COMPLETION_TIMEOUT", "3600"))
|
||||
PRO_QUEUE_TIMEOUT = float(os.getenv("PRO_QUEUE_TIMEOUT", "600"))
|
||||
PRO_MAX_CONCURRENCY = max(1, int(os.getenv("PRO_MAX_CONCURRENCY", "1")))
|
||||
PRO_QUEUE_MAX_SIZE = max(0, int(os.getenv("PRO_QUEUE_MAX_SIZE", "5")))
|
||||
PRO_STATUS_RETENTION_SECONDS = float(os.getenv("PRO_STATUS_RETENTION_SECONDS", "600"))
|
||||
PRO_CANCEL_ACK_TIMEOUT = 5.0
|
||||
PUBLIC_PRO_ERROR = "PRO generation failed. Please retry or adjust the instruction."
|
||||
|
||||
|
||||
class ProCompletionRequest(BaseModel):
|
||||
prefix: str
|
||||
suffix: str
|
||||
languageId: str = "markdown"
|
||||
instruction: str = ""
|
||||
pro_thinking: str = "medium"
|
||||
privacy_mode: bool = False
|
||||
user_preferences: Optional[UserPreferences] = None
|
||||
|
||||
|
||||
class ProCancelRequest(BaseModel):
|
||||
request_id: str
|
||||
reason: str = "abort"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProRequestState:
|
||||
request_id: str
|
||||
status: str = "queued"
|
||||
created_at: float = field(default_factory=time.time)
|
||||
updated_at: float = field(default_factory=time.time)
|
||||
error: str = ""
|
||||
task: asyncio.Task | None = None
|
||||
cancel_requested: bool = False
|
||||
done_event: asyncio.Event = field(default_factory=asyncio.Event)
|
||||
|
||||
def touch(self, status: str | None = None, error: str = "") -> None:
|
||||
if status:
|
||||
self.status = status
|
||||
if error:
|
||||
self.error = error
|
||||
self.updated_at = time.time()
|
||||
|
||||
def request_cancel(self) -> None:
|
||||
self.cancel_requested = True
|
||||
self.touch("cancelled")
|
||||
|
||||
|
||||
PRO_STATES: dict[str, ProRequestState] = {}
|
||||
PRO_STATES_LOCK = asyncio.Lock()
|
||||
PRO_SEMAPHORE = asyncio.Semaphore(PRO_MAX_CONCURRENCY)
|
||||
|
||||
|
||||
def _iso_timestamp(value: float) -> str:
|
||||
return datetime.fromtimestamp(value, tz=timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _clamp_thinking(value: str | None) -> str | None:
|
||||
normalized = (value or "medium").strip().lower()
|
||||
if normalized in {"none", "off", "false"}:
|
||||
return None
|
||||
if normalized in {"low", "medium", "high"}:
|
||||
return normalized
|
||||
return "medium"
|
||||
|
||||
|
||||
def _queued_states() -> list[ProRequestState]:
|
||||
return [state for state in PRO_STATES.values() if state.status == "queued"]
|
||||
|
||||
|
||||
def _queue_position(request_id: str) -> int | None:
|
||||
queued = sorted(_queued_states(), key=lambda item: item.created_at)
|
||||
for index, state in enumerate(queued, start=1):
|
||||
if state.request_id == request_id:
|
||||
return index
|
||||
return None
|
||||
|
||||
|
||||
async def _cleanup_states() -> None:
|
||||
now = time.time()
|
||||
expired = [
|
||||
request_id
|
||||
for request_id, state in PRO_STATES.items()
|
||||
if state.status in {"done", "error", "cancelled"}
|
||||
and now - state.updated_at > PRO_STATUS_RETENTION_SECONDS
|
||||
]
|
||||
for request_id in expired:
|
||||
PRO_STATES.pop(request_id, None)
|
||||
|
||||
|
||||
def _state_payload(state: ProRequestState) -> dict:
|
||||
return {
|
||||
"request_id": state.request_id,
|
||||
"status": state.status,
|
||||
"queue_position": _queue_position(state.request_id),
|
||||
"created_at": _iso_timestamp(state.created_at),
|
||||
"updated_at": _iso_timestamp(state.updated_at),
|
||||
"error": state.error,
|
||||
}
|
||||
|
||||
|
||||
def _build_pro_prompts(
|
||||
*,
|
||||
prefix: str,
|
||||
suffix: str,
|
||||
language_id: str,
|
||||
instruction: str,
|
||||
location: str = "",
|
||||
preferences: UserPreferences | None = None,
|
||||
) -> tuple[str, str]:
|
||||
safe_language = (language_id or "markdown").strip() or "markdown"
|
||||
safe_instruction = (instruction or "").strip()
|
||||
preference_lines: list[str] = []
|
||||
if preferences:
|
||||
if preferences.language and preferences.language != "auto":
|
||||
preference_lines.append(f"- Preferred language: {preferences.language}")
|
||||
if preferences.currency and preferences.currency != "auto":
|
||||
preference_lines.append(f"- Preferred currency: {preferences.currency}")
|
||||
if preferences.timezone and preferences.timezone != "auto":
|
||||
preference_lines.append(f"- Timezone: {preferences.timezone}")
|
||||
if location:
|
||||
preference_lines.append(f"- Location hint: {location}")
|
||||
|
||||
system_prompt = f"""You edit Markdown documents.
|
||||
Return only the Markdown text to insert at the cursor.
|
||||
Do not explain, analyze, label the answer, or wrap the whole answer in a code fence.
|
||||
Match the document language, style, and Markdown structure.
|
||||
Language: {safe_language}."""
|
||||
|
||||
preferences_text = "\n".join(preference_lines) if preference_lines else "- none"
|
||||
instruction_text = safe_instruction or "Continue the Markdown naturally."
|
||||
user_prompt = f"""Instruction:
|
||||
{instruction_text}
|
||||
|
||||
User preferences:
|
||||
{preferences_text}
|
||||
|
||||
Markdown before cursor:
|
||||
{prefix}
|
||||
|
||||
Markdown after cursor:
|
||||
{suffix}
|
||||
|
||||
Write only the Markdown that belongs at the cursor."""
|
||||
return system_prompt.strip(), user_prompt.strip()
|
||||
|
||||
|
||||
def _get_client_ip(request: Request) -> str:
|
||||
if request.client:
|
||||
return request.headers.get("X-Client-IP") or request.client.host
|
||||
return request.headers.get("X-Client-IP") or "unknown"
|
||||
|
||||
|
||||
async def _send_sse_event(queue: asyncio.Queue, event_name: str, data: dict) -> None:
|
||||
await queue.put((event_name, json.dumps(data, ensure_ascii=False)))
|
||||
|
||||
|
||||
async def _wait_for_cancel_cleanup(state: ProRequestState, request_tag: str, reason: str) -> None:
|
||||
if state.done_event.is_set():
|
||||
return
|
||||
try:
|
||||
await asyncio.wait_for(state.done_event.wait(), timeout=PRO_CANCEL_ACK_TIMEOUT)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"[%s] /v1/pro/completions cancel cleanup not confirmed request_id=%s reason=%s",
|
||||
request_tag,
|
||||
state.request_id,
|
||||
reason,
|
||||
)
|
||||
|
||||
|
||||
def register_pro_completion_routes(app: FastAPI, get_api_key):
|
||||
@app.post("/v1/pro/completions")
|
||||
async def create_pro_completion(
|
||||
request: Request,
|
||||
req: ProCompletionRequest,
|
||||
api_key: str = Security(get_api_key),
|
||||
):
|
||||
request_id = request.headers.get("X-Request-Id") or str(uuid.uuid4())
|
||||
request_tag = request_id[:8]
|
||||
event_queue: asyncio.Queue[tuple[str, str] | None] = asyncio.Queue()
|
||||
previous_state: ProRequestState | None = None
|
||||
|
||||
async with PRO_STATES_LOCK:
|
||||
await _cleanup_states()
|
||||
queued_count = len(_queued_states())
|
||||
if queued_count >= PRO_QUEUE_MAX_SIZE:
|
||||
logger.info("[%s] /v1/pro/completions rejected queue_full request_id=%s", request_tag, request_id)
|
||||
return JSONResponse(
|
||||
content={"error": "PRO queue is full", "request_id": request_id},
|
||||
status_code=429,
|
||||
)
|
||||
|
||||
existing = PRO_STATES.get(request_id)
|
||||
if existing and existing.task and not existing.task.done():
|
||||
existing.request_cancel()
|
||||
existing.task.cancel()
|
||||
previous_state = existing
|
||||
|
||||
state = ProRequestState(request_id=request_id)
|
||||
PRO_STATES[request_id] = state
|
||||
|
||||
if previous_state:
|
||||
await _wait_for_cancel_cleanup(previous_state, request_tag, "replace")
|
||||
|
||||
client_ip = "hidden"
|
||||
location = ""
|
||||
if not req.privacy_mode: # pragma: no cover
|
||||
client_ip = _get_client_ip(request)
|
||||
location = get_ip_location_text(client_ip)
|
||||
|
||||
prefix = req.prefix or ""
|
||||
suffix = req.suffix or ""
|
||||
|
||||
system_prompt, user_prompt = _build_pro_prompts(
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
language_id=req.languageId,
|
||||
instruction=req.instruction,
|
||||
location=location,
|
||||
preferences=req.user_preferences,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"[%s] /v1/pro/completions request_id=%s client_ip=%s prefix_chars=%d suffix_chars=%d instruction_chars=%d lang=%s thinking=%s",
|
||||
request_tag,
|
||||
request_id,
|
||||
client_ip,
|
||||
len(prefix),
|
||||
len(suffix),
|
||||
len(req.instruction or ""),
|
||||
req.languageId,
|
||||
req.pro_thinking,
|
||||
)
|
||||
|
||||
async def producer() -> None:
|
||||
acquired = False
|
||||
chunks: list[str] = []
|
||||
try:
|
||||
async with PRO_STATES_LOCK:
|
||||
if state.cancel_requested:
|
||||
raise asyncio.CancelledError()
|
||||
state.touch("queued")
|
||||
queue_position = _queue_position(request_id)
|
||||
await _send_sse_event(event_queue, "queued", {"request_id": request_id, "queue_position": queue_position})
|
||||
|
||||
await asyncio.wait_for(PRO_SEMAPHORE.acquire(), timeout=PRO_QUEUE_TIMEOUT)
|
||||
acquired = True
|
||||
|
||||
async with PRO_STATES_LOCK:
|
||||
if state.cancel_requested:
|
||||
raise asyncio.CancelledError()
|
||||
state.touch("started")
|
||||
await _send_sse_event(event_queue, "started", {"request_id": request_id})
|
||||
|
||||
async for event_type, payload in stream_ollama_events(
|
||||
user_prompt,
|
||||
system_prompt=system_prompt,
|
||||
tag=f"{request_tag}-pro",
|
||||
temperature=0.7,
|
||||
thinking=_clamp_thinking(req.pro_thinking),
|
||||
use_pro_model=True,
|
||||
enable_thinking=True,
|
||||
timeout=PRO_COMPLETION_TIMEOUT,
|
||||
):
|
||||
if event_type == "thinking":
|
||||
await _send_sse_event(event_queue, "thinking", {"request_id": request_id})
|
||||
continue
|
||||
|
||||
if not payload:
|
||||
continue
|
||||
chunks.append(payload)
|
||||
await _send_sse_event(event_queue, "chunk", {"delta": payload, "request_id": request_id})
|
||||
|
||||
content = "".join(chunks)
|
||||
async with PRO_STATES_LOCK:
|
||||
if state.cancel_requested:
|
||||
raise asyncio.CancelledError()
|
||||
if not content:
|
||||
raise ValueError("PRO returned empty content")
|
||||
|
||||
async with PRO_STATES_LOCK:
|
||||
state.touch("done")
|
||||
logger.info("[%s] /v1/pro/completions done request_id=%s content_chars=%d", request_tag, request_id, len(content))
|
||||
await _send_sse_event(event_queue, "done", {"content": content, "request_id": request_id})
|
||||
except asyncio.CancelledError:
|
||||
async with PRO_STATES_LOCK:
|
||||
state.request_cancel()
|
||||
logger.info("[%s] /v1/pro/completions cancelled request_id=%s", request_tag, request_id)
|
||||
await _send_sse_event(event_queue, "cancelled", {"cancelled": True, "request_id": request_id})
|
||||
raise
|
||||
except Exception as exc:
|
||||
async with PRO_STATES_LOCK:
|
||||
state.touch("error", PUBLIC_PRO_ERROR)
|
||||
logger.exception("[%s] /v1/pro/completions failed request_id=%s", request_tag, request_id)
|
||||
await _send_sse_event(event_queue, "error", {"error": PUBLIC_PRO_ERROR, "request_id": request_id})
|
||||
finally:
|
||||
if acquired:
|
||||
PRO_SEMAPHORE.release()
|
||||
state.done_event.set()
|
||||
await event_queue.put(None)
|
||||
|
||||
producer_task = asyncio.create_task(producer())
|
||||
async with PRO_STATES_LOCK:
|
||||
state.task = producer_task
|
||||
|
||||
async def event_stream():
|
||||
try:
|
||||
while True:
|
||||
item = await event_queue.get()
|
||||
if item is None:
|
||||
break
|
||||
event_name, data = item
|
||||
yield f"event: {event_name}\ndata: {data}\n\n"
|
||||
except asyncio.CancelledError:
|
||||
async with PRO_STATES_LOCK:
|
||||
state.request_cancel()
|
||||
producer_task.cancel()
|
||||
raise
|
||||
finally:
|
||||
if not producer_task.done() and not state.done_event.is_set():
|
||||
async with PRO_STATES_LOCK:
|
||||
state.request_cancel()
|
||||
producer_task.cancel()
|
||||
with contextlib.suppress(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(state.done_event.wait(), timeout=PRO_CANCEL_ACK_TIMEOUT)
|
||||
|
||||
return StreamingResponse(
|
||||
event_stream(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
@app.post("/v1/pro/completions/cancel")
|
||||
async def cancel_pro_completion(req: ProCancelRequest, api_key: str = Security(get_api_key)):
|
||||
request_id = req.request_id or ""
|
||||
request_tag = request_id[:8]
|
||||
state_to_wait: ProRequestState | None = None
|
||||
async with PRO_STATES_LOCK:
|
||||
await _cleanup_states()
|
||||
state = PRO_STATES.get(request_id)
|
||||
if not state:
|
||||
return {"cancelled": False, "status": "not_found"}
|
||||
if state.task and not state.task.done():
|
||||
state.request_cancel()
|
||||
state.task.cancel()
|
||||
state_to_wait = state
|
||||
if state.status in {"done", "error", "cancelled"}:
|
||||
if not state_to_wait:
|
||||
return {"cancelled": False, "status": state.status}
|
||||
else:
|
||||
state.request_cancel()
|
||||
|
||||
if state_to_wait:
|
||||
await _wait_for_cancel_cleanup(state_to_wait, request_tag, req.reason)
|
||||
|
||||
return {"cancelled": True, "status": "ok"}
|
||||
|
||||
@app.get("/v1/pro/completions/status/{request_id}")
|
||||
async def get_pro_completion_status(request_id: str, api_key: str = Security(get_api_key)):
|
||||
async with PRO_STATES_LOCK:
|
||||
await _cleanup_states()
|
||||
state = PRO_STATES.get(request_id)
|
||||
if not state:
|
||||
raise HTTPException(status_code=404, detail="PRO request not found")
|
||||
return _state_payload(state)
|
||||
Reference in New Issue
Block a user