Add document search functionality

- Create minimal vector search implementation
- Add embedding-based document retrieval
- Integrate search results into chat context
- Use simple Blocks interface instead of ChatInterface

app.py

@@ -1,9 +1,38 @@
 import gradio as gr
 import os
 from openai import OpenAI
+from simple_vector_search import SimpleVectorSearch
+
+# Initialize vector search
+vector_search = SimpleVectorSearch()
+
+def get_relevant_context(query: str, client: OpenAI) -> str:
+    """Get relevant context from documents"""
+    try:
+        # Get query embedding
+        response = client.embeddings.create(
+            model="text-embedding-ada-002",
+            input=query
+        )
+        query_embedding = response.data[0].embedding
+        
+        # Search for relevant documents
+        results = vector_search.search(query_embedding, k=3)
+        
+        # Format context
+        if results:
+            context = "Relevant information from documentation:\n\n"
+            for i, result in enumerate(results, 1):
+                context += f"{i}. {result['content'][:500]}...\n\n"
+            return context
+        else:
+            return ""
+    except Exception as e:
+        print(f"Error in vector search: {e}")
+        return ""
 
 def chat_function(message, history):
-    """Chat function with OpenAI integration"""
+    """Chat function with OpenAI integration and document search"""
     # Check if API key is set
     api_key = os.environ.get("OPENAI_API_KEY", "")
     
@@ -14,8 +43,16 @@ def chat_function(message, history):
             # Initialize OpenAI client
             client = OpenAI(api_key=api_key)
             
+            # Get relevant context
+            context = get_relevant_context(message, client)
+            
+            # Build system message
+            system_content = "You are a helpful AI assistant with access to technical documentation."
+            if context:
+                system_content += f"\n\n{context}"
+            
             # Convert history to OpenAI format
-            messages = [{"role": "system", "content": "You are a helpful AI assistant."}]
+            messages = [{"role": "system", "content": system_content}]
             
             for user_msg, assistant_msg in history:
                 messages.append({"role": "user", "content": user_msg})
@@ -40,10 +77,27 @@ def chat_function(message, history):
     return history
 
 # Create chat interface
-demo = gr.ChatInterface(
-    fn=chat_function,
-    title="AI Assistant Multi-Agent System",
-    description="Chat with OpenAI-powered AI assistant"
-)
+with gr.Blocks(title="AI Assistant") as demo:
+    gr.Markdown("# 🤖 AI Assistant Multi-Agent System")
+    gr.Markdown("Chat with an AI assistant that has access to technical documentation.")
+    
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(
+        label="Your Message",
+        placeholder="Ask about Harmony or Chorus documentation...",
+        lines=2
+    )
+    
+    with gr.Row():
+        submit = gr.Button("Send", variant="primary")
+        clear = gr.Button("Clear")
+    
+    # Event handlers
+    def respond(message, history):
+        return chat_function(message, history), ""
+    
+    msg.submit(respond, [msg, chatbot], [chatbot, msg])
+    submit.click(respond, [msg, chatbot], [chatbot, msg])
+    clear.click(lambda: ([], ""), outputs=[chatbot, msg])
 
 demo.launch()

requirements.txt

@@ -1,2 +1,4 @@
 gradio>=4.0.0
-openai>=1.0.0
+openai>=1.0.0
+numpy>=1.24.0
+tiktoken>=0.5.0

simple_vector_search.py

@@ -0,0 +1,65 @@
+"""
+Minimal vector search implementation for HuggingFace deployment
+"""
+import json
+import numpy as np
+from pathlib import Path
+from typing import List, Dict, Tuple
+
+class SimpleVectorSearch:
+    """Simple in-memory vector search"""
+    
+    def __init__(self, data_dir: str = "py/backend/data/embeddings"):
+        self.data_dir = Path(data_dir)
+        self.documents = []
+        self.embeddings = None
+        self._load_embeddings()
+    
+    def _load_embeddings(self):
+        """Load all embedding files"""
+        all_docs = []
+        all_embeddings = []
+        
+        # Load all JSON files
+        for json_file in self.data_dir.glob("*.json"):
+            try:
+                with open(json_file, 'r') as f:
+                    data = json.load(f)
+                    for item in data:
+                        all_docs.append({
+                            'content': item['content'],
+                            'metadata': item.get('metadata', {})
+                        })
+                        all_embeddings.append(item['embedding'])
+            except Exception as e:
+                print(f"Error loading {json_file}: {e}")
+        
+        self.documents = all_docs
+        self.embeddings = np.array(all_embeddings) if all_embeddings else None
+    
+    def search(self, query_embedding: List[float], k: int = 3) -> List[Dict]:
+        """Search for similar documents"""
+        if self.embeddings is None or len(self.embeddings) == 0:
+            return []
+        
+        # Convert query to numpy array
+        query_vec = np.array(query_embedding)
+        
+        # Compute cosine similarity
+        query_norm = query_vec / (np.linalg.norm(query_vec) + 1e-10)
+        embeddings_norm = self.embeddings / (np.linalg.norm(self.embeddings, axis=1, keepdims=True) + 1e-10)
+        similarities = np.dot(embeddings_norm, query_norm)
+        
+        # Get top k indices
+        top_indices = np.argsort(similarities)[-k:][::-1]
+        
+        # Return documents with scores
+        results = []
+        for idx in top_indices:
+            results.append({
+                'content': self.documents[idx]['content'],
+                'metadata': self.documents[idx]['metadata'],
+                'score': float(similarities[idx])
+            })
+        
+        return results