Update app.py

app.py

@@ -30,20 +30,23 @@ kb_collection = None
 # KB Setup Endpoint
 # -------------------------------
 @app.post("/setup")
-async def setup_endpoint(req: SetupRequest):
+async def setup_endpoint(kb_file: UploadFile = File(...)):
     global kb_collection
 
-    if not os.path.exists(req.kb_file):
-        raise HTTPException(status_code=404, detail="KB.json file not found")
+    if not kb_file.filename.endswith(".json"):
+        raise HTTPException(status_code=400, detail="Only JSON files are allowed")
 
-    # Load KB.json
-    with open(req.kb_file, "r") as f:
-        kb_data = json.load(f)
+    # Read uploaded KB.json
+    contents = await kb_file.read()
+    try:
+        kb_data = json.loads(contents)
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid JSON file: {e}")
 
-    # Create Chroma client (new API)
+    # Initialize Chroma client
     chroma_client = chromadb.Client(
         Settings(
-            persist_directory="/data/chroma_db",  # make sure this dir exists and writable
+            persist_directory="/data/chroma_db",
             anonymized_telemetry=False
         )
     )
@@ -52,12 +55,14 @@ async def setup_endpoint(req: SetupRequest):
     try:
         kb_collection = chroma_client.get_collection("kb_collection")
     except:
-        kb_collection = chroma_client.create_collection("kb_collection")
-
-    # Embedding function
-    embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
-        model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1"
-    )
+        # Create collection with embedding function (new API)
+        embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
+            model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1"
+        )
+        kb_collection = chroma_client.create_collection(
+            name="kb_collection",
+            embedding_function=embedding_func
+        )
 
     # Add KB entries
     for entry in kb_data["knowledge_base"]:
@@ -69,14 +74,13 @@ async def setup_endpoint(req: SetupRequest):
                 "question_variations": entry.get("question_variations", []),
                 "keywords": entry.get("keywords", [])
             }],
-            ids=[entry["id"]],
-            embedding_function=embedding_func
+            ids=[entry["id"]]
         )
 
-    # Persist to disk
+    # Persist collection
     kb_collection.persist()
 
-    return {"status": "KB embedded and stored successfully"}
+    return {"status": "KB uploaded, embedded, and stored successfully"}
 
 # -------------------------------
 # Step-by-Step Endpoints