108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
#!/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()
|