Update app.py

app.py

@@ -36,12 +36,8 @@ kb_collection = None
 @app.post("/setup")
 async def setup_kb(kb_file: UploadFile = File(...)):
     global kb_collection
-    """
-    Uploads a JSON file, generates embeddings, and populates a persistent ChromaDB knowledge base.
-    Only runs when explicitly called.
-    """
     try:
-        # Read JSON directly from the uploaded file
+        # Read JSON directly
         content_bytes = await kb_file.read()
         data = json.loads(content_bytes)
 
@@ -50,38 +46,42 @@ async def setup_kb(kb_file: UploadFile = File(...)):
 
         print(f"📘 Loaded {len(data)} items from uploaded file '{kb_file.filename}'")
 
-        # Initialize encoder and Chroma client (only here)
+        # Initialize Sentence Transformer
         encoder = SentenceTransformer("all-MiniLM-L6-v2")
-        chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
-        kb_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
-        collection = kb_collection
 
-        # Clear existing data (optional)
-        if collection.count() > 0:
-            print(f"🧹 Clearing {collection.count()} existing records...")
-            collection.delete(ids=collection.get()['ids'])
+        # Initialize ChromaDB (persistent)
+        chroma_client = chromadb.PersistentClient(path="/tmp/chroma")  # must be writable
+        kb_collection = chroma_client.get_or_create_collection(
+            name="knowledge_base",
+            embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
+                model_name="all-MiniLM-L6-v2"
+            )
+        )
 
-        # Prepare data for embeddings
+        # Clear old data
+        if kb_collection.count() > 0:
+            kb_collection.delete(ids=kb_collection.get()['ids'])
+
+        # Prepare data
         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}"
+            content = item.get("answer") or item.get("content") or ""
+            title = item.get("id") or f"Document {i+1}"
             texts.append(f"{title}. {content}")
             ids.append(str(i))
-            metadatas.append(item)
+            # Only keep str/int/bool/float for metadata
+            metadata = {k: str(v) if isinstance(v, list) else v for k, v in item.items()}
+            metadatas.append(metadata)
 
-        print("🧠 Generating embeddings...")
+        # Generate 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}.")
+        # Add to Chroma
+        kb_collection.add(ids=ids, embeddings=embeddings, metadatas=metadatas)
+        print(f"✅ Successfully added {kb_collection.count()} records.")
 
-        return {"message": "Knowledge base successfully initialized.", "count": collection.count()}
+        return {"message": "Knowledge base successfully initialized.", "count": kb_collection.count()}
 
-    except json.JSONDecodeError:
-        raise HTTPException(status_code=400, detail="Invalid JSON file.")
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Setup failed: {e}")