Update app.py

app.py

@@ -1,5 +1,6 @@
 
 try:
+    import os
     import gradio as gr
     import torch
     from sentence_transformers import SentenceTransformer
@@ -241,6 +242,111 @@ class EmbeddingGemmaPrompts:
 
 
 
+def search_knowledge_base(
+    query: str, 
+    num_results: int = 5, 
+    source_filter: Optional[str] = None,
+    task_type: str = "search"
+) -> Dict[str, Any]:
+    """
+    Search the RS Studies knowledge base using semantic similarity
+    
+    Args:
+        query: The search query
+        num_results: Number of results to return
+        source_filter: Optional source folder filter
+        task_type: Type of task for query formatting
+    
+    Returns:
+        Dictionary with search results and metadata
+    """
+    if not ensure_initialized():
+        return {"error": "Server not properly initialized", "results": []}
+    
+    try:
+        # Create query embedding with task-specific formatting using EmbeddingGemmaPrompts
+        query_formatted = EmbeddingGemmaPrompts.encode_query(query, task_type)
+        query_embedding = model.encode([query_formatted], device=device)
+        
+        # Prepare search parameters
+        search_params = {
+            "query_embeddings": query_embedding.tolist(),
+            "n_results": min(num_results, config.MAX_NUM_RESULTS),
+            "include": ["documents", "metadatas", "distances"]
+        }
+        
+        # Add source filter if specified
+        if source_filter and source_filter in config.VALID_SOURCES:
+            search_params["where"] = {"source_folder": {"$eq": source_filter}}
+        
+        # Perform search
+        results = collection.query(**search_params)
+        
+        # Format results
+        formatted_results = []
+        if results["documents"] and len(results["documents"]) > 0:
+            for i in range(len(results["documents"][0])):
+                result = {
+                    "rank": i + 1,
+                    "content": results["documents"][0][i],
+                    "source_folder": results["metadatas"][0][i].get("source_folder", "unknown"),
+                    "chunk_file": results["metadatas"][0][i].get("chunk_file", "unknown"),
+                    "chunk_number": results["metadatas"][0][i].get("chunk_number", "unknown"),
+                    "similarity_score": float(1 - results["distances"][0][i]),
+                    "distance": float(results["distances"][0][i]),
+                    "chunk_length": results["metadatas"][0][i].get("chunk_length", 0),
+                    "metadata": results["metadatas"][0][i]
+                }
+                formatted_results.append(result)
+        
+        return {
+            "query": query,
+            "task_type": task_type,
+            "num_results": len(formatted_results),
+            "source_filter": source_filter,
+            "results": formatted_results,
+            "success": True
+        }
+        
+    except Exception as e:
+        return {"error": f"Search failed: {str(e)}", "results": [], "success": False}
+
+def get_available_sources() -> Dict[str, Any]:
+    """Get list of available source folders in the knowledge base"""
+    if not ensure_initialized():
+        return {"error": "Server not properly initialized", "sources": []}
+    
+    try:
+        # Get all metadata to find unique source folders
+        all_results = collection.get(include=["metadatas"])
+        sources = set()
+        
+        for metadata in all_results["metadatas"]:
+            source = metadata.get("source_folder")
+            if source:
+                sources.add(source)
+        
+        # Get statistics for each source
+        source_stats = {}
+        for source in sources:
+            source_results = collection.get(
+                where={"source_folder": {"$eq": source}},
+                include=["metadatas"]
+            )
+            source_stats[source] = len(source_results["metadatas"])
+        
+        return {
+            "sources": sorted(list(sources)),
+            "source_stats": source_stats,
+            "total_sources": len(sources),
+            "total_chunks": collection.count(),
+            "success": True
+        }
+        
+    except Exception as e:
+        return {"error": f"Failed to get sources: {str(e)}", "sources": [], "success": False}
+
+
 def slice_list(lst: list, start: int, end: int) -> list:
     """
     A tool that slices a list given a start and end index.