“ydy0615”
2c7a02f587
- Added support for PRO mode in LLM with specific instruction handling and context awareness. - Updated prompt building functions to include prefill options for better context management. - Introduced new inline examples for PRO mode in JSON format. - Enhanced system prompts to reflect PRO mode capabilities and rules. - Modified API endpoints to accommodate new parameters and ensure backward compatibility. - Improved test cases to validate new functionality and ensure comprehensive coverage.
367 lines
13 KiB
Python
367 lines
13 KiB
Python
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
|
|
from prompt import build_pro_completion_prompts
|
|
|
|
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,
|
|
pro_thinking: str = "medium",
|
|
location: str = "",
|
|
preferences: UserPreferences | None = None,
|
|
) -> tuple[str, str]:
|
|
return build_pro_completion_prompts(
|
|
prefix=prefix,
|
|
suffix=suffix,
|
|
instruction=instruction,
|
|
language_id=language_id,
|
|
location=location,
|
|
pro_thinking_level=pro_thinking,
|
|
preferences=preferences,
|
|
)
|
|
|
|
|
|
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,
|
|
pro_thinking=req.pro_thinking,
|
|
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)
|