AI-markdown/services/unlimited-ocr-service/app/controllers/ConvertController.py

115 lines
3.6 KiB
Python

from fastapi import APIRouter, UploadFile, File, Depends, Query
from sqlalchemy.ext.asyncio import AsyncSession
from app.models.ConvertModel import ConvertResponse, HealthResponse, ConversionRecord
from app.services import UnlimitedOcrService as ocr_service
from app.database import get_db
from pydantic import BaseModel
router = APIRouter()
SUPPORTED_INPUT_FORMATS = sorted([
"pdf", "jpg", "jpeg", "png", "tiff", "tif", "bmp", "webp"
])
class SettingsRequest(BaseModel):
ocr_base_url: str | None = None
ocr_api_key: str | None = None
ocr_model: str = "Unlimited-OCR"
ocr_dpi: int = 300
class SettingsResponse(BaseModel):
ocr_enabled: bool
ocr_base_url: str | None
ocr_model: str
ocr_dpi: int
default_prompt: str | None = None
@router.get("/health", response_model=HealthResponse)
def health():
return HealthResponse(
status="ok",
supported_formats=SUPPORTED_INPUT_FORMATS,
output_formats=["markdown"],
ocr_enabled=ocr_service.OCR_ACTIVE,
ocr_model=ocr_service.OCR_MODEL if ocr_service.OCR_ACTIVE else None,
)
@router.get("/settings", response_model=SettingsResponse)
def get_settings():
return SettingsResponse(
ocr_enabled=ocr_service.OCR_ACTIVE,
ocr_base_url=ocr_service.OCR_BASE_URL,
ocr_model=ocr_service.OCR_MODEL,
ocr_dpi=ocr_service.OCR_DPI,
default_prompt=ocr_service.DEFAULT_OCR_PROMPT,
)
@router.post("/settings", response_model=SettingsResponse)
def update_settings(req: SettingsRequest):
ocr_service._init_ocr(req.ocr_base_url or None, req.ocr_model, req.ocr_api_key or "none")
ocr_service.OCR_DPI = req.ocr_dpi
return SettingsResponse(
ocr_enabled=ocr_service.OCR_ACTIVE,
ocr_base_url=ocr_service.OCR_BASE_URL,
ocr_model=ocr_service.OCR_MODEL,
ocr_dpi=ocr_service.OCR_DPI,
default_prompt=ocr_service.DEFAULT_OCR_PROMPT,
)
@router.post("/convert", response_model=ConvertResponse)
async def convert(
file: UploadFile = File(...),
llm_prompt: str | None = Query(default=None, description="Custom OCR prompt"),
db: AsyncSession = Depends(get_db),
):
record = await ocr_service.convert_file(file, db, llm_prompt=llm_prompt)
return ConvertResponse(
id=record.id,
filename=record.filename,
output_format=record.output_format,
content=record.content,
page_count=record.page_count,
llm_enabled=record.llm_enabled,
)
@router.get("/history", response_model=list[ConversionRecord])
async def history(limit: int = 20, db: AsyncSession = Depends(get_db)):
records = await ocr_service.get_history(db, limit)
return [
ConversionRecord(
id=r.id,
filename=r.filename,
file_type=r.file_type,
output_format=r.output_format,
page_count=r.page_count,
llm_enabled=r.llm_enabled,
created_at=str(r.created_at),
)
for r in records
]
@router.get("/conversions/{conversion_id}", response_model=ConvertResponse)
async def get_conversion(conversion_id: int, db: AsyncSession = Depends(get_db)):
record = await ocr_service.get_conversion(conversion_id, db)
return ConvertResponse(
id=record.id,
filename=record.filename,
output_format=record.output_format,
content=record.content,
page_count=record.page_count,
llm_enabled=record.llm_enabled,
)
@router.delete("/conversions/{conversion_id}")
async def delete_conversion(conversion_id: int, db: AsyncSession = Depends(get_db)):
return await ocr_service.delete_conversion(conversion_id, db)