chore: 更新项目配置和前后端代码优化
This commit is contained in:
+19
-11
@@ -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
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user