Update app.py

app.py

@@ -25,56 +25,62 @@ class SetupRequest(BaseModel):
 # -------------------------------
 # Persistent Chroma client
 # -------------------------------
+CHROMA_PATH = "/data/chroma"
 COLLECTION_NAME = "knowledge_base"
 
-# Initialize in-memory client
-chroma_client = chromadb.Client()
-
-# Create or get collection
-kb_collection: Collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
-
-print("✅ In-memory Knowledge Base initialized successfully.")
-
 # -------------------------------
 # KB Setup Endpoint
 # -------------------------------
 @app.post("/setup")
-async def setup_endpoint(kb_file: UploadFile = File(...)):
+def setup_kb(json_file_path: str):
     """
-    Upload KB JSON file and embed it in ChromaDB.
-    JSON format: {"knowledge_base": [{"id": "1", "answer": "...", "category": "..."}]}
+    Loads a JSON file, generates embeddings, and populates a persistent ChromaDB knowledge base.
+    Only runs when explicitly called.
     """
-    global kb_collection
-    if not kb_collection:
-        raise HTTPException(status_code=500, detail="Chroma KB not initialized")
+    if not os.path.exists(json_file_path):
+        raise HTTPException(status_code=404, detail=f"File not found: {json_file_path}")
 
     try:
-        content = await kb_file.read()
-        kb_data = json.loads(content)
+        # Load JSON data
+        with open(json_file_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        if not isinstance(data, list):
+            raise HTTPException(status_code=400, detail="JSON must contain a list of knowledge items.")
+
+        print(f"📘 Loaded {len(data)} items from {json_file_path}")
+
+        # Initialize encoder and Chroma client (only here)
+        encoder = SentenceTransformer("all-MiniLM-L6-v2")
+        chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
+        collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
+
+        # Clear existing data (optional)
+        if collection.count() > 0:
+            print(f"🧹 Clearing {collection.count()} existing records...")
+            collection.delete(ids=collection.get()['ids'])
+
+        # Prepare data for embeddings
+        texts, ids, metadatas = [], [], []
+        for i, item in enumerate(data):
+            content = item.get("content") or item.get("text") or ""
+            title = item.get("title") or f"Document {i+1}"
+            texts.append(f"{title}. {content}")
+            ids.append(str(i))
+            metadatas.append(item)
+
+        print("🧠 Generating embeddings...")
+        embeddings = encoder.encode(texts, show_progress_bar=True).tolist()
+
+        # Add to ChromaDB
+        print("💾 Adding to ChromaDB...")
+        collection.add(ids=ids, embeddings=embeddings, metadatas=metadatas)
+        print(f"✅ Successfully added {collection.count()} records to {COLLECTION_NAME}.")
+
+        return {"message": "Knowledge base successfully initialized.", "count": collection.count()}
+
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Failed to read JSON: {e}")
-
-    if "knowledge_base" not in kb_data:
-        raise HTTPException(status_code=400, detail="Invalid KB format, missing 'knowledge_base' key")
-
-    # Clear existing entries before adding new KB
-    kb_collection.delete(where={"id": {"$ne": None}})  # Delete all previous entries
-
-    # Add KB entries
-    for entry in kb_data["knowledge_base"]:
-        kb_collection.add(
-            documents=[entry["answer"]],
-            metadatas=[{
-                "id": entry["id"],
-                "category": entry.get("category", ""),
-                "question_variations": entry.get("question_variations", []),
-                "keywords": entry.get("keywords", [])
-            }],
-            ids=[entry["id"]]
-        )
-
-    kb_collection.persist()
-    return {"status": f"KB uploaded and stored successfully, {len(kb_data['knowledge_base'])} entries added"}
+        raise HTTPException(status_code=500, detail=f"Setup failed: {e}")
 
 # -------------------------------
 # Step-by-Step Endpoints