140 lines
5.0 KiB
Python
140 lines
5.0 KiB
Python
"""
|
|
HTR API - Simplified Local vLLM Version
|
|
Reads configuration from environment variables (.env file)
|
|
"""
|
|
|
|
import os
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel
|
|
import httpx
|
|
import base64
|
|
|
|
# ===========================================
|
|
# Configuration from Environment
|
|
# ===========================================
|
|
VLLM_ENDPOINT = os.getenv("VLLM_ENDPOINT", "http://127.0.0.1:8001/v1/chat/completions")
|
|
VLLM_MODEL_NAME = os.getenv("VLLM_MODEL_NAME")
|
|
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", 10 * 1024 * 1024))
|
|
|
|
# Fail fast if critical env vars aren't set
|
|
if not VLLM_MODEL_NAME:
|
|
raise RuntimeError("VLLM_MODEL_NAME not set! Check .env file and htr-api.service EnvironmentFile path.")
|
|
|
|
# ===========================================
|
|
# HTR Prompt
|
|
# ===========================================
|
|
HTR_PROMPT = """You are a handwriting transcription assistant.
|
|
|
|
Your task is to transcribe the handwritten pages into plain text paragraphs, preserving the original wording exactly.
|
|
|
|
Rules:
|
|
- Do not change, correct, or improve any spelling, grammar, punctuation, or sentence structure.
|
|
- Preserve paragraph breaks as they appear in the handwriting. Start a new paragraph only where the handwriting clearly starts a new paragraph that has a blank line. Do NOT insert line breaks on every sentence because the sentence is not fitting in the journal image. Just the paragraph break with the clear blank line between paragraphs.
|
|
- Also, insert 2 paragraphs breaks to clearly indicate a new paragraph.
|
|
- If you are unsure of a word or cannot read it, write ??? in its place.
|
|
- Do not add, remove, or rearrange any sentences.
|
|
|
|
Only output the transcribed text following these rules."""
|
|
|
|
# ===========================================
|
|
# FastAPI App
|
|
# ===========================================
|
|
app = FastAPI(
|
|
title="HTR API",
|
|
description="Handwriting Text Recognition API for Prometheus Cafe",
|
|
version="2.0.0"
|
|
)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# ===========================================
|
|
# Request/Response Models
|
|
# ===========================================
|
|
class ExtractRequest(BaseModel):
|
|
image: str # Base64 encoded image
|
|
|
|
class ExtractResponse(BaseModel):
|
|
text: str
|
|
|
|
# ===========================================
|
|
# Endpoints
|
|
# ===========================================
|
|
@app.get("/api/health")
|
|
async def health_check():
|
|
"""Health check endpoint"""
|
|
return {
|
|
"status": "healthy",
|
|
"backend": "local-vllm",
|
|
"model": VLLM_MODEL_NAME,
|
|
"endpoint": VLLM_ENDPOINT
|
|
}
|
|
|
|
@app.post("/api/extract", response_model=ExtractResponse)
|
|
async def extract_text(request: ExtractRequest):
|
|
"""Extract text from handwritten image using local vLLM."""
|
|
|
|
# Validate and clean image data
|
|
try:
|
|
image_data = request.image
|
|
if "," in image_data:
|
|
image_data = image_data.split(",")[1]
|
|
|
|
decoded = base64.b64decode(image_data)
|
|
if len(decoded) > MAX_IMAGE_SIZE:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"Image too large. Max size: {MAX_IMAGE_SIZE // 1024 // 1024}MB"
|
|
)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=400, detail=f"Invalid image data: {str(e)}")
|
|
|
|
# Call local vLLM
|
|
try:
|
|
async with httpx.AsyncClient(timeout=180.0) as client:
|
|
response = await client.post(
|
|
VLLM_ENDPOINT,
|
|
headers={"Content-Type": "application/json"},
|
|
json={
|
|
"model": VLLM_MODEL_NAME,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": HTR_PROMPT},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/jpeg;base64,{image_data}"
|
|
}
|
|
}
|
|
]
|
|
}
|
|
],
|
|
"max_tokens": 2048,
|
|
"temperature": 0.1,
|
|
}
|
|
)
|
|
|
|
if response.status_code != 200:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"vLLM error {response.status_code}: {response.text}"
|
|
)
|
|
|
|
result = response.json()
|
|
extracted_text = result["choices"][0]["message"]["content"]
|
|
|
|
return ExtractResponse(text=extracted_text)
|
|
|
|
except httpx.TimeoutException:
|
|
raise HTTPException(status_code=504, detail="vLLM request timed out")
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Extraction failed: {str(e)}")
|