Browse Source

fix: error handling

Timothy J. Baek 1 year ago
parent
commit
784ee6f521
1 changed files with 20 additions and 14 deletions
  1. 20 14
      backend/main.py

+ 20 - 14
backend/main.py

@@ -108,20 +108,26 @@ class RAGMiddleware(BaseHTTPMiddleware):
 
 
                 for doc in docs:
                 for doc in docs:
                     context = None
                     context = None
-                    if doc["type"] == "collection":
-                        context = query_collection(
-                            collection_names=doc["collection_names"],
-                            query=query,
-                            k=rag_app.state.TOP_K,
-                            embedding_function=rag_app.state.sentence_transformer_ef,
-                        )
-                    else:
-                        context = query_doc(
-                            collection_name=doc["collection_name"],
-                            query=query,
-                            k=rag_app.state.TOP_K,
-                            embedding_function=rag_app.state.sentence_transformer_ef,
-                        )
+
+                    try:
+                        if doc["type"] == "collection":
+                            context = query_collection(
+                                collection_names=doc["collection_names"],
+                                query=query,
+                                k=rag_app.state.TOP_K,
+                                embedding_function=rag_app.state.sentence_transformer_ef,
+                            )
+                        else:
+                            context = query_doc(
+                                collection_name=doc["collection_name"],
+                                query=query,
+                                k=rag_app.state.TOP_K,
+                                embedding_function=rag_app.state.sentence_transformer_ef,
+                            )
+                    except Exception as e:
+                        print(e)
+                        context = None
+
                     relevant_contexts.append(context)
                     relevant_contexts.append(context)
 
 
                 context_string = ""
                 context_string = ""