diff --git a/aborist/cli.py b/aborist/cli.py index 405bca0..11d8201 100644 --- a/aborist/cli.py +++ b/aborist/cli.py @@ -262,6 +262,67 @@ def _cmd_ask(args: argparse.Namespace) -> int: return 0 if result.get("status") in ("cache_hit", "cache_miss_then_written") else 1 +def _cmd_query(args: argparse.Namespace) -> int: + """Multi-source RAG: question -> top-K corpus docs -> Hermes -> cache.""" + import os + + from aborist.qa.client import OpenAICompatibleClient, StubClient + from aborist.qa.query import query + + base_url = args.endpoint or os.environ.get( + "ABORIST_LLM_ENDPOINT", "https://hermes.ai.unturf.com/v1" + ) + model = args.model or os.environ.get( + "ABORIST_LLM_MODEL", + "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic", + ) + revision = os.environ.get("ABORIST_LLM_REVISION", "") + quantization = os.environ.get("ABORIST_LLM_QUANTIZATION", "fp8-dynamic") + api_key = os.environ.get("ABORIST_LLM_API_KEY") + + client: object + if args.dry_run: + client = StubClient( + answer="[STUB] dry-run: would have asked Hermes-3 with the assembled context." + ) + else: + client = OpenAICompatibleClient(base_url=base_url, api_key=api_key) + + qa_db = args.qa_db + if qa_db is None: + if args.global_shards_dir: + qa_db = Path(args.global_shards_dir) / "qa.db" + else: + qa_db = Path.home() / ".aborist" / "qa.db" + qa_db = Path(qa_db) + + shards_dir = ( + Path(args.global_shards_dir) if args.global_shards_dir else None + ) + single_db = None if shards_dir else args.db + + result = query( + question=args.question, + qa_db=qa_db, + chat_client=client, + model_id=model, + revision=revision, + quantization=quantization, + shards_dir=shards_dir, + single_db=single_db, + top_k=args.top_k, + over_fetch=args.over_fetch, + max_context_chars=args.max_context_chars, + ) + + print(json.dumps(result, indent=2)) + return ( + 0 + if result.get("status") in ("cache_hit", "cache_miss_then_written") + else 1 + ) + + def _cmd_providence(args: argparse.Namespace) -> int: """List providence_cache records for a document URI or source_root.""" conn = ( @@ -707,6 +768,44 @@ def build_parser() -> argparse.ArgumentParser: ) ask_cmd.set_defaults(func=_cmd_ask) + query_cmd = sub.add_parser( + "query", + help="multi-source RAG: question -> top-K corpus docs -> Hermes -> cache", + ) + query_cmd.add_argument("question", help="the question to ask") + query_cmd.add_argument( + "--top-k", dest="top_k", type=int, default=8, + help="max distinct source documents in context (default 8)", + ) + query_cmd.add_argument( + "--over-fetch", dest="over_fetch", type=int, default=32, + help="FTS5 hits to fetch per shard before dedup (default 32)", + ) + query_cmd.add_argument( + "--max-context-chars", dest="max_context_chars", type=int, default=60000, + help="cap on assembled context bytes (default 60000)", + ) + query_cmd.add_argument( + "--qa-db", dest="qa_db", default=None, + help=( + "providence_cache target DB. default: /qa.db, or " + "~/.aborist/qa.db when no shards-dir" + ), + ) + query_cmd.add_argument( + "--model", default=None, + help="model_id (default $ABORIST_LLM_MODEL or hermes-3)", + ) + query_cmd.add_argument( + "--endpoint", default=None, + help="OpenAI-compatible base URL (default $ABORIST_LLM_ENDPOINT)", + ) + query_cmd.add_argument( + "--dry-run", dest="dry_run", action="store_true", + help="use StubClient — assembles context but skips the LLM call", + ) + query_cmd.set_defaults(func=_cmd_query) + prov_cmd = sub.add_parser( "providence", help="list providence_cache records", diff --git a/aborist/qa/__init__.py b/aborist/qa/__init__.py index 75e5923..9e4d025 100644 --- a/aborist/qa/__init__.py +++ b/aborist/qa/__init__.py @@ -8,6 +8,7 @@ from aborist.qa.keys import ( model_profile_hash, question_hash, ) +from aborist.qa.query import DEFAULT_QUERY_POLICY, query from aborist.qa.runner import DEFAULT_POLICY, ask __all__ = [ @@ -20,5 +21,7 @@ __all__ = [ "model_profile_hash", "question_hash", "DEFAULT_POLICY", + "DEFAULT_QUERY_POLICY", "ask", + "query", ] diff --git a/aborist/qa/query.py b/aborist/qa/query.py new file mode 100644 index 0000000..c97d748 --- /dev/null +++ b/aborist/qa/query.py @@ -0,0 +1,339 @@ +"""Multi-source corpus Q&A. + +Pose a question, the tree finds related cached docs, assembles them as context, +asks Hermes, caches the answer. + +The flow: + 1. FTS5 search across all shards (chunks_fts can't be UNION'd in views, + so we query each shard's index independently and merge by score). + 2. Pick top-K distinct documents within a character budget. + 3. Compute `context_root` = Merkle root over the sorted source + document_roots — that's the "source" dimension of v9.8's 8-dim + cache_key for this multi-source answer. + 4. Cache lookup; hit returns STRICT immediately. + 5. Miss calls Hermes via the OpenAI-compatible client. + 6. Persist record with merkle_proof = {context_root, sources: [...]} + so a verifier can recompute context_root and check each source + document_root against its shard. + +Per-source proofs are not bundled here (the source roots themselves are +already content-addressed). A verifier asks the shards for any specific +chunk's proof on demand. +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass +from pathlib import Path + +from aborist import ( + CANONICALIZATION_VERSION, + CHUNKING_VERSION, + SCHEMA_VERSION, +) +from aborist.merkle import MerkleTree +from aborist.qa.client import ChatClient +from aborist.qa.keys import ( + cache_key, + conversation_hash, + governance_policy_hash, + model_profile_hash, + question_hash, +) +from aborist.search import FTS5Backend +from aborist.store import ( + append_audit, + connect, + discover_shards, + transaction, +) + + +DEFAULT_QUERY_POLICY = { + "system_prompt": ( + "You are answering a question using ONLY the sources provided below. " + "Each source is delimited by '=== Source: ===' headers. " + "Cite which source supports each claim by quoting it inline. " + "If the sources do not contain the answer, say 'I don't know based on " + "the provided sources.' Do not speculate beyond what is written." + ), + "temperature": 0.1, + "top_p": 1.0, + "max_tokens": 768, +} + + +@dataclass +class _Hit: + document_root: str + document_uri: str + title: str | None + score: float + shard_path: str + chunk_idx: int + + +def _search_corpus( + shards_dir: Path | None, + single_db: Path | None, + question: str, + over_fetch: int, +) -> list[_Hit]: + """FTS5 across shards (or single DB). Dedupe by document_root.""" + paths: list[Path] + if shards_dir is not None: + paths = discover_shards(shards_dir) + elif single_db is not None: + paths = [Path(single_db)] + else: + paths = [] + + raw: list[tuple] = [] + for p in paths: + conn = connect(p) + try: + backend = FTS5Backend(conn) + hits = backend.search(question, limit=over_fetch) + for h in hits: + raw.append( + ( + h.score, + h.document_root, + h.document_uri, + h.title, + h.chunk_idx, + str(p.resolve()), + ) + ) + finally: + conn.close() + + raw.sort(key=lambda r: -r[0]) + seen: set[str] = set() + out: list[_Hit] = [] + for score, root, uri, title, idx, sp in raw: + if root in seen: + continue + seen.add(root) + out.append( + _Hit( + document_root=root, + document_uri=uri, + title=title, + score=score, + shard_path=sp, + chunk_idx=idx, + ) + ) + return out + + +def _load_doc_text(shard_path: str, document_root: str) -> str | None: + """Concatenate all hot chunks of a document. Returns None if cold or missing.""" + conn = connect(shard_path) + try: + rows = conn.execute( + "SELECT content FROM chunks " + "WHERE document_root = ? AND content IS NOT NULL " + "ORDER BY idx ASC", + (document_root,), + ).fetchall() + finally: + conn.close() + if not rows: + return None + return "\n\n".join(r["content"] for r in rows) + + +def _context_root(source_roots: list[str]) -> str: + """Merkle root over sorted source document_roots — the v9.8 'source' dim + for multi-source answers. Sorting makes the root deterministic regardless + of search ranking order.""" + if not source_roots: + return "00" * 32 + sorted_roots = sorted(source_roots) + if len(sorted_roots) == 1: + return sorted_roots[0] + leaves = [bytes.fromhex(r) for r in sorted_roots] + return MerkleTree.build(leaves).root.hex() + + +def query( + *, + question: str, + qa_db: Path, + chat_client: ChatClient, + model_id: str, + revision: str = "", + quantization: str = "", + shards_dir: Path | None = None, + single_db: Path | None = None, + top_k: int = 8, + over_fetch: int = 32, + max_context_chars: int = 60000, + policy: dict | None = None, + chain: str = "private", +) -> dict: + """Answer `question` using the corpus. Cache to qa_db. Returns a result dict.""" + policy = policy or DEFAULT_QUERY_POLICY + + # 1. Search. + hits = _search_corpus(shards_dir, single_db, question, over_fetch) + if not hits: + return {"status": "no_sources", "msg": "FTS5 search returned no hits"} + + # 2. Pull doc texts within budget. + chosen: list[_Hit] = [] + context_parts: list[str] = [] + char_budget = max_context_chars + for h in hits[:top_k]: + text = _load_doc_text(h.shard_path, h.document_root) + if not text: + continue + if len(text) > char_budget: + text = text[:char_budget] + context_parts.append( + f"=== Source: {h.document_uri} ===\n{text}" + ) + chosen.append(h) + char_budget -= len(text) + if char_budget <= 0: + break + + if not chosen: + return {"status": "no_sources", "msg": "top-k hits had cold or empty content"} + + context = "\n\n".join(context_parts) + + # 3. Build messages + hashes. + messages = [ + {"role": "system", "content": policy["system_prompt"]}, + { + "role": "user", + "content": f"Sources:\n\n{context}\n\n---\n\nQuestion: {question}", + }, + ] + context_root = _context_root([h.document_root for h in chosen]) + qhash = question_hash(question) + mhash = model_profile_hash(model_id, revision, quantization) + chash = conversation_hash(messages) + ghash = governance_policy_hash(policy) + ckey = cache_key( + context_root, + qhash, + mhash, + chash, + ghash, + SCHEMA_VERSION, + CANONICALIZATION_VERSION, + CHUNKING_VERSION, + ) + + qa_conn = connect(qa_db) + try: + # 4. Cache lookup. + cached = qa_conn.execute( + "SELECT * FROM providence_cache " + "WHERE cache_key = ? AND falsification_state = 'live'", + (ckey,), + ).fetchone() + if cached is not None: + now = int(time.time()) + with transaction(qa_conn): + qa_conn.execute( + "UPDATE providence_cache " + "SET hit_count = hit_count + 1, last_hit_at = ? " + "WHERE cache_key = ?", + (now, ckey), + ) + return { + "status": "cache_hit", + "audit_mode": "STRICT", + "cache_key": ckey, + "context_root": context_root, + "answer_text": cached["answer_text"], + "sources": json.loads(cached["merkle_proof"])["sources"], + } + + # 5. Cache miss — call LLM. + answer_text = chat_client.chat_completion( + messages, + model=model_id, + temperature=policy["temperature"], + max_tokens=policy["max_tokens"], + top_p=policy.get("top_p", 1.0), + ) + + # 6. Persist record + audit event. + proof_obj = { + "context_root": context_root, + "sources": [ + { + "document_root": h.document_root, + "document_uri": h.document_uri, + "title": h.title, + "score": h.score, + "chunk_idx": h.chunk_idx, + "shard": Path(h.shard_path).name, + } + for h in chosen + ], + } + proof_blob = json.dumps(proof_obj, separators=(",", ":")) + now = int(time.time()) + with transaction(qa_conn): + event_hash = append_audit( + qa_conn, + event_type="providence_query", + subject_root=ckey, + body={ + "context_root": context_root, + "n_sources": len(chosen), + "model_id": model_id, + "revision": revision, + "quantization": quantization, + "answer_chars": len(answer_text), + "context_chars": len(context), + }, + ts=now, + ) + qa_conn.execute( + "INSERT INTO providence_cache " + "(cache_key, source_root, document_uri, question_hash, question_text, " + " answer_text, merkle_proof, model_profile_hash, conversation_hash, " + " governance_policy_hash, schema_version, canonicalization_version, " + " chunking_version, falsification_state, chain, audit_event_hash, " + " created_at, hit_count) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'live', ?, ?, ?, 0)", + ( + ckey, + context_root, + "corpus://multi-source", + qhash, + question, + answer_text, + proof_blob, + mhash, + chash, + ghash, + SCHEMA_VERSION, + CANONICALIZATION_VERSION, + CHUNKING_VERSION, + chain, + event_hash, + now, + ), + ) + finally: + qa_conn.close() + + return { + "status": "cache_miss_then_written", + "audit_mode": "STRICT", + "cache_key": ckey, + "context_root": context_root, + "answer_text": answer_text, + "sources": proof_obj["sources"], + } diff --git a/aborist/search/fts5.py b/aborist/search/fts5.py index 5331d61..6abbac6 100644 --- a/aborist/search/fts5.py +++ b/aborist/search/fts5.py @@ -5,20 +5,36 @@ from __future__ import annotations from aborist.search.base import AuditMode, Hit, SearchBackend -def _escape_fts5(query: str) -> str: - """Wrap each token in double quotes and escape internal quotes. +import re - Prevents user-supplied FTS5 operators from breaking syntax. We accept - space-separated terms and AND them implicitly. +_FTS5_TOKEN_RE = re.compile(r"[A-Za-z][A-Za-z0-9]*") +_FTS5_STOPWORDS = frozenset( """ - tokens = [t for t in query.split() if t] + the a an is are was were be been being + of to in on at for with by from as about into through during + and or but not no nor so yet + what who where when why how which this that these those such + i you he she it we they me him her us them + do does did have has had can could should would will may might + """.split() +) + + +def _escape_fts5(query: str) -> str: + """Build an FTS5 MATCH expression from a free-text query. + + Tokenizes by alpha runs (drops punctuation entirely, including the + `?` that breaks quoted phrases), drops common stopwords, and ORs the + remaining tokens. OR-semantics is intentional for RAG: the LLM sees + BM25-ranked top-K, so partial matches are useful evidence rather than + noise. Stopwords are dropped because they'd otherwise dominate scoring. + """ + raw = _FTS5_TOKEN_RE.findall(query) + tokens = [t for t in raw if t.lower() not in _FTS5_STOPWORDS and len(t) > 1] if not tokens: - return '""' - quoted = [] - for t in tokens: - # FTS5 doubles internal quotes to escape them. - quoted.append('"' + t.replace('"', '""') + '"') - return " ".join(quoted) + # Fallback: keep the query alive even if it was all stopwords. + tokens = [t for t in raw if len(t) > 1] or ['""'] + return " OR ".join('"' + t.replace('"', '""') + '"' for t in tokens) class FTS5Backend(SearchBackend): diff --git a/tests/test_query.py b/tests/test_query.py new file mode 100644 index 0000000..f080d98 --- /dev/null +++ b/tests/test_query.py @@ -0,0 +1,167 @@ +"""Multi-source corpus Q&A: search → context → cache → Hermes. + +Stub client only — no network. Validates: +- FTS5 finds the right docs across a small corpus +- context_root is deterministic (sorted source roots, then Merkle) +- cache hit on identical question returns same record without calling client +- different question -> different cache_key +""" + +from __future__ import annotations + +import json +from typing import Iterator + +from aborist.document import Document +from aborist.ingest import ingest_source +from aborist.qa import query +from aborist.qa.client import StubClient +from aborist.source import Source +from aborist.store import connect + + +class FakeSource(Source): + source_type = "test" + + def __init__(self, docs: list[Document]): + self.docs = docs + + def iter_documents(self) -> Iterator[Document]: + yield from self.docs + + +def _doc(uri: str, content: str) -> Document: + return Document(uri=uri, content=content, source_type="test", title=uri.rsplit("/", 1)[-1]) + + +# Three docs with distinguishable content so FTS5 can pick winners. +DOCS = [ + _doc( + "test://anarchism", + "Anarchism is a political philosophy that opposes the state. " * 12 + + "Mutual aid is central to anarchist theory. " * 8, + ), + _doc( + "test://capitalism", + "Capitalism is an economic system based on private ownership. " * 12 + + "Market exchange and capital accumulation drive growth. " * 8, + ), + _doc( + "test://anarcho-capitalism", + "Anarcho-capitalism combines anarchism's opposition to the state with capitalism's private property. " * 12 + + "Murray Rothbard developed many of its core ideas. " * 8, + ), +] + + +def test_query_picks_relevant_docs_and_writes_record(tmp_path): + main_db = tmp_path / "corpus.db" + qa_db = tmp_path / "qa.db" + conn = connect(main_db) + try: + ingest_source(conn, FakeSource(DOCS)) + finally: + conn.close() + + client = StubClient(answer="anarcho-capitalism rests on private property without a state.") + result = query( + question="What is anarcho-capitalism?", + qa_db=qa_db, + chat_client=client, + model_id="test-model", + single_db=main_db, + top_k=3, + ) + + assert result["status"] == "cache_miss_then_written" + assert result["audit_mode"] == "STRICT" + assert "anarcho-capitalism" in result["answer_text"] + assert len(result["sources"]) >= 1 + assert any("anarcho-capitalism" in s["document_uri"] for s in result["sources"]) + + # context_root is deterministic. + sorted_roots = sorted(s["document_root"] for s in result["sources"]) + if len(sorted_roots) == 1: + assert result["context_root"] == sorted_roots[0] + # Repeat call with same question → cache hit, no LLM call. + n_calls_before = len(client.calls) + result2 = query( + question="What is anarcho-capitalism?", + qa_db=qa_db, + chat_client=client, + model_id="test-model", + single_db=main_db, + top_k=3, + ) + assert result2["status"] == "cache_hit" + assert result2["cache_key"] == result["cache_key"] + assert len(client.calls) == n_calls_before # no new call + + +def test_query_different_question_different_cache(tmp_path): + main_db = tmp_path / "corpus.db" + qa_db = tmp_path / "qa.db" + conn = connect(main_db) + try: + ingest_source(conn, FakeSource(DOCS)) + finally: + conn.close() + + client = StubClient(answer="answer text") + r1 = query( + question="What is anarchism?", + qa_db=qa_db, + chat_client=client, + model_id="m", + single_db=main_db, + ) + r2 = query( + question="What is capitalism?", + qa_db=qa_db, + chat_client=client, + model_id="m", + single_db=main_db, + ) + assert r1["cache_key"] != r2["cache_key"] + assert len(client.calls) == 2 # both missed cache, both called client + + +def test_query_no_sources_when_empty_corpus(tmp_path): + main_db = tmp_path / "empty.db" + qa_db = tmp_path / "qa.db" + connect(main_db).close() # creates schema, no docs + result = query( + question="What is anything?", + qa_db=qa_db, + chat_client=StubClient(), + model_id="m", + single_db=main_db, + ) + assert result["status"] == "no_sources" + + +def test_query_persists_audit_event(tmp_path): + main_db = tmp_path / "corpus.db" + qa_db = tmp_path / "qa.db" + conn = connect(main_db) + try: + ingest_source(conn, FakeSource(DOCS)) + finally: + conn.close() + query( + question="What is anarchism?", + qa_db=qa_db, + chat_client=StubClient(answer="X"), + model_id="m", + single_db=main_db, + ) + # qa.db should now have one providence_query event in its audit chain. + qc = connect(qa_db) + try: + events = qc.execute( + "SELECT event_type FROM audit_events ORDER BY seq" + ).fetchall() + finally: + qc.close() + types = [e["event_type"] for e in events] + assert "providence_query" in types