RAG: search your documents, cite the passages

import os

from shipit_agent import Agent
from shipit_agent.llms import build_llm_from_env
from shipit_agent.llms.factory import build_llm_from_settings

os.environ["SHIPIT_ANTHROPIC_MODEL"] = "claude-haiku-5-5"
haiku = build_llm_from_env("anthropic")
gemma = build_llm_from_settings({"model": "google.gemma-4-31b"}, provider="bedrock")

from pathlib import Path

DOCS = {
    "refunds.md": """# Refund policy
Customers can request a full refund within 30 days of a charge. After 30 days, refunds are prorated
by unused days. Refunds go back to the original payment method and take 5 to 10 business days.
Annual plans cancelled after 90 days are not refundable, but stay active until the term ends.""",
    "plans.md": """# Plans and pricing
Starter costs $12 per seat per month and includes 3 projects. Team costs $29 per seat per month with
unlimited projects and SSO. Enterprise is priced per contract and adds audit logs, a 99.95% uptime SLA
and a named support engineer. Annual billing saves two months on every plan.""",
    "security.md": """# Security
All data is encrypted at rest with AES-256 and in transit with TLS 1.3. Customer data stays in the
region chosen at signup: EU (Frankfurt) or US (Virginia). Access to production needs hardware keys,
and every access is logged and reviewed weekly. We run a third-party penetration test twice a year.""",
    "support.md": """# Support
Starter gets email support with a 2 business day response. Team gets chat and email with a 4 hour
response during business hours. Enterprise gets 24/7 phone support with a 1 hour response for
severity-1 incidents. Status updates are posted at status.example.com.""",
}

folder = Path("kb")
folder.mkdir(exist_ok=True)
for name, text in DOCS.items():
    (folder / name).write_text(text)
print(sorted(p.name for p in folder.iterdir()))

from shipit_agent.rag import RAG, HashingEmbedder

rag = RAG.default(embedder=HashingEmbedder(dimension=512))
for path in sorted(folder.glob("*.md")):
    chunks = rag.index_file(str(path))
    print(f"{path.name:12} {len(chunks)} chunk(s)")
print("chunks in the index:", rag.count())
print("sources:", rag.list_sources())

chunk = rag.index_text(
    "Invoices are emailed on the 1st of each month and can be downloaded as PDF from Settings > Billing.",
    source="invoices.md", title="Invoices", metadata={"team": "billing"},
)[0]
print("id:      ", chunk.id)
print("document:", chunk.document_id, "| index:", chunk.chunk_index, "| chars:", chunk.start_char, "-", chunk.end_char)
print("source:  ", chunk.source, "| metadata:", chunk.metadata)
print("text:    ", chunk.text)

from shipit_agent.rag import DocumentChunker, InMemoryBM25Store, InMemoryVectorStore

