756 lines
30 KiB
Python
Executable File
756 lines
30 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Context Enhancer — HYBRID search (semantic + BM25) on knowledge_base_hybrid
|
|
for prompt enrichment.
|
|
|
|
Runs as a synchronous function (fast, <1s) before each Hermes response.
|
|
If Qdrant is offline or embedding fails, returns "" (fail-open).
|
|
|
|
Usage:
|
|
python3 context_enhancer.py "your query here"
|
|
python3 context_enhancer.py --top-k 5 --threshold 0.50 "deploy docker"
|
|
python3 context_enhancer.py --hybrid-off "your query here" # forces dense-only
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import uuid
|
|
import hashlib
|
|
import sqlite3
|
|
import subprocess
|
|
import requests
|
|
import argparse
|
|
import re
|
|
import glob
|
|
from typing import List, Dict, Optional, Tuple
|
|
from pathlib import Path
|
|
import time
|
|
from datetime import datetime, timezone
|
|
|
|
# ─── Config ────────────────────────────────────────────────────────────────
|
|
OPENROUTER_KEY = os.environ.get("OPENROUTER_API_KEY")
|
|
if not OPENROUTER_KEY:
|
|
env_path = Path.home() / ".env"
|
|
if env_path.exists():
|
|
for line in env_path.read_text().splitlines():
|
|
if line.startswith("OPENROUTER_API_KEY="):
|
|
OPENROUTER_KEY = line.split("=", 1)[1].strip().strip('"')
|
|
break
|
|
|
|
QDRANT_URL = os.environ.get("QDRANT_URL", "http://localhost:6333")
|
|
COLLECTION = os.environ.get("QDRANT_COLLECTION", "knowledge_base")
|
|
|
|
EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "qwen/qwen3-embedding-8b")
|
|
TOP_K_DEFAULT = 3
|
|
SCORE_THRESHOLD_DEFAULT = 0.55
|
|
MAX_TEXT_LEN = 8000
|
|
REQUEST_TIMEOUT = 10
|
|
|
|
# FastEmbed BM25 config
|
|
FASTEMBED_VENV = os.environ.get("FASTEMBED_VENV", "")
|
|
|
|
# Default: use current Python if venv not configured
|
|
_FASTEMBED_PYTHON = FASTEMBED_VENV if FASTEMBED_VENV else sys.executable
|
|
_FASTEMBED_SITEPKGS = os.environ.get(
|
|
"FASTEMBED_SITEPKGS",
|
|
os.path.join(os.path.dirname(sys.executable), "../lib/python3.12/site-packages")
|
|
)
|
|
BM25_MODEL = "Qdrant/bm25"
|
|
|
|
# Lineage config
|
|
LINEAGE_DB = os.environ.get(
|
|
"STATE_DB_PATH",
|
|
os.path.expanduser("~/.hermes/state.db")
|
|
)
|
|
|
|
# Telemetry config
|
|
TELEMETRY_LOG = os.environ.get(
|
|
"TELEMETRY_LOG_PATH",
|
|
os.path.expanduser("~/.hermes/logs/query-telemetry.jsonl")
|
|
)
|
|
TELEMETRY_MAX_BYTES = 10 * 1024 * 1024 # 10MB rotation
|
|
|
|
# ─── Lineage Registration ───────────────────────────────────────────────────
|
|
|
|
def register_lineage(
|
|
session_id: str,
|
|
query: str,
|
|
retrieved_chunk_ids: List[str],
|
|
generation_context_hash: str,
|
|
generation_model: str = "unknown",
|
|
) -> Optional[str]:
|
|
"""
|
|
Register generation provenance in the lineage DB.
|
|
Fail-open: if it fails, log error and return None. Never breaks the critical path.
|
|
"""
|
|
lineage_id = str(uuid.uuid4())
|
|
try:
|
|
with sqlite3.connect(LINEAGE_DB) as conn:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO lineage (lineage_id, session_id, query, retrieved_chunk_ids,
|
|
generation_model, generation_context_hash, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, datetime('now'))
|
|
""",
|
|
(lineage_id, session_id, query,
|
|
json.dumps(retrieved_chunk_ids, ensure_ascii=False),
|
|
generation_model, generation_context_hash)
|
|
)
|
|
conn.commit()
|
|
return lineage_id
|
|
except Exception as e:
|
|
print(f"[LINEAGE-WARNING] Failed to register lineage: {e}", file=sys.stderr)
|
|
return None
|
|
|
|
|
|
# ─── Telemetry ─────────────────────────────────────────────────────────────
|
|
|
|
def _rotate_telemetry_if_needed() -> None:
|
|
"""Rotate telemetry file if >10MB. Rename to .1 and restart."""
|
|
try:
|
|
if os.path.exists(TELEMETRY_LOG) and os.path.getsize(TELEMETRY_LOG) > TELEMETRY_MAX_BYTES:
|
|
rotated = TELEMETRY_LOG + ".1"
|
|
if os.path.exists(rotated):
|
|
os.remove(rotated)
|
|
os.rename(TELEMETRY_LOG, rotated)
|
|
except Exception as e:
|
|
print(f"[TELEMETRY-WARNING] Rotation failed: {e}", file=sys.stderr)
|
|
|
|
|
|
def emit_telemetry(record: dict) -> None:
|
|
"""
|
|
Append-only JSONL for query telemetry.
|
|
Fail-open: never breaks the critical path.
|
|
"""
|
|
try:
|
|
os.makedirs(os.path.dirname(TELEMETRY_LOG), exist_ok=True)
|
|
_rotate_telemetry_if_needed()
|
|
with open(TELEMETRY_LOG, "a", encoding="utf-8") as f:
|
|
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
except Exception as e:
|
|
print(f"[TELEMETRY-WARNING] Failed to write telemetry: {e}", file=sys.stderr)
|
|
|
|
|
|
def estimate_tokens(text: str) -> int:
|
|
"""Simple estimate: ~1.3 tokens per word (basic heuristic)."""
|
|
return int(len(text.split()) * 1.3)
|
|
|
|
|
|
# ─── Prompt Injection Sanitization ──────────────────────────────────────────
|
|
|
|
_INJECTION_PATTERNS_CE = [
|
|
# "ignore all previous/prior instructions/directives"
|
|
(re.compile(r"(?i)\bignore\s+all\s+(previous|prior)\s+(instructions|directives|commands|messages|prompts|context)"),
|
|
"[REDACTED]"),
|
|
# "you are/will now become/act/acting as (a/an) AI/assistant..."
|
|
(re.compile(r"(?i)\byou\s+(are|will\s+now)\s+(now\s+)?(become|act|acting)\s+as\s+(a\s+|an\s+)?(AI\s+assistant|assistant|AI|agent|LLM|chatbot|model|system)"),
|
|
"[REDACTED]"),
|
|
# "new instructions/directives/commands follow/above/below"
|
|
(re.compile(r"(?i)\bnew\s+(instructions|directives|commands)\s+(follow|above|below)"),
|
|
"[REDACTED]"),
|
|
# Template injection: {{...}}, ${...}
|
|
(re.compile(r"\{\{.*?\}\}|\$\{.*?\}"), "[REDACTED]"),
|
|
# Markdown/javascript data: URLs in links and images
|
|
(re.compile(r"(?i)(javascript|data)\s*:"), "sanitized:"),
|
|
# Known system prefixes
|
|
(re.compile(r"(?i)\[IMPORTANT:.*?\]|\[SYSTEM:.*?\]"), "[REDACTED]"),
|
|
# Control characters (keep newlines and tabs)
|
|
(re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]"), ""),
|
|
]
|
|
|
|
|
|
def _strip_prompt_injection(text: str) -> str:
|
|
"""Strip known prompt-injection patterns from retrieved text.
|
|
Fail-open: never raises. Empty/null input returns empty string.
|
|
"""
|
|
if not text:
|
|
return ""
|
|
try:
|
|
result = str(text)
|
|
for pattern, replacement in _INJECTION_PATTERNS_CE:
|
|
result = pattern.sub(replacement, result)
|
|
# Normalize excessive whitespace (4+ newlines → 3, 8+ spaces → 1)
|
|
result = re.sub(r"\n{4,}", "\n\n\n", result)
|
|
result = re.sub(r" {8,}", " ", result)
|
|
return result.strip()
|
|
except Exception:
|
|
return str(text)[:400]
|
|
|
|
|
|
# ─── Core ───────────────────────────────────────────────────────────────────
|
|
|
|
def embed_query(text: str) -> Optional[List[float]]:
|
|
"""Generate dense embedding via OpenRouter qwen/qwen3-embedding-8b."""
|
|
if not OPENROUTER_KEY:
|
|
return None
|
|
try:
|
|
resp = requests.post(
|
|
"https://openrouter.ai/api/v1/embeddings",
|
|
headers={
|
|
"Authorization": f"Bearer {OPENROUTER_KEY}",
|
|
"Content-Type": "application/json"
|
|
},
|
|
json={
|
|
"model": EMBEDDING_MODEL,
|
|
"input": text[:MAX_TEXT_LEN]
|
|
},
|
|
timeout=REQUEST_TIMEOUT
|
|
)
|
|
resp.raise_for_status()
|
|
return resp.json()["data"][0]["embedding"]
|
|
except Exception as e:
|
|
# Fail-open: log silently, return None
|
|
print(f"[CE-ERROR] Embedding dense failed: {e}", file=sys.stderr)
|
|
return None
|
|
|
|
|
|
def embed_query_sparse(text: str) -> Optional[Tuple[List[int], List[float]]]:
|
|
"""
|
|
Generate sparse BM25 embedding via FastEmbed (subprocess in ai-lab venv).
|
|
Query text passes via stdin — never embedded in a -c code string.
|
|
Fail-open: if it fails, return None. Caller falls back to dense-only.
|
|
"""
|
|
try:
|
|
result = subprocess.run(
|
|
[_FASTEMBED_PYTHON, "-c", """\
|
|
import os, sys, json
|
|
sys.path.insert(0, os.environ["FASTEMBED_SITEPKGS"])
|
|
from fastembed.sparse import SparseTextEmbedding
|
|
query = sys.stdin.read()
|
|
model = SparseTextEmbedding(model_name="Qdrant/bm25")
|
|
sparse = list(model.embed([query]))[0]
|
|
print(json.dumps({"indices": sparse.indices.tolist(), "values": sparse.values.tolist()}))
|
|
"""],
|
|
input=text,
|
|
capture_output=True, text=True, timeout=15
|
|
)
|
|
data = json.loads(result.stdout.strip())
|
|
return data["indices"], data["values"]
|
|
except Exception as e:
|
|
print(f"[CE-ERROR] Embedding sparse failed: {e}", file=sys.stderr)
|
|
return None
|
|
|
|
|
|
def search_knowledge_base(
|
|
dense_vector: List[float],
|
|
sparse_vector: Optional[Tuple[List[int], List[float]]] = None,
|
|
top_k: int = TOP_K_DEFAULT,
|
|
score_threshold: float = SCORE_THRESHOLD_DEFAULT
|
|
) -> List[Dict]:
|
|
"""
|
|
Hybrid search in Qdrant: dense (semantic) + sparse (BM25 keyword)
|
|
via prefetch + RRF. If sparse fails, fall back to dense-only.
|
|
"""
|
|
try:
|
|
if sparse_vector is not None:
|
|
# Hybrid: prefetch dense + prefetch sparse → RRF
|
|
resp = requests.post(
|
|
f"{QDRANT_URL}/collections/{COLLECTION}/points/query",
|
|
headers={"Content-Type": "application/json"},
|
|
json={
|
|
"prefetch": [
|
|
{"query": dense_vector, "using": "dense", "limit": top_k * 3},
|
|
{"query": {"indices": sparse_vector[0], "values": sparse_vector[1]},
|
|
"using": "sparse", "limit": top_k * 3},
|
|
],
|
|
"query": {"fusion": "rrf"},
|
|
"limit": top_k * 2,
|
|
"with_payload": True,
|
|
},
|
|
timeout=REQUEST_TIMEOUT + 5,
|
|
)
|
|
else:
|
|
# Fallback: dense-only (collections with compatible named vectors)
|
|
resp = requests.post(
|
|
f"{QDRANT_URL}/collections/{COLLECTION}/points/query",
|
|
headers={"Content-Type": "application/json"},
|
|
json={
|
|
"query": dense_vector,
|
|
"using": "dense",
|
|
"limit": top_k,
|
|
"with_payload": True,
|
|
},
|
|
timeout=REQUEST_TIMEOUT,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
results = []
|
|
raw_results = data.get("result", {})
|
|
# API /points/search returns direct list; API /points/query returns {"points": [...]}
|
|
points = raw_results if isinstance(raw_results, list) else raw_results.get("points", [])
|
|
for r in points:
|
|
score = r.get("score", 0)
|
|
if score < score_threshold:
|
|
continue
|
|
payload = r.get("payload", {})
|
|
results.append({
|
|
"id": r.get("id", "unknown"),
|
|
"score": score,
|
|
"title": payload.get("title", "Untitled"),
|
|
"content_preview": _strip_prompt_injection((payload.get("text", "") or "")[:400]),
|
|
"source": payload.get("source", "unknown"),
|
|
"tags": payload.get("tags", [])
|
|
})
|
|
return results
|
|
except Exception as e:
|
|
print(f"[CE-ERROR] Qdrant search failed: {e}", file=sys.stderr)
|
|
return []
|
|
|
|
|
|
# ─── Fallback: Lexical Search in Vault ──────────────────────────────────────
|
|
|
|
def tokenize_query(text: str) -> List[str]:
|
|
"""Tokenize query into relevant terms (lowercase, alphanumeric)."""
|
|
# Remove basic Portuguese/English stopwords
|
|
stopwords = {"o", "a", "os", "as", "um", "uma", "de", "da", "do", "em", "no", "na", "para", "com", "por", "que", "se", "e", "ou", "mas", "the", "a", "an", "is", "are", "was", "were", "be", "been", "have", "has", "had", "do", "does", "did", "will", "would", "could", "should"}
|
|
tokens = re.findall(r'\b[a-zA-Z0-9]+\b', text.lower())
|
|
return [t for t in tokens if t not in stopwords and len(t) > 2]
|
|
|
|
|
|
def lexical_search_in_vault(
|
|
query_terms: List[str],
|
|
top_k: int = TOP_K_DEFAULT,
|
|
vault_root: str = os.environ.get("WIKI_PATH", os.path.expanduser("~/vault/wiki"))
|
|
) -> List[Dict]:
|
|
"""
|
|
Lexical search in .md files under vault/wiki/.
|
|
Rank by term match density / document size.
|
|
Return top-k simulated chunks (or entire file if small).
|
|
"""
|
|
if not query_terms:
|
|
return []
|
|
|
|
md_files = glob.glob(f"{vault_root}/**/*.md", recursive=True)
|
|
if not md_files:
|
|
return []
|
|
|
|
scored = []
|
|
for filepath in md_files:
|
|
try:
|
|
text = Path(filepath).read_text(encoding="utf-8", errors="replace").lower()
|
|
if not text:
|
|
continue
|
|
|
|
matches = 0
|
|
for term in query_terms:
|
|
matches += text.count(term)
|
|
|
|
if matches == 0:
|
|
continue
|
|
|
|
# Heuristic: matches / sqrt(word_count) — favors concise documents
|
|
word_count = max(1, len(text.split()))
|
|
density = matches / (word_count ** 0.5) # sqrt(word_count) to not over-penalize medium texts
|
|
scored.append({
|
|
"filepath": filepath,
|
|
"matches": matches,
|
|
"word_count": word_count,
|
|
"density": density,
|
|
"title": Path(filepath).stem,
|
|
"text": text[:2000], # truncate for return
|
|
})
|
|
except Exception:
|
|
continue
|
|
|
|
if not scored:
|
|
return []
|
|
|
|
# Sort by descending density
|
|
scored.sort(key=lambda x: x["density"], reverse=True)
|
|
|
|
results = []
|
|
for i, item in enumerate(scored[:top_k]):
|
|
# Simulate a result chunk
|
|
results.append({
|
|
"id": f"lexical-{hashlib.md5(item['filepath'].encode()).hexdigest()[:16]}",
|
|
"score": round(min(1.0, item["density"]), 2),
|
|
"title": item["title"],
|
|
"content_preview": _strip_prompt_injection(item["text"][:400]),
|
|
"source": f"vault-{item['filepath'].replace(vault_root, '').lstrip('/')[:40]}",
|
|
"tags": ["fallback", "lexical"],
|
|
"fallback_level": "lexical",
|
|
})
|
|
|
|
if results:
|
|
print(f"[CE-FALLBACK] Lexical search returned {len(results)} results from vault", file=sys.stderr)
|
|
return results
|
|
|
|
|
|
# ─── Fallback: SQLite Keyword Search ────────────────────────────────────────
|
|
|
|
def sqlite_keyword_search(
|
|
query_terms: List[str],
|
|
top_k: int = TOP_K_DEFAULT
|
|
) -> List[Dict]:
|
|
"""
|
|
Search terms in the lineage table (query field) and other state tables.
|
|
Last resort — does not replace vault.
|
|
"""
|
|
if not query_terms:
|
|
return []
|
|
|
|
try:
|
|
with sqlite3.connect(LINEAGE_DB) as conn:
|
|
conn.row_factory = sqlite3.Row
|
|
c = conn.cursor()
|
|
|
|
results = []
|
|
# Search in lineage.query
|
|
placeholders = " OR ".join(["query LIKE ?"] * len(query_terms))
|
|
params = [f"%{term}%" for term in query_terms]
|
|
c.execute(
|
|
f"""
|
|
SELECT lineage_id, session_id, query, generation_context_hash, created_at
|
|
FROM lineage
|
|
WHERE {placeholders}
|
|
ORDER BY created_at DESC
|
|
LIMIT ?
|
|
""",
|
|
params + [top_k * 2]
|
|
)
|
|
for row in c.fetchall():
|
|
results.append({
|
|
"id": f"sqlite-{row['lineage_id'][:16]}",
|
|
"score": 0.5,
|
|
"title": f"Lineage {row['lineage_id'][:8]}...",
|
|
"content_preview": _strip_prompt_injection((row["query"] or "")[:400]),
|
|
"source": f"sqlite-history-{row['session_id']}",
|
|
"tags": ["fallback", "sqlite"],
|
|
"fallback_level": "sqlite",
|
|
})
|
|
|
|
if results:
|
|
print(f"[CE-FALLBACK] SQLite keyword search returned {len(results)} results", file=sys.stderr)
|
|
return results[:top_k]
|
|
except Exception as e:
|
|
print(f"[CE-FALLBACK] SQLite search failed: {e}", file=sys.stderr)
|
|
return []
|
|
|
|
|
|
# ─── Fallback Wrapper ──────────────────────────────────────────────────────
|
|
|
|
def search_with_fallback(
|
|
dense_vector: Optional[List[float]] = None,
|
|
sparse_vector: Optional[Tuple[List[int], List[float]]] = None,
|
|
query_text: str = "",
|
|
top_k: int = TOP_K_DEFAULT,
|
|
score_threshold: float = SCORE_THRESHOLD_DEFAULT
|
|
) -> Tuple[List[Dict], str, float, float]:
|
|
"""
|
|
4-level fallback cascade:
|
|
1. Hybrid (dense + sparse + RRF) — normal mode
|
|
2. Dense-only — if sparse fails or is None
|
|
3. Lexical-only — if Qdrant goes down (ConnectionError, Timeout)
|
|
4. SQLite keyword — if vault is inaccessible
|
|
|
|
Returns (results, fallback_level, qdrant_latency_ms, fallback_latency_ms).
|
|
"""
|
|
fallback_level = "hybrid"
|
|
qdrant_latency_ms = 0.0
|
|
t0 = time.perf_counter()
|
|
|
|
# Level 1 or 2: Qdrant (hybrid if sparse available, otherwise dense-only)
|
|
try:
|
|
t_q0 = time.perf_counter()
|
|
if sparse_vector is not None:
|
|
resp = requests.post(
|
|
f"{QDRANT_URL}/collections/{COLLECTION}/points/query",
|
|
headers={"Content-Type": "application/json"},
|
|
json={
|
|
"prefetch": [
|
|
{"query": dense_vector, "using": "dense", "limit": top_k * 3},
|
|
{"query": {"indices": sparse_vector[0], "values": sparse_vector[1]},
|
|
"using": "sparse", "limit": top_k * 3},
|
|
],
|
|
"query": {"fusion": "rrf"},
|
|
"limit": top_k * 2,
|
|
"with_payload": True,
|
|
},
|
|
timeout=REQUEST_TIMEOUT + 5,
|
|
)
|
|
fallback_level = "hybrid"
|
|
else:
|
|
resp = requests.post(
|
|
f"{QDRANT_URL}/collections/{COLLECTION}/points/query",
|
|
headers={"Content-Type": "application/json"},
|
|
json={
|
|
"query": dense_vector,
|
|
"using": "dense",
|
|
"limit": top_k,
|
|
"with_payload": True,
|
|
},
|
|
timeout=REQUEST_TIMEOUT,
|
|
)
|
|
fallback_level = "dense-only"
|
|
|
|
resp.raise_for_status()
|
|
qdrant_latency_ms = (time.perf_counter() - t_q0) * 1000
|
|
data = resp.json()
|
|
results = []
|
|
raw_results = data.get("result", {})
|
|
points = raw_results if isinstance(raw_results, list) else raw_results.get("points", [])
|
|
for r in points:
|
|
score = r.get("score", 0)
|
|
if score < score_threshold:
|
|
continue
|
|
payload = r.get("payload", {})
|
|
results.append({
|
|
"id": r.get("id", "unknown"),
|
|
"score": score,
|
|
"title": payload.get("title", "Untitled"),
|
|
"content_preview": _strip_prompt_injection((payload.get("text", "") or "")[:400]),
|
|
"source": payload.get("source", "unknown"),
|
|
"tags": payload.get("tags", [])
|
|
})
|
|
|
|
if results:
|
|
fallback_latency_ms = (time.perf_counter() - t0) * 1000
|
|
return results, fallback_level, qdrant_latency_ms, fallback_latency_ms
|
|
except requests.exceptions.ConnectionError:
|
|
print("[CE-FALLBACK] Qdrant unavailable (ConnectionError), falling back to lexical search.", file=sys.stderr)
|
|
except requests.exceptions.Timeout:
|
|
print("[CE-FALLBACK] Qdrant timeout, falling back to lexical search.", file=sys.stderr)
|
|
except Exception as e:
|
|
# Unexpected Qdrant error (e.g. SparseIndexError, 5xx, etc.)
|
|
# If sparse_vector existed, it might be a sparse error — try dense-only
|
|
if sparse_vector is not None and "sparse" in str(e).lower():
|
|
print(f"[CE-FALLBACK] Sparse index failed ('{e}'), trying dense-only...", file=sys.stderr)
|
|
try:
|
|
t_q2 = time.perf_counter()
|
|
resp = requests.post(
|
|
f"{QDRANT_URL}/collections/{COLLECTION}/points/search",
|
|
headers={"Content-Type": "application/json"},
|
|
json={
|
|
"vector": dense_vector,
|
|
"using": "dense",
|
|
"limit": top_k,
|
|
"with_payload": True,
|
|
},
|
|
timeout=REQUEST_TIMEOUT,
|
|
)
|
|
resp.raise_for_status()
|
|
qdrant_latency_ms = (time.perf_counter() - t_q2) * 1000
|
|
data = resp.json()
|
|
results = []
|
|
raw_results = data.get("result", {})
|
|
points = raw_results if isinstance(raw_results, list) else raw_results.get("points", [])
|
|
for r in points:
|
|
score = r.get("score", 0)
|
|
if score < score_threshold:
|
|
continue
|
|
payload = r.get("payload", {})
|
|
results.append({
|
|
"id": r.get("id", "unknown"),
|
|
"score": score,
|
|
"title": payload.get("title", "Untitled"),
|
|
"content_preview": _strip_prompt_injection((payload.get("text", "") or "")[:400]),
|
|
"source": payload.get("source", "unknown"),
|
|
"tags": payload.get("tags", [])
|
|
})
|
|
if results:
|
|
fallback_latency_ms = (time.perf_counter() - t0) * 1000
|
|
print("[CE-FALLBACK] Dense-only worked after sparse error.", file=sys.stderr)
|
|
return results, "dense-only", qdrant_latency_ms, fallback_latency_ms
|
|
except (requests.exceptions.ConnectionError, requests.exceptions.Timeout):
|
|
print("[CE-FALLBACK] Qdrant also unavailable for dense-only — lexical fallback.", file=sys.stderr)
|
|
except Exception as e2:
|
|
print(f"[CE-FALLBACK] Dense-only also failed: {e2}", file=sys.stderr)
|
|
else:
|
|
print(f"[CE-FALLBACK] Qdrant general error ({e}), falling back to lexical.", file=sys.stderr)
|
|
|
|
# Level 3: Lexical search in vault
|
|
terms = tokenize_query(query_text)
|
|
lexical_results = lexical_search_in_vault(terms, top_k=top_k)
|
|
if lexical_results:
|
|
fallback_latency_ms = (time.perf_counter() - t0) * 1000
|
|
return lexical_results, "lexical", qdrant_latency_ms, fallback_latency_ms
|
|
|
|
# Level 4: SQLite keyword search
|
|
sqlite_results = sqlite_keyword_search(terms, top_k=top_k)
|
|
if sqlite_results:
|
|
fallback_latency_ms = (time.perf_counter() - t0) * 1000
|
|
return sqlite_results, "sqlite", qdrant_latency_ms, fallback_latency_ms
|
|
|
|
# Nothing worked
|
|
fallback_latency_ms = (time.perf_counter() - t0) * 1000
|
|
print("[CE-FALLBACK] All fallback levels exhausted.", file=sys.stderr)
|
|
return [], "none", qdrant_latency_ms, fallback_latency_ms
|
|
|
|
|
|
def update_last_accessed_at(chunk_ids: list) -> None:
|
|
"""
|
|
Update last_accessed_at on chunks returned by search.
|
|
Reset decay for chunks that are actually being used.
|
|
Fail-open: never breaks the critical query path.
|
|
"""
|
|
if not chunk_ids:
|
|
return
|
|
try:
|
|
now = datetime.now().astimezone().isoformat()
|
|
requests.post(
|
|
f"{QDRANT_URL}/collections/{COLLECTION}/points/payload",
|
|
headers={"Content-Type": "application/json"},
|
|
json={
|
|
"points": chunk_ids,
|
|
"payload": {"last_accessed_at": now},
|
|
},
|
|
timeout=5,
|
|
)
|
|
except Exception as e:
|
|
print(f"[CE-WARNING] Failed to update last_accessed_at: {e}", file=sys.stderr)
|
|
|
|
|
|
def enhance_prompt(
|
|
user_message: str,
|
|
top_k: int = TOP_K_DEFAULT,
|
|
score_threshold: float = SCORE_THRESHOLD_DEFAULT,
|
|
format_mode: str = "markdown",
|
|
hybrid: bool = True,
|
|
session_id: Optional[str] = None,
|
|
generation_model: str = "unknown",
|
|
) -> str:
|
|
"""
|
|
Return a context block with relevant vault content (hybrid dense+BM25),
|
|
or an empty string if nothing relevant found / services offline.
|
|
Register lineage with chunk provenance and generated context hash.
|
|
"""
|
|
# Resolve session_id
|
|
if session_id is None:
|
|
session_id = os.environ.get("HERMES_SESSION_ID", "standalone")
|
|
|
|
# Skip very short or irrelevant queries
|
|
if len(user_message.strip()) < 5:
|
|
return ""
|
|
|
|
# Very short or social query — skip
|
|
social_keywords = {"hi", "hello", "hey", "yo", "sup", "ok", "thanks", "bye", "okay", "great"}
|
|
if user_message.strip().lower() in social_keywords:
|
|
return ""
|
|
|
|
dense_vector = embed_query(user_message)
|
|
sparse_vector = None
|
|
if hybrid:
|
|
sparse_vector = embed_query_sparse(user_message)
|
|
|
|
# Fallback cascade: hybrid → dense-only → lexical → sqlite
|
|
hits, fallback_level, qdrant_latency_ms, fallback_latency_ms = search_with_fallback(
|
|
dense_vector=dense_vector,
|
|
sparse_vector=sparse_vector,
|
|
query_text=user_message,
|
|
top_k=top_k,
|
|
score_threshold=score_threshold
|
|
)
|
|
|
|
# If empty and we had no embedding, try pure lexical
|
|
if not hits and dense_vector is None:
|
|
terms = tokenize_query(user_message)
|
|
lexical_results = lexical_search_in_vault(terms, top_k=top_k)
|
|
if lexical_results:
|
|
hits = lexical_results
|
|
fallback_level = "lexical"
|
|
|
|
if not hits:
|
|
# Emit telemetry even when empty (records that search ran and returned nothing)
|
|
emit_telemetry({
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
"session_id": session_id,
|
|
"query": user_message,
|
|
"retrieval_mode": fallback_level,
|
|
"retrieved_chunk_ids": [],
|
|
"semantic_scores": [],
|
|
"rerank_scores": [],
|
|
"final_context_tokens": 0,
|
|
"fallback_level": {"hybrid": 0, "dense-only": 1, "lexical": 2, "sqlite": 3, "none": -1}.get(fallback_level, -1),
|
|
"llm_response_quality": None,
|
|
"qdrant_latency_ms": round(qdrant_latency_ms, 2),
|
|
"fallback_latency_ms": round(fallback_latency_ms, 2),
|
|
})
|
|
return ""
|
|
|
|
# Update last_accessed_at for returned chunks (reset decay on real usage)
|
|
retrieved_chunk_ids = [h["id"] for h in hits]
|
|
update_last_accessed_at(retrieved_chunk_ids)
|
|
|
|
if format_mode == "markdown":
|
|
lines = ["\n## Relevant Context from Vault\n"]
|
|
for i, h in enumerate(hits, 1):
|
|
lines.append(f"### [{i}] {h['title']} (score: {h['score']:.2f})")
|
|
lines.append(f"- **Source:** `{h['source']}` | Tags: {', '.join(str(t) for t in h['tags'])}")
|
|
lines.append(f"- **Excerpt:** {h['content_preview']}...")
|
|
lines.append("")
|
|
context_str = "\n".join(lines)
|
|
elif format_mode == "compact":
|
|
parts = []
|
|
for h in hits:
|
|
parts.append(f"[{h['title'][:40]}] (s:{h['score']:.2f})")
|
|
context_str = "\n".join(parts)
|
|
else:
|
|
context_str = json.dumps(hits, indent=2, ensure_ascii=False)
|
|
|
|
generation_context_hash = hashlib.sha256(context_str.encode("utf-8")).hexdigest()[:32]
|
|
|
|
register_lineage(
|
|
session_id=session_id,
|
|
query=user_message,
|
|
retrieved_chunk_ids=retrieved_chunk_ids,
|
|
generation_context_hash=generation_context_hash,
|
|
generation_model=generation_model,
|
|
)
|
|
|
|
# Extract scores for telemetry
|
|
semantic_scores = [h.get("score", 0) for h in hits]
|
|
# For rerank_scores, use the same scores (internal Qdrant RRF, not exposed via simple API)
|
|
rerank_scores = semantic_scores[:]
|
|
|
|
final_context_tokens = estimate_tokens(context_str)
|
|
|
|
emit_telemetry({
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
"session_id": session_id,
|
|
"query": user_message,
|
|
"retrieval_mode": fallback_level,
|
|
"retrieved_chunk_ids": retrieved_chunk_ids,
|
|
"semantic_scores": semantic_scores,
|
|
"rerank_scores": rerank_scores,
|
|
"final_context_tokens": final_context_tokens,
|
|
"fallback_level": {"hybrid": 0, "dense-only": 1, "lexical": 2, "sqlite": 3, "none": -1}.get(fallback_level, -1),
|
|
"llm_response_quality": None,
|
|
"qdrant_latency_ms": round(qdrant_latency_ms, 2),
|
|
"fallback_latency_ms": round(fallback_latency_ms, 2),
|
|
})
|
|
|
|
return context_str
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Context Enhancer — Hybrid search in knowledge_base")
|
|
parser.add_argument("query", nargs="?", help="User message/query")
|
|
parser.add_argument("--top-k", type=int, default=TOP_K_DEFAULT, help=f"Number of results (default: {TOP_K_DEFAULT})")
|
|
parser.add_argument("--threshold", type=float, default=SCORE_THRESHOLD_DEFAULT, help=f"Minimum score (default: {SCORE_THRESHOLD_DEFAULT})")
|
|
parser.add_argument("--format", choices=["markdown", "compact", "json"], default="markdown", help="Output format")
|
|
parser.add_argument("--silent", action="store_true", help="Silent — return empty on error")
|
|
parser.add_argument("--hybrid-off", action="store_true", help="Force dense-only (ignore BM25)")
|
|
args = parser.parse_args()
|
|
|
|
if not args.query:
|
|
args.query = sys.stdin.read().strip() or "how ARQ worker works"
|
|
|
|
try:
|
|
result = enhance_prompt(
|
|
args.query,
|
|
top_k=args.top_k,
|
|
score_threshold=args.threshold,
|
|
format_mode=args.format,
|
|
hybrid=not args.hybrid_off,
|
|
generation_model=os.environ.get("HERMES_MODEL", "unknown"),
|
|
)
|
|
print(result)
|
|
except Exception as e:
|
|
if not args.silent:
|
|
print(f"[CE-ERROR] {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|