59 lines
1.6 KiB
Python
59 lines
1.6 KiB
Python
from typing import List
|
|
|
|
from fastapi import FastAPI, UploadFile, File, Depends
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from starlette.middleware.cors import CORSMiddleware
|
|
from starlette.responses import StreamingResponse
|
|
|
|
from database import get_db
|
|
from schemas.agent import QueryRequest
|
|
from schemas.library import LibraryFileRequest, LibraryFileResponse
|
|
from service.agent_service import query_agent
|
|
from service.library_service import add_file, get_file, list_files
|
|
from storage import upload_rustfs
|
|
from ocr import ocr_from_url
|
|
|
|
app = FastAPI(title="AI Study Service")
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
@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"
|
|
)
|
|
|
|
|
|
@app.post("/files", response_model=bool)
|
|
async def add_library_file(data: LibraryFileRequest, db: AsyncSession = Depends(get_db)):
|
|
await add_file(db, data)
|
|
return True
|
|
|
|
|
|
@app.get("/files", response_model=List[LibraryFileResponse])
|
|
async def list_library_file(db: AsyncSession = Depends(get_db)):
|
|
return await list_files(db)
|
|
|
|
|
|
@app.get("/files/{file_id}", response_model=LibraryFileResponse)
|
|
async def get_library_file(file_id: int, db: AsyncSession = Depends(get_db)):
|
|
return await get_file(db, file_id)
|