from fastapi import FastAPI, UploadFile, File
from starlette.responses import StreamingResponse

from models.agent import QueryRequest
from service import query_agent
from storage import upload_rustfs
from ocr import ocr_from_url

app = FastAPI(title="AI Study Service")


@app.post("/upload")
def upload(file: UploadFile = File(...)):
    url = upload_rustfs(file.file, file.filename)
    text = ocr_from_url(url)

    return {
        "text": text,
    }


@app.post("/generate")
def generate(query: QueryRequest):
    """流式对话"""
    return StreamingResponse(
        query_agent(query),
        media_type="text/event-stream"
    )