195 lines
6.5 KiB
Python
Executable File
195 lines
6.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Query the local homelab retrieval index and optionally ask Ollama."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import os
|
|
import pathlib
|
|
import re
|
|
import sys
|
|
import urllib.error
|
|
import urllib.request
|
|
from collections import Counter
|
|
|
|
|
|
DEFAULT_INDEX_DIR = pathlib.Path(os.environ.get("LAB_AI_KNOWLEDGE_INDEX_DIR", "/data/homelab-ai/index"))
|
|
TOKEN_RE = re.compile(r"[A-Za-z0-9_./:-]{2,}")
|
|
|
|
|
|
def tokenise(text: str) -> list[str]:
|
|
return [token.lower() for token in TOKEN_RE.findall(text)]
|
|
|
|
|
|
def load_index(index_dir: pathlib.Path) -> dict[str, object]:
|
|
index_path = index_dir / "index.json"
|
|
if not index_path.is_file():
|
|
raise FileNotFoundError(f"missing index: {index_path}")
|
|
return json.loads(index_path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def score_chunks(index: dict[str, object], query: str, limit: int) -> list[tuple[float, dict[str, object]]]:
|
|
query_terms = tokenise(query)
|
|
if not query_terms:
|
|
return []
|
|
|
|
query_counts = Counter(query_terms)
|
|
idf = index.get("idf", {})
|
|
average_length = float(index.get("average_length") or 1.0)
|
|
chunks = index.get("chunks", [])
|
|
results: list[tuple[float, dict[str, object]]] = []
|
|
k1 = 1.5
|
|
b = 0.75
|
|
|
|
if not isinstance(idf, dict) or not isinstance(chunks, list):
|
|
return []
|
|
|
|
for chunk in chunks:
|
|
if not isinstance(chunk, dict):
|
|
continue
|
|
tokens = chunk.get("tokens", [])
|
|
if not isinstance(tokens, list) or not tokens:
|
|
continue
|
|
term_counts = Counter(str(token) for token in tokens)
|
|
doc_len = len(tokens)
|
|
score = 0.0
|
|
for term, query_count in query_counts.items():
|
|
frequency = term_counts.get(term, 0)
|
|
if frequency == 0:
|
|
continue
|
|
term_idf = float(idf.get(term, 0.0))
|
|
denominator = frequency + k1 * (1 - b + b * doc_len / average_length)
|
|
score += query_count * term_idf * (frequency * (k1 + 1)) / denominator
|
|
if score > 0:
|
|
results.append((score, chunk))
|
|
|
|
return sorted(results, key=lambda item: item[0], reverse=True)[:limit]
|
|
|
|
|
|
def render_context(results: list[tuple[float, dict[str, object]]], max_chars: int) -> str:
|
|
sections: list[str] = []
|
|
remaining = max_chars
|
|
|
|
for score, chunk in results:
|
|
header = f"[{chunk.get('path')}:{chunk.get('start_line')}] score={score:.3f}"
|
|
text = str(chunk.get("text", "")).strip()
|
|
section = f"{header}\n{text}"
|
|
if len(section) > remaining:
|
|
section = section[: max(0, remaining - 20)].rstrip() + "\n[truncated]"
|
|
sections.append(section)
|
|
remaining -= len(section)
|
|
if remaining <= 0:
|
|
break
|
|
|
|
return "\n\n---\n\n".join(sections)
|
|
|
|
|
|
def render_citations(results: list[tuple[float, dict[str, object]]], max_chars: int) -> str:
|
|
sections = ["Jeannie Memory With Provenance", "=============================", ""]
|
|
remaining = max_chars
|
|
for index, (score, chunk) in enumerate(results, start=1):
|
|
path = chunk.get("path")
|
|
start_line = chunk.get("start_line")
|
|
end_line = chunk.get("end_line")
|
|
text = str(chunk.get("text", "")).strip()
|
|
excerpt = text[:900].rstrip()
|
|
section = [
|
|
f"[{index}] {path}:{start_line}-{end_line} score={score:.3f}",
|
|
excerpt,
|
|
]
|
|
rendered = "\n".join(section)
|
|
if len(rendered) > remaining:
|
|
rendered = rendered[: max(0, remaining - 20)].rstrip() + "\n[truncated]"
|
|
sections.append(rendered)
|
|
sections.append("")
|
|
remaining -= len(rendered)
|
|
if remaining <= 0:
|
|
break
|
|
sections.append("Rule: answers must cite these source paths or say the context is insufficient.")
|
|
return "\n".join(sections)
|
|
|
|
|
|
def ask_ollama(question: str, context: str, endpoint: str, model: str, timeout: int) -> str:
|
|
prompt = f"""You are helping operate a personal homelab.
|
|
Use only the context below. If the context is insufficient, say what to inspect next.
|
|
Be concise, factual, and command-oriented.
|
|
|
|
Question:
|
|
{question}
|
|
|
|
Context:
|
|
{context}
|
|
"""
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
"options": {
|
|
"temperature": 0.1,
|
|
"num_predict": 360,
|
|
},
|
|
}
|
|
request = urllib.request.Request(
|
|
f"{endpoint.rstrip('/')}/api/generate",
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
result = json.loads(response.read().decode("utf-8"))
|
|
return str(result.get("response", "")).strip()
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("query", nargs="+", help="Question or search text")
|
|
parser.add_argument("--index-dir", type=pathlib.Path, default=DEFAULT_INDEX_DIR)
|
|
parser.add_argument("--limit", type=int, default=5)
|
|
parser.add_argument("--max-chars", type=int, default=9000)
|
|
parser.add_argument("--context-only", action="store_true")
|
|
parser.add_argument("--citations", action="store_true")
|
|
parser.add_argument("--ask", action="store_true")
|
|
parser.add_argument("--ollama-url", default=os.environ.get("LAB_AI_GATEWAY_URL", "http://127.0.0.1:11434"))
|
|
parser.add_argument("--model", default=os.environ.get("LAB_AI_GATEWAY_MODEL", "qwen2.5:0.5b"))
|
|
parser.add_argument("--timeout", type=int, default=int(os.environ.get("LAB_AI_GATEWAY_TIMEOUT_SECONDS", "20")))
|
|
args = parser.parse_args()
|
|
|
|
query = " ".join(args.query)
|
|
try:
|
|
index = load_index(args.index_dir)
|
|
except FileNotFoundError as exc:
|
|
print(str(exc), file=sys.stderr)
|
|
return 1
|
|
|
|
results = score_chunks(index, query, args.limit)
|
|
if not results:
|
|
print("no matching homelab knowledge chunks found", file=sys.stderr)
|
|
return 1
|
|
|
|
context = render_context(results, args.max_chars)
|
|
if args.citations:
|
|
print(render_citations(results, args.max_chars))
|
|
return 0
|
|
|
|
if args.context_only:
|
|
print(context)
|
|
return 0
|
|
|
|
if args.ask:
|
|
try:
|
|
answer = ask_ollama(query, context, args.ollama_url, args.model, args.timeout)
|
|
except (OSError, TimeoutError, urllib.error.URLError, json.JSONDecodeError) as exc:
|
|
print(f"ollama request failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
print(answer)
|
|
return 0
|
|
|
|
print(context)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|