Spaces:
Runtime error
Runtime error
File size: 6,344 Bytes
524e6d5 9ed3424 524e6d5 9ed3424 524e6d5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | import uuid
from pathlib import Path
import aiofiles
from arq import ArqRedis
from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile
from fastapi.responses import FileResponse
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.api.collections import get_owned_collection
from app.api.schemas import DocumentOut, SearchHit, SearchResponse, UrlIngest
from app.core.config import get_settings
from app.core.security import Principal, get_principal
from app.db.models import Document
from app.db.session import get_db
from app.ingest.parsers import detect_source_type
from app.retrieval import vectorstore
router = APIRouter(prefix="/collections/{collection_id}", tags=["documents"])
def _queue(request: Request) -> ArqRedis:
return request.app.state.arq
async def _enqueue(request: Request, document_id: str) -> None:
await _queue(request).enqueue_job("ingest_document", document_id)
@router.post("/documents", response_model=DocumentOut, status_code=202)
async def upload_document(
collection_id: str,
file: UploadFile,
request: Request,
db: AsyncSession = Depends(get_db),
principal: Principal = Depends(get_principal),
):
await get_owned_collection(collection_id, db, principal)
settings = get_settings()
source_type = detect_source_type(file.filename or "")
upload_dir = Path(settings.upload_dir) / principal.tenant_id
upload_dir.mkdir(parents=True, exist_ok=True)
dest = upload_dir / f"{uuid.uuid4()}{Path(file.filename or '').suffix.lower()}"
size = 0
max_bytes = settings.max_upload_mb * 1024 * 1024
async with aiofiles.open(dest, "wb") as out:
while chunk := await file.read(1024 * 1024):
size += len(chunk)
if size > max_bytes:
await out.close()
dest.unlink(missing_ok=True)
raise HTTPException(status_code=413, detail=f"File over {settings.max_upload_mb}MB")
await out.write(chunk)
doc = Document(
tenant_id=principal.tenant_id,
collection_id=collection_id,
name=file.filename or dest.name,
source_type=source_type,
source_uri=str(dest),
)
db.add(doc)
await db.commit()
await db.refresh(doc)
await _enqueue(request, doc.id)
return doc
@router.post("/urls", response_model=DocumentOut, status_code=202)
async def ingest_url(
collection_id: str,
body: UrlIngest,
request: Request,
db: AsyncSession = Depends(get_db),
principal: Principal = Depends(get_principal),
):
await get_owned_collection(collection_id, db, principal)
doc = Document(
tenant_id=principal.tenant_id,
collection_id=collection_id,
name=str(body.url),
source_type="url",
source_uri=str(body.url),
)
db.add(doc)
await db.commit()
await db.refresh(doc)
await _enqueue(request, doc.id)
return doc
@router.get("/documents", response_model=list[DocumentOut])
async def list_documents(
collection_id: str,
db: AsyncSession = Depends(get_db),
principal: Principal = Depends(get_principal),
):
await get_owned_collection(collection_id, db, principal)
rows = await db.scalars(
select(Document)
.where(Document.collection_id == collection_id)
.order_by(Document.created_at.desc())
)
return list(rows)
@router.delete("/documents/{document_id}", status_code=204)
async def delete_document(
collection_id: str,
document_id: str,
db: AsyncSession = Depends(get_db),
principal: Principal = Depends(get_principal),
):
await get_owned_collection(collection_id, db, principal)
doc = await db.get(Document, document_id)
if doc is None or doc.tenant_id != principal.tenant_id:
raise HTTPException(status_code=404, detail="Document not found")
vectorstore.delete_document(document_id)
if doc.source_type != "url":
Path(doc.source_uri).unlink(missing_ok=True)
await db.delete(doc)
await db.commit()
@router.get("/documents/{document_id}/file")
async def get_document_file(
collection_id: str,
document_id: str,
db: AsyncSession = Depends(get_db),
principal: Principal = Depends(get_principal),
):
await get_owned_collection(collection_id, db, principal)
doc = await db.get(Document, document_id)
if doc is None or doc.tenant_id != principal.tenant_id or doc.source_type == "url":
raise HTTPException(status_code=404, detail="File not found")
path = Path(doc.source_uri)
if not path.is_file():
raise HTTPException(status_code=404, detail="File missing from storage")
return FileResponse(path, filename=doc.name)
@router.get("/images/search")
async def search_images(
collection_id: str,
q: str,
limit: int = 12,
db: AsyncSession = Depends(get_db),
principal: Principal = Depends(get_principal),
):
"""Text-to-image semantic search over CLIP embeddings."""
await get_owned_collection(collection_id, db, principal)
points = vectorstore.image_search(q, principal.tenant_id, collection_id, limit=limit)
return {
"query": q,
"hits": [
{
"score": p.score,
"document_id": str(p.payload.get("document_id", "")),
"document_name": str(p.payload.get("document_name", "")),
}
for p in points
if p.payload
],
}
@router.get("/search", response_model=SearchResponse)
async def search(
collection_id: str,
q: str,
limit: int = 10,
db: AsyncSession = Depends(get_db),
principal: Principal = Depends(get_principal),
):
"""Direct hybrid retrieval — debugging/eval endpoint; chat agent uses same path."""
await get_owned_collection(collection_id, db, principal)
points = vectorstore.hybrid_search(q, principal.tenant_id, collection_id, limit=limit)
hits = [
SearchHit(
score=p.score,
text=str(p.payload.get("text", "")),
document_id=str(p.payload.get("document_id", "")),
document_name=str(p.payload.get("document_name", "")),
chunk_index=int(p.payload.get("chunk_index", 0)),
)
for p in points
if p.payload
]
return SearchResponse(query=q, hits=hits)
|