From 142269374f2add70084eca5f163afadf8039c1be Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Sun, 7 Jan 2024 01:59:00 -0800
Subject: [PATCH] feat: vectordb query error handling

---
 backend/apps/rag/main.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/backend/apps/rag/main.py b/backend/apps/rag/main.py
index 03073bf69..ff4bbb6ce 100644
--- a/backend/apps/rag/main.py
+++ b/backend/apps/rag/main.py
@@ -85,12 +85,19 @@ async def get_status():
 
 @app.get("/query/{collection_name}")
 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")