55 lines
2.0 KiB
Python
55 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
"""RAG test - embed query, search, answer via Qwen."""
|
|
import json, os, chromadb, urllib.request
|
|
|
|
def 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=30) as resp:
|
|
body = json.loads(resp.read())
|
|
if isinstance(body, list):
|
|
return [item["embedding"][0] for item in body]
|
|
return body.get("embedding", [])
|
|
|
|
os.chdir(os.path.expanduser("~/nubes"))
|
|
client = chromadb.PersistentClient(path=os.path.expanduser("~/nubes/chroma_db"))
|
|
collection = client.get_collection("elsa_docs")
|
|
|
|
# 1. Embed query
|
|
query = "What is the tightening torque for fuel filter on VW Phaeton 3.2 VR6 2004?"
|
|
print(f"Query: {query}")
|
|
q_emb = embed([query])[0]
|
|
|
|
# 2. Search
|
|
results = collection.query(query_embeddings=[q_emb], n_results=3)
|
|
print(f"Matches: {len(results['ids'][0])}")
|
|
for i in range(len(results['ids'][0])):
|
|
meta = results['metadatas'][0][i] or {}
|
|
dist = results['distances'][0][i]
|
|
print(f"\n[{i}] dist={dist:.3f}")
|
|
print(f" source: {meta.get('source','')}")
|
|
print(f" title: {meta.get('title','')[:120]}")
|
|
print(f" text: {results['documents'][0][i][:300]}...")
|
|
|
|
# 3. Build prompt for Qwen
|
|
context = "\n\n".join([f"[{m.get('title','')[:80] if m else ''}]\n{d[:1500]}"
|
|
for d, m in zip(results['documents'][0], results['metadatas'][0])])
|
|
|
|
prompt = f"""You are a VAG automotive diagnostician. Answer using ONLY the documentation below.
|
|
DOCUMENTATION:
|
|
{context}
|
|
|
|
Question: {query}
|
|
Answer:"""
|
|
|
|
data = json.dumps({"prompt": prompt, "temperature": 0.1, "n_predict": 300,
|
|
"stop": ["\n\n", "Answer:", "\nThe answer", "\nDocumentation"]}).encode()
|
|
req = urllib.request.Request("http://localhost:8080/completion", data=data,
|
|
headers={"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(req, timeout=120) as resp:
|
|
answer = json.loads(resp.read()).get("content", "")
|
|
|
|
print(f"\n\n=== LLM Answer ===\n{answer}")
|
|
|