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

scripts/rag_query.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""RAG Query — search ChromaDB + answer via llama.cpp."""
+import json, os, sys, urllib.request, time
+import chromadb
+
+LLAMA_URL = "http://localhost:8080/completion"
+
+def llama_embed(texts):
+    data = json.dumps({"content": texts}).encode()
+    req = urllib.request.Request("http://localhost:8081/embedding", data=data,
+        headers={"Content-Type": "application/json"})
+    with urllib.request.urlopen(req, timeout=60) as resp:
+        body = json.loads(resp.read())
+    if isinstance(body, list):
+        return [item["embedding"][0] for item in body]
+    return body.get("embedding", [])
+
+
+class LlamaEmbedFn:
+    def __call__(self, input):
+        all_embs = []
+        for i in range(0, len(input), 32):
+            batch = input[i:i+32]
+            embs = llama_embed(batch)
+            all_embs.extend(embs if isinstance(embs[0], list) else [embs])
+        return all_embs
+
+
+def search(collection, query, n=5):
+    results = collection.query(query_texts=[query], n_results=n)
+    docs = []
+    for i in range(len(results["ids"][0])):
+        docs.append({
+            "id": results["ids"][0][i],
+            "doc": results["documents"][0][i],
+            "meta": results["metadatas"][0][i],
+            "distance": results["distances"][0][i],
+        })
+    return docs
+
+
+def ask_llama(prompt, max_tokens=256):
+    data = json.dumps({
+        "prompt": prompt, "temperature": 0.1,
+        "n_predict": max_tokens,
+        "stop": ["\n\n", "Documentation:", "Question:"],
+    }).encode()
+    req = urllib.request.Request(LLAMA_URL, data=data,
+        headers={"Content-Type": "application/json"})
+    with urllib.request.urlopen(req, timeout=120) as resp:
+        return json.loads(resp.read()).get("content", "")
+
+
+def main():
+    db_dir = os.path.expanduser("~/nubes/chroma_db")
+    collection_name = "elsa_docs"
+    
+    ef = LlamaEmbedFn()
+    
+    client = chromadb.PersistentClient(path=db_dir)
+    collection = client.get_collection(collection_name, embedding_function=ef)
+    
+    print(f"Collection '{collection_name}': {collection.count()} documents")
+    
+    while True:
+        try:
+            query = input("\n>>> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            break
+        
+        if not query:
+            continue
+        if query.lower() in ("exit", "quit", "q"):
+            break
+        
+        # 1. Поиск
+        results = search(collection, query, n_results=5)
+        
+        # 2. Формируем контекст
+        context_parts = []
+        for r in results:
+            title = r["meta"].get("title", "")[:120]
+            ctx = f"[{title}] {r['doc'][:1000]}"
+            context_parts.append(ctx)
+        
+        context = "\n\n---\n\n".join(context_parts)
+        
+        # 3. Промпт
+        prompt = f"""You are an expert VAG (VW/Audi/Skoda/SEAT) automotive diagnostician. 
+Answer the user's question using ONLY the provided repair documentation below.
+If the answer is not in the documentation, say "Not found in documentation."
+Answer in Russian.
+
+DOCUMENTATION:
+{context}
+
+Question: {query}
+Answer:"""
+        
+        print(f"\n[Found {len(results)} matches, sending to LLM...]")
+        answer = ask_llama(prompt)
+        print(f"\n{answer}")
+
+
+if __name__ == "__main__":
+    main()