From 784ee6f52183ab2158359d2bf75535ca7ad740d2 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Fri, 8 Mar 2024 23:21:00 -0800 Subject: [PATCH] fix: error handling --- backend/main.py | 34 ++++++++++++++++++++-------------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/backend/main.py b/backend/main.py index 11ca81fc..e63f91a0 100644 --- a/backend/main.py +++ b/backend/main.py @@ -108,20 +108,26 @@ class RAGMiddleware(BaseHTTPMiddleware): for doc in docs: 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) context_string = ""