long_text = " ".join(DOCS.values())
for target in (512, 64, 32):
    small = RAG(vector_store=InMemoryVectorStore(), keyword_store=InMemoryBM25Store(),
                embedder=HashingEmbedder(dimension=256),
                chunker=DocumentChunker(target_tokens=target, overlap_tokens=target // 8))
    print(f"target_tokens={target:3} -> {len(small.index_text(long_text, source='all.md'))} chunks")

context = rag.search("how long does a refund take", top_k=3)
print(f"{context.total_found} found in {context.timing_ms['total_ms']:.2f} ms", {k: round(v, 3) for k, v in context.timing_ms.items()})
def fmt(score) -> str:
    return "  -  " if score is None else f"{score:.3f}"   # None: that search didn't return this chunk

for r in context.results:
    print(f"{fmt(r.score)}  vec {fmt(r.vector_score)}  kw {fmt(r.keyword_score)}  {r.chunk.source:12} {r.chunk.text[:60]}…")

for alpha in (0.0, 0.5, 1.0):
    hits = rag.search("which plan includes SSO", top_k=3, hybrid_alpha=alpha).results
    print(f"alpha={alpha}: {[h.chunk.source for h in hits]}")

from shipit_agent.rag import IndexFilters

for f in (IndexFilters(sources=["kb/support.md"]), IndexFilters(metadata_match={"team": "billing"})):
    hits = rag.search("response time", top_k=3, filters=f).results
    print([h.chunk.source for h in hits])

fine = RAG(vector_store=InMemoryVectorStore(), keyword_store=InMemoryBM25Store(),
           embedder=HashingEmbedder(dimension=256), chunker=DocumentChunker(target_tokens=24, overlap_tokens=0))
fine.index_text(DOCS["security.md"], source="security.md")
hit = fine.search("penetration test", top_k=1, chunks_above=1).results[0]
print("expanded above:", len(hit.expanded_above))
for c in hit.expanded_above:
    print("  above:", c.text[:70])
print("  hit:  ", hit.chunk.text[:70])

first = rag.search("encryption", top_k=1).results[0].chunk
again = rag.fetch_chunk(first.id)
print(again.id == first.id, "|", again.source, "|", again.text[:80])

print(rag.search("uptime SLA", top_k=2).to_prompt_context()[:600])

import json

import boto3

bedrock = boto3.client("bedrock-runtime", region_name=os.getenv("AWS_REGION_NAME", "us-east-1"))

def titan(texts: list[str]) -> list[list[float]]:
    vectors = []
    for text in texts:
        body = json.dumps({"inputText": text, "dimensions": 256})
        response = bedrock.invoke_model(modelId="amazon.titan-embed-text-v2:0", body=body)
        vectors.append(json.loads(response["body"].read())["embedding"])
    return vectors

from shipit_agent.rag import CallableEmbedder

semantic = RAG.default(embedder=CallableEmbedder(titan, dimension=256))
for path in sorted(folder.glob("*.md")):
    semantic.index_file(str(path))
print("indexed", semantic.count(), "chunks with Titan")

question = "can I get my money back if I stop using it after a few weeks?"
for name, index in (("hashing", rag), ("titan", semantic)):
    hits = index.search(question, top_k=4, hybrid_alpha=1.0).results
    print(f"{name:8}", [f"{h.chunk.source.removeprefix('kb/')} {fmt(h.vector_score)}" for h in hits])

from shipit_agent.rag import LLMReranker

reranked = RAG.default(embedder=CallableEmbedder(titan, dimension=256), reranker=LLMReranker(llm=haiku))
for path in sorted(folder.glob("*.md")):
    reranked.index_file(str(path))

q = "which plan gets help fastest when production is down?"
for flag in (False, True):
    hits = reranked.search(q, top_k=3, enable_reranking=flag).results
    print("reranked" if flag else "original", [h.chunk.source for h in hits],
          [fmt(h.rerank_score) for h in hits] if flag else "")

doc_id = rag.search("Starter costs", top_k=1).results[0].chunk.document_id
rag.reindex(doc_id, content="# Plans and pricing\nStarter now costs $15 per seat per month.")
print(rag.search("Starter price per seat", top_k=1).results[0].chunk.text)

before = rag.count()
rag.delete_document(doc_id)
print("chunks:", before, "->", rag.count())
rag.index_file(str(folder / "plans.md"))

agent = Agent(llm=haiku, rag=semantic)
print("tools:", [t.name for t in agent.tools])
result = agent.run("If I cancel a monthly Team plan after 45 days, what do I get back, and how long does it take?")
print(result.output)

for s in result.rag_sources:
    print(f"[{s.index}] {s.source:12} score {s.score:.3f}  {s.text[:70]}…")

for e in Agent(llm=haiku, rag=semantic).stream("Where is EU customer data stored, and how is it protected?"):
    if e.type == "tool_called":
        print("search:", e.payload["tool"], e.payload["arguments"])
    elif e.type == "tool_completed":
        print("  ->", e.payload["output"][:110].replace("\n", " "), "…")
    elif e.type == "run_completed":
        print("\n" + e.payload["output"])

print(Agent(llm=haiku, rag=semantic).run("Do you offer a student discount?").output)

result = Agent(llm=gemma, rag=semantic).run(
    "If I cancel a monthly Team plan after 45 days, what do I get back, and how long does it take?")
print(result.output)
print("sources:", [s.source for s in result.rag_sources])

custom = Agent(llm=haiku, tools=semantic.as_tools(),
               prompt="Answer from the knowledge base in one sentence, then cite like [1].")
print([t.name for t in custom.tools])
print(custom.run("What does Enterprise add over Team?").output)

import shutil

shutil.rmtree(folder)