Create ingest/cluster.py

ingest/cluster.py

@@ -1,33 +1,20 @@
-from typing import List, Dict
-
-def build_embeddings(results: List[Dict]):
-    try:
-        from sentence_transformers import SentenceTransformer
-        model = SentenceTransformer("all-MiniLM-L6-v2")
-        texts = [r.get("snippet", "") for r in results]
-        return model.encode(texts, show_progress_bar=False)
-    except Exception:
-        return None
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
 
+model = SentenceTransformer("all-MiniLM-L6-v2")
 
-def cluster_embeddings(results: List[Dict], embeddings):
-    if embeddings is None or len(results) == 0:
-        return {}
-
-    try:
-        import faiss
-        import numpy as np
+def semantic_cluster(results):
+    texts = [r["title"] + " " + r["snippet"] for r in results]
+    if not texts:
+        return None
 
-        dim = embeddings.shape[1]
-        index = faiss.IndexFlatL2(dim)
-        index.add(embeddings)
+    embeddings = model.encode(texts)
+    dim = embeddings.shape[1]
+    index = faiss.IndexFlatL2(dim)
+    index.add(np.array(embeddings))
 
-        clusters = {}
-        for i, r in enumerate(results):
-            clusters.setdefault(r["source"], []).append({
-                "title": r["title"],
-                "index": i
-            })
-        return clusters
-    except Exception:
-        return {}
+    return {
+        "points": embeddings.tolist(),
+        "labels": list(range(len(results)))
+    }