Update app.py

app.py

@@ -4,46 +4,80 @@ from sentence_transformers import SentenceTransformer
 import google.generativeai as genai
 import os
 
-# Get API key from Hugging Face secrets
+# Get API key
 GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
 genai.configure(api_key=GOOGLE_API_KEY)
 
-# Load models (cached)
+# Load models
 print("Loading embedding model...")
 embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
-print("Loading database...")
+print("Checking database...")
+
+# Create or load database
 chroma_client = chromadb.PersistentClient(path="./deadcells_db_free")
-collection = chroma_client.get_collection(name="deadcells_wiki")
-print("Ready!")
 
-# Find Gemini model
+try:
+    collection = chroma_client.get_collection(name="deadcells_wiki")
+    print(f"Database loaded! {collection.count()} chunks available")
+except:
+    print("Database not found! Building from wiki files...")
+    
+    # Create collection
+    collection = chroma_client.create_collection(name="deadcells_wiki")
+    
+    # Load wiki files
+    import glob
+    wiki_files = glob.glob("wiki_content/*.txt")
+    
+    if not wiki_files:
+        raise Exception("No wiki files found! Upload wiki_content folder")
+    
+    chunk_id = 0
+    for filepath in wiki_files:
+        filename = os.path.basename(filepath)
+        with open(filepath, 'r', encoding='utf-8') as f:
+            content = f.read()
+        
+        lines = content.split('\n')
+        url = lines[0].replace('URL: ', '') if lines[0].startswith('URL:') else ''
+        text = '\n'.join(lines[2:])
+        
+        # Simple chunking
+        words = text.split()
+        for i in range(0, len(words), 800):
+            chunk = ' '.join(words[i:i + 1000])
+            if len(chunk.split()) < 50:
+                continue
+                
+            embedding = embedding_model.encode(chunk).tolist()
+            
+            collection.add(
+                documents=[chunk],
+                embeddings=[embedding],
+                metadatas=[{"source": filename, "url": url, "chunk_index": i}],
+                ids=[f"chunk-{chunk_id}"]
+            )
+            chunk_id += 1
+    
+    print(f"Database created! {chunk_id} chunks indexed")
+
 available_model = 'gemini-1.5-flash-latest'
 
 def get_embedding(text):
     return embedding_model.encode(text).tolist()
 
 def chat(message, history):
-    """Handle chat messages"""
-    
-    # Search wiki
     question_embedding = get_embedding(message)
-    results = collection.query(
-        query_embeddings=[question_embedding],
-        n_results=5
-    )
+    results = collection.query(query_embeddings=[question_embedding], n_results=5)
     
     relevant_chunks = results['documents'][0]
-    sources = results['metadatas'][0]
-    
     if not relevant_chunks:
         return "I couldn't find any relevant information in the wiki."
     
-    # Build context
     context = "\n\n---\n\n".join(relevant_chunks)
     
-    # Ask Gemini
     model = genai.GenerativeModel(available_model)
-    prompt = f"""You are a Dead Cells expert. Answer using ONLY the provided wiki content.
+    prompt = f"""You are a Dead Cells expert. Answer using ONLY the wiki content provided.
 
 Wiki Content:
 {context}
@@ -53,29 +87,18 @@ Question: {message}
 Answer:"""
     
     response = model.generate_content(prompt)
-    answer = response.text
-    
-    # Add sources
-    source_list = "\n\n📚 **Sources:** " + ", ".join([s['source'] for s in sources[:3]])
-    
-    return answer + source_list
+    return response.text
 
-# Create Gradio interface
 demo = gr.ChatInterface(
     fn=chat,
     title="🎮 Dead Cells Wiki Bot",
-    description="Ask me anything about Dead Cells! I answer using the official wiki.",
+    description="Ask me anything about Dead Cells!",
     examples=[
-        "What achievements are there for beating bosses?",
-        "Tell me about the Hand of the King",
+        "What achievements are there?",
+        "Tell me about bosses",
         "How does malaise work?",
-        "What are boss stem cells?",
     ],
-    theme="soft",
-    retry_btn=None,
-    undo_btn=None,
-    clear_btn="Clear Chat"
+    theme="soft"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()