Преглед изворни кода

feat: vectordb query error handling

Timothy J. Baek пре 1 година
родитељ
комит
142269374f
1 измењених фајлова са 12 додато и 5 уклоњено
  1. 12 5
      backend/apps/rag/main.py

+ 12 - 5
backend/apps/rag/main.py

@@ -85,12 +85,19 @@ async def get_status():
 
 
 @app.get("/query/{collection_name}")
 @app.get("/query/{collection_name}")
 def query_collection(collection_name: str, query: str, k: Optional[int] = 4):
 def query_collection(collection_name: str, query: str, k: Optional[int] = 4):
-    collection = CHROMA_CLIENT.get_collection(
-        name=collection_name,
-    )
-    result = collection.query(query_texts=[query], n_results=k)
+    try:
+        collection = CHROMA_CLIENT.get_collection(
+            name=collection_name,
+        )
+        result = collection.query(query_texts=[query], n_results=k)
 
 
-    return result
+        return result
+    except Exception as e:
+        print(e)
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=ERROR_MESSAGES.DEFAULT(e),
+        )
 
 
 
 
 @app.post("/web")
 @app.post("/web")