Update app.py

app.py

@@ -1,78 +1,183 @@
-import gradio as gr
-import chromadb
-from sentence_transformers import SentenceTransformer
-import google.generativeai as genai
-import os
-
-# Get API key from Hugging Face secrets
-GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
-genai.configure(api_key=GOOGLE_API_KEY)
-
-# Load models (cached)
-embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
-chroma_client = chromadb.PersistentClient(path="./deadcells_db_free")
-collection = chroma_client.get_collection(name="deadcells_wiki")
-
-# Find Gemini model
-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
-    )
-    
-    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.
-
-Wiki Content:
-{context}
-
-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])
-    
-    return answer + source_list
-
-# 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.",
-    examples=[
-        "What achievements are there for beating bosses?",
-        "Tell me about the Hand of the King",
-        "How does malaise work?",
-        "What are boss stem cells?",
-    ],
-    theme="soft",
-    retry_btn=None,
-    undo_btn=None,
-    clear_btn="Clear Chat"
-)
-
-if __name__ == "__main__":
+import gradio as gr
+import chromadb
+from sentence_transformers import SentenceTransformer
+import google.generativeai as genai
+import os
+import glob
+
+# Get API key
+GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
+genai.configure(api_key=GOOGLE_API_KEY)
+
+# Load embedding model
+print("Loading embedding model...")
+embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+print("Model loaded!")
+
+# Create or load database
+print("Checking database...")
+chroma_client = chromadb.PersistentClient(path="./deadcells_db_free")
+
+try:
+    collection = chroma_client.get_collection(name="deadcells_wiki")
+    print(f"✅ Database loaded! {collection.count()} chunks available")
+    
+except Exception as e:
+    print(f"Database not found! Building from wiki files...")
+    
+    # Create collection
+    collection = chroma_client.create_collection(name="deadcells_wiki")
+    
+    # Load wiki files
+    wiki_files = glob.glob("wiki_content/*.txt")
+    
+    if not wiki_files:
+        raise Exception(f"❌ No wiki files found! Please upload wiki_content folder")
+    
+    print(f"Found {len(wiki_files)} wiki files. Starting vectorization...")
+    
+    chunk_id = 0
+    batch_size = 100
+    batch_texts = []
+    batch_embeddings = []
+    batch_metadatas = []
+    batch_ids = []
+    
+    for idx, filepath in enumerate(wiki_files):
+        filename = os.path.basename(filepath)
+        
+        try:
+            with open(filepath, 'r', encoding='utf-8') as f:
+                content = f.read()
+            
+            # Parse file
+            lines = content.split('\n')
+            url = lines[0].replace('URL: ', '') if lines and lines[0].startswith('URL:') else ''
+            text = '\n'.join(lines[2:]) if len(lines) > 2 else content
+            
+            # Skip if empty
+            if not text or len(text.strip()) < 100:
+                continue
+            
+            # 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
+                
+                # Create embedding
+                embedding = embedding_model.encode(chunk).tolist()
+                
+                # Add to batch
+                batch_texts.append(chunk)
+                batch_embeddings.append(embedding)
+                batch_metadatas.append({
+                    "source": filename,
+                    "url": url,
+                    "chunk_index": i
+                })
+                batch_ids.append(f"chunk-{chunk_id}")
+                chunk_id += 1
+                
+                # Upload batch when full
+                if len(batch_texts) >= batch_size:
+                    collection.add(
+                        documents=batch_texts,
+                        embeddings=batch_embeddings,
+                        metadatas=batch_metadatas,
+                        ids=batch_ids
+                    )
+                    batch_texts = []
+                    batch_embeddings = []
+                    batch_metadatas = []
+                    batch_ids = []
+        
+        except Exception as file_error:
+            print(f"⚠️ Error processing {filename}: {file_error}")
+            continue
+        
+        # Progress update
+        if (idx + 1) % 50 == 0:
+            print(f"Processed {idx + 1}/{len(wiki_files)} files...")
+    
+    # Upload remaining batch
+    if batch_texts:
+        collection.add(
+            documents=batch_texts,
+            embeddings=batch_embeddings,
+            metadatas=batch_metadatas,
+            ids=batch_ids
+        )
+    
+    print(f"✅ Database created! {chunk_id} chunks indexed from {len(wiki_files)} files")
+
+# Gemini model
+available_model = 'gemini-1.5-flash-latest'
+
+def get_embedding(text):
+    return embedding_model.encode(text).tolist()
+
+def chat(message, history):
+    """Handle chat messages"""
+    
+    try:
+        # Search wiki
+        question_embedding = get_embedding(message)
+        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.
+
+Wiki Content:
+{context}
+
+Question: {message}
+
+Answer based ONLY on the wiki content above:"""
+        
+        response = model.generate_content(prompt)
+        answer = response.text
+        
+        # Add sources
+        source_list = "\n\n📚 **Sources:** " + ", ".join(set([s['source'] for s in sources]))
+        
+        return answer + source_list
+        
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+# Create Gradio interface
+demo = gr.ChatInterface(
+    fn=chat,
+    title="🎮 Dead Cells Wiki Bot",
+    description="Ask me anything about Dead Cells! Powered by 557 wiki pages.",
+    examples=[
+        "What achievements are there for beating bosses?",
+        "Tell me about the Hand of the King",
+        "How does malaise work?",
+        "What are boss stem cells?",
+        "List some brutality weapons"
+    ],
+    theme="soft",
+    retry_btn=None,
+    undo_btn=None,
+    clear_btn="Clear Chat"
+)
+
+if __name__ == "__main__":
     demo.launch()