ElsaWin RAG pipeline: full parsing chain + ChromaDB ingest + docs

scripts/rag_ingest.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+"""Ingest JSONL into ChromaDB — uses llama.cpp /embedding endpoint."""
+import json, os, sys, glob, urllib.request, time
+import chromadb
+
+LLAMA_URL = "http://localhost:8081/embedding"
+
+def llama_embed(texts):
+    """Get embeddings from nomic-embed server with retry."""
+    data = json.dumps({"content": texts}).encode()
+    req = urllib.request.Request(LLAMA_URL, data=data,
+        headers={"Content-Type": "application/json"})
+    for attempt in range(3):
+        try:
+            with urllib.request.urlopen(req, timeout=120) as resp:
+                body = json.loads(resp.read())
+            if isinstance(body, list):
+                return [item["embedding"][0] for item in body]
+            return body.get("embedding", [])
+        except Exception as e:
+            if attempt < 2:
+                print(f"[retry {attempt+1}] {e}")
+                time.sleep(5)
+            else:
+                raise
+
+
+class LlamaEmbedFn:
+    """Custom embedding function using llama.cpp."""
+    def __call__(self, input):
+        all_embs = []
+        for i in range(0, len(input), 32):
+            batch = input[i:i+32]
+            embs = llama_embed(batch)
+            if isinstance(embs[0], float):
+                all_embs.append(embs)
+            else:
+                all_embs.extend(embs)
+        return all_embs
+
+
+def main():
+    jsonl_dir = sys.argv[1] if len(sys.argv) > 1 else os.path.expanduser("~/nubes/data/elsa_jsonl_wi_en")
+    db_dir = sys.argv[2] if len(sys.argv) > 2 else os.path.expanduser("~/nubes/chroma_db")
+    collection_name = sys.argv[3] if len(sys.argv) > 3 else "elsa_docs"
+    
+    ef = LlamaEmbedFn()
+    
+    client = chromadb.PersistentClient(path=db_dir)
+    
+    try:
+        client.delete_collection(collection_name)
+    except:
+        pass
+    
+    collection = client.create_collection(
+        name=collection_name,
+        embedding_function=ef,
+        metadata={"hnsw:space": "cosine"}
+    )
+    
+    jsonl_files = sorted(glob.glob(os.path.join(jsonl_dir, "*.jsonl")))
+    print(f"Found {len(jsonl_files)} JSONL files")
+    print(f"Using llama.cpp at {LLAMA_URL}")
+    
+    total = 0
+    for jf in jsonl_files:
+        ids, docs, metas = [], [], []
+        
+        with open(jf, "r", encoding="utf-8") as f:
+            for line in f:
+                try:
+                    d = json.loads(line)
+                    title = d.get("title", "").strip()
+                    text = d.get("text", "").strip()
+                    ft = d.get("full_title", "").strip()
+                    
+                    if ft and text:
+                        doc = f"{ft}\n{text}"
+                    elif title and text:
+                        doc = f"{title}\n{text}"
+                    else:
+                        doc = text or title or ft
+                    
+                    if not doc or len(doc) < 20:
+                        continue
+                    
+                    # Split long docs into chunks of ~1000 chars
+                    chunks = [doc[i:i+1000] for i in range(0, len(doc), 1000)]
+                    for chunk in chunks:
+                        if len(chunk) < 20:
+                            continue
+                        ids.append(str(total))
+                        docs.append(chunk)
+                        metas.append({
+                            "source": d.get("source", ""),
+                            "file": d.get("file", ""),
+                            "title": title or ft or doc[:100],
+                        })
+                        total += 1
+                except:
+                    continue
+        
+        if ids:
+            # Send in smaller batches of 500
+            for i in range(0, len(ids), 500):
+                batch_ids = ids[i:i+500]
+                batch_docs = docs[i:i+500]
+                batch_metas = metas[i:i+500]
+                for attempt in range(3):
+                    try:
+                        collection.add(ids=batch_ids, documents=batch_docs, metadatas=batch_metas)
+                        break
+                    except Exception as e:
+                        if attempt < 2:
+                            print(f"  [retry {os.path.basename(jf)} batch {attempt+1}] {e}")
+                            time.sleep(10)
+                        else:
+                            print(f"  [SKIP batch] {e}")
+            print(f"  {os.path.basename(jf)}: +{len(ids)} (total {total})")
+    
+    print(f"\nDone. {total} documents → {db_dir}/{collection_name}")
+
+if __name__ == "__main__":
+    main()