Merge pull request #237 from tigergraph/GML-1726-invalid-key

GML-1726: handle error when openai API key is not valid

app/config.py

@@ -2,6 +2,7 @@
 import os
 
 from fastapi.security import HTTPBasic
+from pymilvus.exceptions import MilvusException
 
 from app.embeddings.embedding_services import (
     AWS_Bedrock_Embedding,
@@ -27,6 +28,7 @@
 security = HTTPBasic()
 session_handler = SessionHandler()
 status_manager = StatusManager()
+service_status = {}
 
 # Configs
 LLM_SERVICE = os.getenv("LLM_CONFIG", "configs/llm_config.json")
@@ -128,35 +130,51 @@ def get_llm_service(llm_config):
 )
 
 LogWriter.info("Setting up Milvus embedding store for InquiryAI")
-embedding_store = MilvusEmbeddingStore(
-    embedding_service,
-    host=milvus_config["host"],
-    port=milvus_config["port"],
-    collection_name="tg_inquiry_documents",
-    support_ai_instance=False,
-    username=milvus_config.get("username", ""),
-    password=milvus_config.get("password", ""),
-    alias=milvus_config.get("alias", "default"),
-)
+try:
+    embedding_store = MilvusEmbeddingStore(
+        embedding_service,
+        host=milvus_config["host"],
+        port=milvus_config["port"],
+        collection_name="tg_inquiry_documents",
+        support_ai_instance=False,
+        username=milvus_config.get("username", ""),
+        password=milvus_config.get("password", ""),
+        alias=milvus_config.get("alias", "default"),
+    )
+    service_status["embedding_store"] = {"status": "ok", "error": None}
+except MilvusException as e:
+    embedding_store = None
+    service_status["embedding_store"] = {"status": "milvus error", "error": str(e)}
+except Exception as e:
+    embedding_store = None
+    service_status["embedding_store"] = {"status": "embedding error", "error": str(e)}
 
 support_collection_name = milvus_config.get("collection_name", "tg_support_documents")
 LogWriter.info(
     f"Setting up Milvus embedding store for SupportAI with collection_name: {support_collection_name}"
 )
 vertex_field = milvus_config.get("vertex_field", "vertex_id")
-support_ai_embedding_store = MilvusEmbeddingStore(
-    embedding_service,
-    host=milvus_config["host"],
-    port=milvus_config["port"],
-    support_ai_instance=True,
-    collection_name=support_collection_name,
-    username=milvus_config.get("username", ""),
-    password=milvus_config.get("password", ""),
-    vector_field=milvus_config.get("vector_field", "document_vector"),
-    text_field=milvus_config.get("text_field", "document_content"),
-    vertex_field=vertex_field,
-    alias=milvus_config.get("alias", "default"),
-)
+try:
+    support_ai_embedding_store = MilvusEmbeddingStore(
+        embedding_service,
+        host=milvus_config["host"],
+        port=milvus_config["port"],
+        support_ai_instance=True,
+        collection_name=support_collection_name,
+        username=milvus_config.get("username", ""),
+        password=milvus_config.get("password", ""),
+        vector_field=milvus_config.get("vector_field", "document_vector"),
+        text_field=milvus_config.get("text_field", "document_content"),
+        vertex_field=vertex_field,
+        alias=milvus_config.get("alias", "default"),
+    )
+    service_status["support_ai_embedding_store"] = {"status": "ok", "error": None}
+except MilvusException as e:
+    support_ai_embedding_store = None
+    service_status["support_ai_embedding_store"] = {"status": "milvus error", "error": str(e)}
+except Exception as e:
+    support_ai_embedding_store = None
+    service_status["support_ai_embedding_store"] = {"status": "embedding error", "error": str(e)}
 
 
 if DOC_PROCESSING_CONFIG is None or (

app/routers/inquiryai.py

@@ -9,7 +9,7 @@
 from fastapi.security.http import HTTPBase
 
 from app.agent import TigerGraphAgent
-from app.config import embedding_service, embedding_store, llm_config, session_handler
+from app.config import embedding_service, embedding_store, llm_config, session_handler, service_status
 from app.llm_services import (
     AWS_SageMaker_Endpoint,
     AWSBedrock,
@@ -45,10 +45,14 @@ def retrieve_answer(
     conn: Request,
     credentials: Annotated[HTTPBase, Depends(security)]
 ) -> CoPilotResponse:
-    conn = conn.state.conn
     logger.debug_pii(
         f"/{graphname}/query request_id={req_id_cv.get()} question={query.query}"
     )
+    if service_status["embedding_store"]["error"]:
+        return CoPilotResponse(
+            natural_language_response="Something wrong with CoPilot's embedding store. It could be caused by an invalid API key to the embedding service. Please contact support to check the system health for details.", answered_question=False, response_type="inquiryai"
+        )
+    conn = conn.state.conn
     logger.debug(
         f"/{graphname}/query request_id={req_id_cv.get()} database connection created"
     )

app/routers/root.py

@@ -3,7 +3,7 @@
 from fastapi import APIRouter, Response
 from fastapi.responses import FileResponse
 
-from app.config import llm_config
+from app.config import llm_config, service_status
 
 from pymilvus import connections, utility
 
@@ -18,26 +18,12 @@ def read_root():
 
 @router.get("/health")
 async def health():
-    # Check if Milvus is up and running and if the required collections exist
-    connections.connect(host="milvus-standalone", port="19530")
-
-    try:
-        # Check if the required collections exist
-        inquiry_collection_exists = utility.has_collection("tg_inquiry_documents")
-        support_collection_exists = utility.has_collection("tg_support_documents")
-
-        if inquiry_collection_exists or support_collection_exists:
-            return {
-                "status": "healthy",
-                "llm_completion_model": llm_config["completion_service"]["llm_model"],
-                "embedding_service": llm_config["embedding_service"][
-                    "embedding_model_service"
-                ],
-            }
-        else:
-            return {"status": "Milvus is up and running, but no collection exist yet"}
-    except Exception as e:
-        return {"status": "Error checking Milvus health", "error": str(e)}
+    status = {
+        "status": "unhealthy" if any(v["error"] is not None for v in service_status.values()) else "healthy",
+        "details": service_status
+    }
+
+    return status
 
 
 @router.get("/metrics")