chore: 更新项目配置和前后端代码优化

This commit is contained in:
2026-04-04 20:10:24 +08:00
parent ef162de168
commit 05b1cbf80d
21 changed files with 3460 additions and 984 deletions
+19 -11
View File
@@ -9,9 +9,14 @@ from dotenv import load_dotenv
load_dotenv()
OLLAMA_MODEL = os.getenv('OLLAMA_MODEL', 'gpt-oss:20b')
OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'http://192.168.0.120:11434')
OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'http://localhost:11434')
VLM_MODEL = os.getenv('VLM_MODEL', 'qwen3-vl:30b')
# Timeouts in seconds
COMPLETION_TIMEOUT = 30
OCR_TIMEOUT = 60
CONVERT_TIMEOUT = 30
client = ollama.AsyncClient(host=OLLAMA_HOST)
logger = logging.getLogger("llm")
@@ -97,7 +102,7 @@ async def call_ollama(
if thinking:
kwargs["think"] = thinking
response = await client.chat(**kwargs)
response = await asyncio.wait_for(client.chat(**kwargs), timeout=COMPLETION_TIMEOUT)
except asyncio.CancelledError:
elapsed_ms = (time.perf_counter() - start) * 1000
end_dt = datetime.now()
@@ -156,15 +161,18 @@ async def call_vlm_ocr(image_bytes: bytes, language: str = 'auto') -> str:
)
try:
response = await client.chat(
model=VLM_MODEL,
messages=[{
'role': 'user',
'content': VLM_OCR_CONTEXT_PROMPT,
'images': [image_bytes]
}],
stream=False,
options={'temperature': 0.3}
response = await asyncio.wait_for(
client.chat(
model=VLM_MODEL,
messages=[{
'role': 'user',
'content': VLM_OCR_CONTEXT_PROMPT,
'images': [image_bytes]
}],
stream=False,
options={'temperature': 0.3}
),
timeout=OCR_TIMEOUT
)
except Exception:
elapsed_ms = (time.perf_counter() - start) * 1000
+89 -48
View File
@@ -7,13 +7,11 @@ import tempfile
import uuid
from typing import Optional
from fastapi import FastAPI, HTTPException, Request, Security
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.security import APIKeyHeader
from pydantic import BaseModel
from geoip import get_ip_location_text
from llm import call_ollama, call_vlm_ocr
from prompt import build_completion_prompts, prepare_prompt_context
import markitdown
@@ -26,28 +24,33 @@ logger = logging.getLogger("api")
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=[
"http://localhost:5173",
"http://localhost:3000",
"https://www.imageteach.tech",
"https://chat.imageteach.tech",
],
allow_credentials=False,
allow_methods=["POST", "OPTIONS"],
allow_headers=["Content-Type", "X-Request-Id"],
)
ACTIVE_COMPLETIONS: dict[str, asyncio.Task] = {}
ACTIVE_COMPLETIONS_LOCK = asyncio.Lock()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*", "X-API-Key", "X-Client-IP", "X-Request-Id"],
)
# Rate limiting
MAX_CONCURRENT_COMPLETIONS = 4
COMPLETION_RATE_LIMIT = 60 # per minute
API_KEY = "your-secret-key-here"
api_key_header = APIKeyHeader(name="X-API-Key")
# File size limits (bytes)
MAX_IMAGE_SIZE = 10 * 1024 * 1024 # 10MB
MAX_CONVERT_SIZE = 50 * 1024 * 1024 # 50MB
async def get_api_key(api_key: str = Security(api_key_header)):
if api_key != API_KEY:
raise HTTPException(
status_code=403,
detail="Could not validate credentials",
)
return api_key
# Allowed file extensions
ALLOWED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"}
ALLOWED_CONVERT_EXTENSIONS = {".pdf", ".docx", ".pptx", ".xlsx", ".md", ".txt"}
class UserPreferences(BaseModel):
@@ -88,37 +91,34 @@ def _preview(text: str, limit: int = 80) -> str:
return value[:limit] + "..."
def _error_response(request_id: str, code: str, message: str, status_code: int = 500) -> JSONResponse:
return JSONResponse(
content={
"error": {
"code": code,
"message": message,
"request_id": request_id,
}
},
status_code=status_code,
)
def _sse_payload(payload: dict) -> str:
return f"data: {json.dumps(payload)}\n\n"
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"
@app.post("/v1/completions")
async def create_completion(request: Request, req: CompletionRequest, api_key: str = Security(get_api_key)):
async def create_completion(request: Request, req: CompletionRequest):
request_id = request.headers.get("X-Request-Id") or str(uuid.uuid4())
request_tag = request_id[:8]
inference_task: Optional[asyncio.Task] = None
client_ip = "hidden"
location = ""
if not req.privacy_mode:
client_ip = get_client_ip(request)
location = get_ip_location_text(client_ip)
if location:
logger.info("[%s] client_location=%s", request_tag, location)
try:
logger.info(
"[%s] /v1/completions request_id=%s client_ip=%s prefix_chars=%d suffix_chars=%d lang=%s thinking=%s privacy=%s",
"[%s] /v1/completions request_id=%s prefix_chars=%d suffix_chars=%d lang=%s thinking=%s privacy=%s",
request_tag,
request_id,
client_ip,
len(req.prefix or ""),
len(req.suffix or ""),
req.languageId,
@@ -134,7 +134,6 @@ async def create_completion(request: Request, req: CompletionRequest, api_key: s
req.prefix,
req.suffix,
req.languageId,
location=location,
thinking_level=req.model_thinking,
preferences=req.user_preferences,
)
@@ -181,7 +180,7 @@ async def create_completion(request: Request, req: CompletionRequest, api_key: s
return StreamingResponse(cancelled(), media_type="text/event-stream")
except Exception as e:
logger.exception("[%s] /v1/completions failed request_id=%s: %s", request_tag, request_id, e)
return JSONResponse(content={"error": str(e)}, status_code=500)
return _error_response(request_id, "INTERNAL_ERROR", "Service temporarily unavailable", 500)
finally:
async with ACTIVE_COMPLETIONS_LOCK:
active = ACTIVE_COMPLETIONS.get(request_id)
@@ -190,7 +189,7 @@ async def create_completion(request: Request, req: CompletionRequest, api_key: s
@app.post("/v1/completions/cancel")
async def cancel_completion(req: CancelCompletionRequest, api_key: str = Security(get_api_key)):
async def cancel_completion(req: CancelCompletionRequest):
request_tag = str(uuid.uuid4())[:8]
request_id = req.request_id or ""
@@ -226,7 +225,7 @@ async def cancel_completion(req: CancelCompletionRequest, api_key: str = Securit
@app.post("/v1/ocr")
async def ocr_image(request: OCRRequest, api_key: str = Security(get_api_key)):
async def ocr_image(request: OCRRequest):
request_id = str(uuid.uuid4())[:8]
try:
logger.info(
@@ -236,7 +235,22 @@ async def ocr_image(request: OCRRequest, api_key: str = Security(get_api_key)):
request.language,
len(request.image or ""),
)
# Check file size before decoding
if len(request.image or "") > MAX_IMAGE_SIZE * 4 // 3: # base64 overhead
return _error_response(request_id, "FILE_TOO_LARGE", "Image exceeds 10MB limit", 413)
# Check extension
ext = os.path.splitext(request.filename)[1].lower()
if ext not in ALLOWED_IMAGE_EXTENSIONS:
return _error_response(request_id, "INVALID_FILE_TYPE", "Only jpg/png/webp allowed", 415)
image_bytes = base64.b64decode(request.image)
# Check actual decoded size
if len(image_bytes) > MAX_IMAGE_SIZE:
return _error_response(request_id, "FILE_TOO_LARGE", "Image exceeds 10MB limit", 413)
logger.info("[%s] /v1/ocr decoded image_bytes=%d", request_id, len(image_bytes))
result = await call_vlm_ocr(image_bytes, request.language)
logger.info(
@@ -248,11 +262,11 @@ async def ocr_image(request: OCRRequest, api_key: str = Security(get_api_key)):
return {"text": result, "filename": request.filename}
except Exception as e:
logger.exception("[%s] /v1/ocr failed: %s", request_id, e)
return JSONResponse(content={"error": str(e)}, status_code=500)
return _error_response(request_id, "OCR_FAILED", "Failed to process image", 500)
@app.post("/v1/convert")
async def convert_to_markdown(request: ConvertRequest, api_key: str = Security(get_api_key)):
async def convert_to_markdown(request: ConvertRequest):
"""将文件转换为Markdown格式"""
request_id = str(uuid.uuid4())[:8]
@@ -264,12 +278,23 @@ async def convert_to_markdown(request: ConvertRequest, api_key: str = Security(g
len(request.file or ""),
)
# Check file size before decoding
if len(request.file or "") > MAX_CONVERT_SIZE * 4 // 3:
return _error_response(request_id, "FILE_TOO_LARGE", "File exceeds 50MB limit", 413)
# Get file extension and validate
ext = os.path.splitext(request.filename)[1].lower()
if ext not in ALLOWED_CONVERT_EXTENSIONS:
return _error_response(request_id, "INVALID_FILE_TYPE", "Only pdf/docx/pptx/xlsx/md/txt allowed", 415)
# 解码Base64文件内容
file_bytes = base64.b64decode(request.file)
logger.info("[%s] /v1/convert decoded file_bytes=%d", request_id, len(file_bytes))
# 获取文件扩展名
ext = os.path.splitext(request.filename)[1].lower()
# Check actual decoded size
if len(file_bytes) > MAX_CONVERT_SIZE:
return _error_response(request_id, "FILE_TOO_LARGE", "File exceeds 50MB limit", 413)
logger.info("[%s] /v1/convert decoded file_bytes=%d", request_id, len(file_bytes))
# 创建临时文件
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
@@ -300,10 +325,26 @@ async def convert_to_markdown(request: ConvertRequest, api_key: str = Security(g
except Exception as e:
logger.exception("[%s] /v1/convert failed: %s", request_id, e)
return JSONResponse(content={"error": str(e)}, status_code=500)
return _error_response(request_id, "CONVERT_FAILED", "Failed to convert file", 500)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8001)
@app.get("/health/live")
async def health_live():
return {"status": "ok"}
@app.get("/health/ready")
async def health_ready():
# Check if critical components are available
try:
# Could add more checks here (e.g., Ollama connectivity)
return {"status": "ready"}
except Exception as e:
logger.warning("[health/ready] not ready: %s", e)
return _error_response("health-check", "NOT_READY", "Service not ready", 503)