Update app.py

app.py

@@ -4,105 +4,78 @@ from fastapi import FastAPI, Query
 from fastapi.responses import JSONResponse
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
+from typing import List, Dict
 
 app = FastAPI()
 
-FIELDS = (
-    "full_name",
-    "description",
-    "default_branch",
-    "open_issues",
-    "stargazers_count",
-    "forks_count",
-    "watchers_count",
-    "license",
-    "size",
-    "fork",
-    "updated_at",
-    "has_build_zig",
-    "has_build_zig_zon",
-    "created_at",
-)
+# Constants
+EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
+SEARCH_THRESHOLD_RATIO = 0.3
+FIELDS_TO_INCLUDE = [
+    "full_name", "description", "default_branch", "open_issues",
+    "stargazers_count", "forks_count", "watchers_count", "license",
+    "size", "fork", "updated_at", "has_build_zig",
+    "has_build_zig_zon", "created_at"
+]
 
-model = SentenceTransformer("all-MiniLM-L6-v2")
+# Load embedding model
+model = SentenceTransformer(EMBEDDING_MODEL_NAME)
 
-def load_dataset_with_fields(name, include_readme=False):
-    dataset = load_dataset(name)["train"]
-    repo_texts = [
-        " ".join(str(x.get(field, "")) for field in FIELDS) + 
-        (" " + x.get("readme_content", "")) * include_readme + 
-        " " + " ".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]
-    return dataset, repo_texts
+# Helper functions
+def prepare_text(entry: Dict, include_readme: bool = True) -> str:
+    parts = [str(entry.get(field, "")) for field in FIELDS_TO_INCLUDE]
+    if include_readme:
+        parts.append(entry.get("readme_content", ""))
+    parts.extend(entry.get("topics", []))
+    return " ".join(parts)
 
-datasets = {
-    "packages": load_dataset_with_fields("zigistry/packages", include_readme=True),
-    "programs": load_dataset_with_fields("zigistry/programs", include_readme=True),
-}
+def load_and_encode_dataset(name: str, include_readme: bool = True):
+    raw_dataset = load_dataset(name)["train"]
+    texts = [prepare_text(item, include_readme) for item in raw_dataset]
+    embeddings = model.encode(texts)
+    return raw_dataset, np.array(embeddings)
 
-indices = {}
-for key, (dataset, repo_texts) in datasets.items():
-    repo_embeddings = model.encode(repo_texts)
-    index = faiss.IndexFlatL2(repo_embeddings.shape[1])
-    index.add(np.array(repo_embeddings))
-    indices[key] = (index, dataset)
+def build_faiss_index(embeddings: np.ndarray) -> faiss.IndexFlatL2:
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(embeddings)
+    return index
 
-scroll_data = {
-    "infiniteScrollPackages": load_dataset_with_fields("zigistry/packages", include_readme=False)[0],
-    "infiniteScrollPrograms": load_dataset_with_fields("zigistry/programs", include_readme=False)[0],
-}
+def search_index(index: faiss.IndexFlatL2, query: str, embeddings: np.ndarray, dataset: List[Dict]) -> List[Dict]:
+    query_vector = model.encode([query])
+    distances, indices = index.search(np.array(query_vector), len(dataset))
+    return filter_by_distance(distances[0], indices[0], dataset)
 
-def filter_results_by_distance(distances, idxs, dataset, max_results=50, threshold=0.6):
-    """
-    Only return results that are likely relevant (distance-based filtering).
-    Lower distance = more similar.
-    Threshold is a fraction of the *minimum* distance found.
-    """
+def filter_by_distance(distances: np.ndarray, indices: np.ndarray, dataset: List[Dict], ratio: float = SEARCH_THRESHOLD_RATIO) -> List[Dict]:
     if len(distances) == 0:
         return []
-    min_dist = np.min(distances)
-    cutoff = min_dist + ((max(distances) - min_dist) * threshold)
-    filtered = [
-        dataset[int(i)]
-        for d, i in zip(distances, idxs)
-        if d <= cutoff
-    ]
-    return filtered[:max_results]
+    min_d, max_d = np.min(distances), np.max(distances)
+    threshold = min_d + (max_d - min_d) * ratio
+    return [dataset[i] for d, i in zip(distances, indices) if d <= threshold]
 
-@app.get("/infiniteScrollPackages/")
-def infinite_scroll_packages(q: int = Query(0, ge=0)):
-    start = q * 10
-    content = scroll_data["infiniteScrollPackages"][start : start + 10]
-    headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"}
-    return JSONResponse(content=content, headers=headers)
+# Load datasets and create indices
+data_configs = {
+    "packages": "zigistry/packages",
+    "programs": "zigistry/programs"
+}
 
-@app.get("/infiniteScrollPrograms/")
-def infinite_scroll_programs(q: int = Query(0, ge=0)):
-    start = q * 10
-    content = scroll_data["infiniteScrollPrograms"][start : start + 10]
-    headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"}
-    return JSONResponse(content=content, headers=headers)
+data_store = {}
 
-@app.get("/searchPackages/")
-def search_packages(q: str):
-    key = "packages"
-    index, dataset = indices[key]
-    query_embedding = model.encode([q])
-    distances, idxs = index.search(np.array(query_embedding), len(dataset))
-    # Only keep results that are likely relevant
-    results = filter_results_by_distance(distances[0], idxs[0], dataset)
-    headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"}
-    return JSONResponse(content=results, headers=headers)
+for key, dataset_name in data_configs.items():
+    dataset, embeddings = load_and_encode_dataset(dataset_name, include_readme=True)
+    index = build_faiss_index(embeddings)
+    data_store[key] = {
+        "dataset": dataset,
+        "index": index,
+        "embeddings": embeddings
+    }
+
+# FastAPI endpoints
+@app.get("/search/{category}/")
+def search(category: str, q: str = Query(...)):
+    if category not in data_store:
+        return JSONResponse(status_code=404, content={"error": "Invalid category"})
 
-@app.get("/searchPrograms/")
-def search_programs(q: str):
-    key = "programs"
-    index, dataset = indices[key]
-    query_embedding = model.encode([q])
-    distances, idxs = index.search(np.array(query_embedding), len(dataset))
-    results = filter_results_by_distance(distances[0], idxs[0], dataset)
+    store = data_store[category]
+    results = search_index(store["index"], q, store["embeddings"], store["dataset"])
     headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"}
     return JSONResponse(content=results, headers=headers)