feat: add web search block functionality and integrate with existing plugins
- Introduced a new web search block plugin to handle web search queries and results. - Updated copilot, doc block, and pro block plugins to include web search context in AI completions. - Implemented utility functions for parsing and building web search markdown. - Enhanced API to support web search requests and responses. - Added configuration for web search URL and timeout settings. - Updated size limit checks to account for web search content.
This commit is contained in:
+58
-3
@@ -5,6 +5,7 @@ import logging
|
||||
import os
|
||||
import uuid
|
||||
from contextlib import suppress
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import FastAPI, File, Form, HTTPException, Request, Response, Security, UploadFile
|
||||
@@ -27,6 +28,7 @@ from job_handlers import (
|
||||
ocr_handler,
|
||||
pro_completion_handler,
|
||||
tts_handler,
|
||||
web_search_handler,
|
||||
)
|
||||
from job_system import (
|
||||
InMemoryJobManager,
|
||||
@@ -87,6 +89,14 @@ class ProCompletionRequest(BaseModel):
|
||||
user_preferences: Optional[UserPreferences] = None
|
||||
|
||||
|
||||
class WebSearchRequest(BaseModel):
|
||||
prefix: str
|
||||
suffix: str
|
||||
languageId: str = "markdown"
|
||||
privacy_mode: bool = False
|
||||
user_preferences: Optional[UserPreferences] = None
|
||||
|
||||
|
||||
class CancelCompletionRequest(BaseModel):
|
||||
request_id: str
|
||||
reason: str = "abort"
|
||||
@@ -267,18 +277,24 @@ async def _authorize_request(
|
||||
|
||||
|
||||
def _register_handlers() -> None:
|
||||
"""注册所有任务处理器。每次调用都会重新获取当前 manager 实例并强制注册,
|
||||
确保 Redis 重连或实例重建后处理器不会丢失。"""
|
||||
global _handlers_registered
|
||||
if _handlers_registered:
|
||||
return
|
||||
manager = get_job_manager()
|
||||
# 强制清空旧 handlers,避免重复注册累积
|
||||
manager.handlers.clear()
|
||||
manager.register_handler("completion", completion_handler)
|
||||
manager.register_handler("pro_completion", pro_completion_handler)
|
||||
manager.register_handler("web_search", web_search_handler)
|
||||
manager.register_handler("compress", compress_handler)
|
||||
manager.register_handler("ocr", ocr_handler)
|
||||
manager.register_handler("convert", convert_handler)
|
||||
manager.register_handler("tts", tts_handler)
|
||||
manager.register_handler("asr", asr_handler)
|
||||
_handlers_registered = True
|
||||
|
||||
# 打印注册信息便于调试
|
||||
logger.info("handlers registered: %s", list(manager.handlers.keys()))tered: %s", list(manager.handlers.keys()))
|
||||
|
||||
|
||||
def _sse(event: str, data: dict) -> str:
|
||||
@@ -375,7 +391,7 @@ async def _guard_api_request(request: Request, *, scope: str) -> tuple[RiskIdent
|
||||
return identity, decision
|
||||
|
||||
|
||||
def _estimate_completion_chars(req: CompletionRequest | ProCompletionRequest) -> int:
|
||||
def _estimate_completion_chars(req: CompletionRequest | ProCompletionRequest | WebSearchRequest) -> int:
|
||||
return len(req.prefix or "") + len(req.suffix or "") + len(getattr(req, "instruction", "") or "")
|
||||
|
||||
|
||||
@@ -551,6 +567,45 @@ async def cancel_pro_completion(req: CancelCompletionRequest, auth: dict = Secur
|
||||
return await _cancel_job(req.request_id or "", req.reason)
|
||||
|
||||
|
||||
@app.post("/v1/web-search")
|
||||
async def create_web_search(
|
||||
request: Request,
|
||||
req: WebSearchRequest,
|
||||
auth: dict = Security(_authorize_request),
|
||||
):
|
||||
del auth
|
||||
body = {
|
||||
"prefix": req.prefix,
|
||||
"suffix": req.suffix,
|
||||
"languageId": req.languageId,
|
||||
"privacy_mode": req.privacy_mode,
|
||||
"user_preferences": _serialize_preferences(req.user_preferences),
|
||||
}
|
||||
try:
|
||||
identity, payload = await _prepare_llm_payload(
|
||||
request,
|
||||
job_type="web_search",
|
||||
request_body=body,
|
||||
raw_size=_estimate_completion_chars(req),
|
||||
token_source_text=f"{req.prefix}\n{req.suffix}",
|
||||
)
|
||||
payload["created_at"] = datetime.utcnow().isoformat()
|
||||
job_id = await _queue_job("web_search", payload, identity.request_id)
|
||||
except RiskRejected as exc:
|
||||
return _risk_json_response(_request_identity(request), exc.decision)
|
||||
except QueueFullError as exc:
|
||||
return JSONResponse({"error": str(exc), "request_id": _request_id(request)}, status_code=429)
|
||||
except JobSystemError as exc:
|
||||
return JSONResponse({"error": str(exc), "request_id": _request_id(request)}, status_code=503)
|
||||
return await _stream_job(job_id)
|
||||
|
||||
|
||||
@app.post("/v1/web-search/cancel")
|
||||
async def cancel_web_search(req: CancelCompletionRequest, auth: dict = Security(_authorize_request)):
|
||||
del auth
|
||||
return await _cancel_job(req.request_id or "", req.reason)
|
||||
|
||||
|
||||
@app.get("/v1/pro/completions/status/{request_id}")
|
||||
async def get_pro_completion_status(request_id: str, auth: dict = Security(_authorize_request)):
|
||||
del auth
|
||||
|
||||
Reference in New Issue
Block a user