Update app.py

app.py

@@ -1,10 +1,10 @@
 import faiss
 import numpy as np
-from fastapi import FastAPI
+from fastapi import FastAPI, Query, HTTPException
 from fastapi.responses import JSONResponse
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
-from typing import List
+from typing import List, Dict, Tuple
 
 app = FastAPI()
 
@@ -29,65 +29,76 @@ print("Loading sentence transformer model (all-MiniLM-L6-v2)...")
 model = SentenceTransformer("all-MiniLM-L6-v2")
 print("Model loaded successfully.")
 
-def load_and_index_dataset(name: str, include_readme: bool = False):
-    print(f"Loading dataset '{name}'...")
-    dataset = load_dataset(name)["train"]
-    
-    repo_texts = [
-        " ".join(str(x.get(field, "")) for field in FIELDS) +
-        (" " + x.get("readme_content", "") if include_readme else "") +
-        " " + " ".join(x.get("topics", []))
-        for x in dataset
-    ]
-    
-    if not include_readme:
-        dataset = [{k: v for k, v in item.items() if k != "readme_content"} for item in dataset]
-        
-    print(f"Creating embeddings for {len(repo_texts)} documents in '{name}'...")
-    repo_embeddings = model.encode(repo_texts, show_progress_bar=True)
-    
-    print(f"Building FAISS index for '{name}'...")
-    embedding_dim = repo_embeddings.shape[1]
-    
-    index = faiss.IndexFlatL2(embedding_dim)
-    index.add(np.array(repo_embeddings, dtype=np.float32))
-    
-    print(f"'{name}' dataset indexed with {index.ntotal} vectors.")
-    return index, list(dataset)
-
-indices = {}
+def load_and_index_dataset(name: str, include_readme: bool = False) -> Tuple[faiss.IndexFlatL2, List[Dict]]:
+    try:
+        print(f"Loading dataset '{name}'...")
+        dataset = load_dataset(name)["train"]
+
+        repo_texts = [
+            " ".join(str(x.get(field, "")) for field in FIELDS) +
+            (" " + x.get("readme_content", "") if include_readme else "") +
+            " " + " ".join(x.get("topics", []))
+            for x in dataset
+        ]
+
+        if not include_readme:
+            dataset = [{k: v for k, v in item.items() if k != "readme_content"} for item in dataset]
+
+        print(f"Creating embeddings for {len(repo_texts)} documents in '{name}'...")
+        repo_embeddings = model.encode(repo_texts, show_progress_bar=True)
+
+        embedding_dim = repo_embeddings.shape[1]
+        index = faiss.IndexFlatL2(embedding_dim)
+        index.add(np.array(repo_embeddings, dtype=np.float32))
+
+        print(f"'{name}' dataset indexed with {index.ntotal} vectors.")
+        return index, list(dataset)
+    except Exception as e:
+        print(f"Error loading dataset '{name}': {e}")
+        raise RuntimeError(f"Dataset loading/indexing failed: {name}")
+
+indices: Dict[str, Tuple[faiss.IndexFlatL2, List[Dict]]] = {}
 
 for key, readme_flag in {"packages": True, "programs": True}.items():
-    index, data = load_and_index_dataset(f"zigistry/{key}", include_readme=readme_flag)
-    indices[key] = (index, data)
-
-def perform_search(query: str, dataset_key: str, k: int):
-    index, dataset = indices[dataset_key]
-    
-    query_embedding = model.encode([query])
-    query_embedding = np.array(query_embedding, dtype=np.float32)
-
-    distances, idxs = index.search(query_embedding, k)
-
-    results = []
-    for dist, idx in zip(distances[0], idxs[0]):
-        if idx == -1:
-            continue
-            
-        item = dataset[int(idx)].copy()
-        item['relevance_score'] = 1.0 - (dist / 2.0)
-        results.append(item)
-        
-    return results
+    try:
+        index, data = load_and_index_dataset(f"zigistry/{key}", include_readme=readme_flag)
+        indices[key] = (index, data)
+    except Exception as e:
+        print(f"Failed to prepare index for {key}: {e}")
+        indices[key] = (None, [])
+
+def perform_search(query: str, dataset_key: str, k: int) -> List[Dict]:
+    index, dataset = indices.get(dataset_key, (None, []))
+    if not index:
+        raise HTTPException(status_code=500, detail=f"Index not available for {dataset_key}")
+
+    try:
+        query_embedding = model.encode([query])
+        distances, idxs = index.search(np.array(query_embedding, dtype=np.float32), k)
+
+        results = []
+        for dist, idx in zip(distances[0], idxs[0]):
+            if idx == -1:
+                continue
+            item = dataset[int(idx)].copy()
+            item["relevance_score"] = float(1.0 - dist / 2.0)
+            results.append(item)
+
+        return results
+    except Exception as e:
+        print(f"Error during search: {e}")
+        raise HTTPException(status_code=500, detail="Search failed")
 
 @app.get("/searchPackages/")
-def search_packages(q: str, k: int = 10):
-    results = perform_search(query=q, dataset_key="packages", k=k)
-    headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"}
-    return JSONResponse(content=results, headers=headers)
+def search_packages(q: str = Query(...), k: int = Query(10)) -> JSONResponse:
+    if not q:
+        raise HTTPException(status_code=400, detail="Query parameter 'q' is required.")
+    results = perform_search(q, "packages", k)
+    return JSONResponse(content=results, headers={"Access-Control-Allow-Origin": "*"})
 
 @app.get("/searchPrograms/")
-def search_programs(q: str, k: int = 10):
-    results = perform_search(query=q, dataset_key="programs", k=k)
-    headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"}
-    return JSONResponse(content=results, headers=headers)
+def search_programs(q: str = Query(...), k: int = Query(10)) -> JSONResponse:
+    if not q:
+        raise HTTPException(status_code=400, detail="Query parameter 'q' is required.")
+    results = perform_search(q, "programs", k)
+    return JSONResponse(content=results, headers={"Access-Control-Allow-Origin": "*"})