100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
from db_setup import get_qdrant_client
|
|
from typing import Annotated
|
|
from fastapi import Depends, HTTPException, APIRouter
|
|
from qdrant_client import AsyncQdrantClient
|
|
from fastapi.responses import JSONResponse
|
|
from .serializers import (
|
|
CreateCollectionSerializer,
|
|
QueryCollectionSerializer,
|
|
UpdateCollectionSerializer,
|
|
DeleteCollectionSerializer
|
|
)
|
|
from .models import CollectionHandler
|
|
|
|
app_router = APIRouter()
|
|
|
|
@app_router.post("/create")
|
|
async def create_collection_endpoint(
|
|
q: Annotated[AsyncQdrantClient, Depends(get_qdrant_client)],
|
|
body: CreateCollectionSerializer = None
|
|
):
|
|
try:
|
|
if body is None:
|
|
raise HTTPException(status_code=400, detail="Collection name is required")
|
|
|
|
print("collection_name: ", body.collection_name)
|
|
|
|
handler = CollectionHandler(
|
|
collection_name=body.collection_name,
|
|
vector=body.vector,
|
|
vector_size=body.vector_size,
|
|
payload=body.payload,
|
|
id=body.id
|
|
)
|
|
|
|
# 1. Create collection
|
|
result = await handler.create_collection()
|
|
|
|
# 2. Automatically call upsert_point
|
|
await handler.upsert_point()
|
|
|
|
return JSONResponse(result)
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
@app_router.get("/query")
|
|
async def query_collection_endpoint(
|
|
q: Annotated[AsyncQdrantClient, Depends(get_qdrant_client)],
|
|
body: QueryCollectionSerializer
|
|
):
|
|
try:
|
|
handler = CollectionHandler(
|
|
collection_name=body.collection_name,
|
|
vector=body.query_vector,
|
|
vector_size=len(body.query_vector),
|
|
payload={},
|
|
id=0
|
|
)
|
|
result = await handler.search(body.query_vector)
|
|
return JSONResponse({"results": str(result)})
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
@app_router.put("/update")
|
|
async def update_collection_endpoint(
|
|
q: Annotated[AsyncQdrantClient, Depends(get_qdrant_client)],
|
|
body: UpdateCollectionSerializer
|
|
):
|
|
try:
|
|
handler = CollectionHandler(
|
|
collection_name=body.collection_name,
|
|
vector=body.vector,
|
|
vector_size=len(body.vector),
|
|
payload=body.payload,
|
|
id=body.id
|
|
)
|
|
result = await handler.update_collection()
|
|
return JSONResponse({"status": "success", "result": result})
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
@app_router.delete("/delete")
|
|
async def delete_collection_endpoint(
|
|
q: Annotated[AsyncQdrantClient, Depends(get_qdrant_client)],
|
|
body: DeleteCollectionSerializer
|
|
):
|
|
try:
|
|
handler = CollectionHandler(
|
|
collection_name=body.collection_name,
|
|
vector=[],
|
|
vector_size=0,
|
|
payload={},
|
|
id=0
|
|
)
|
|
result = await handler.delete_collection()
|
|
return JSONResponse(result)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|