Ask questions of your own documents

curl -s http://127.0.0.1:8081/v1/models   # the text model
curl -s http://127.0.0.1:8082/v1/models   # the embedding model

from dataclasses import dataclass
from pathlib import Path


@dataclass
class Chunk:
    text: str
    source: str   # which file it came from
    index: int    # where in that file, so you can cite it


def split(text: str, size: int = 1200, overlap: int = 200) -> list[str]:
    step = size - overlap
    pieces = []
    for start in range(0, len(text), step):
        piece = text[start:start + size].strip()
        if len(piece) > 50:  # a fragment that short cannot answer anything
            pieces.append(piece)
    return pieces


def load(folder: str) -> list[Chunk]:
    chunks = []
    for path in sorted(Path(folder).rglob("*")):
        if path.suffix.lower() not in {".txt", ".md"} or not path.is_file():
            continue
        text = path.read_text(encoding="utf-8", errors="replace")
        for index, piece in enumerate(split(text)):
            chunks.append(Chunk(text=piece, source=str(path), index=index))
    return chunks

import pickle
import sys

import numpy as np
from openai import OpenAI

from chunks import load

embedder = OpenAI(base_url="http://127.0.0.1:8082/v1", api_key="not-needed")
EMBED_MODEL = embedder.models.list().data[0].id


def embed(texts: list[str], prefix: str) -> np.ndarray:
    """prefix is 'search_document: ' when storing and 'search_query: ' when
    asking. Nomic models are trained to expect it; without it the two kinds of
    text land in slightly different places and retrieval quietly degrades."""
    out = []
    for start in range(0, len(texts), 32):  # batch, or it takes all evening
        batch = [prefix + text for text in texts[start:start + 32]]
        reply = embedder.embeddings.create(model=EMBED_MODEL, input=batch)
        out.extend(item.embedding for item in reply.data)
        print(f"  embedded {min(start + 32, len(texts))}/{len(texts)}")
    matrix = np.array(out, dtype=np.float32)
    # Normalise once, here. Then similarity is a dot product rather than a
    # division per comparison, and every search afterwards is faster.
    matrix /= np.linalg.norm(matrix, axis=1, keepdims=True)
    return matrix


if __name__ == "__main__":
    folder = sys.argv[1] if len(sys.argv) > 1 else "./docs"
    chunks = load(folder)
    if not chunks:
        sys.exit(f"no .txt or .md files found under {folder}")
    print(f"{len(chunks)} chunks to embed")
    vectors = embed([chunk.text for chunk in chunks], "search_document: ")
    with open("index.pkl", "wb") as handle:
        pickle.dump({"chunks": chunks, "vectors": vectors, "model": EMBED_MODEL}, handle)
    print(f"saved index.pkl — {vectors.shape[0]} vectors of {vectors.shape[1]} dimensions")

import pickle

import numpy as np

from index_docs import embed

with open("index.pkl", "rb") as handle:
    INDEX = pickle.load(handle)


def find(question: str, top: int = 5):
    query = embed([question], "search_query: ")[0]
    # Both sides are normalised, so the dot product is the cosine similarity.
    scores = INDEX["vectors"] @ query
    best = np.argsort(scores)[::-1][:top]
    return [(INDEX["chunks"][i], float(scores[i])) for i in best]


if __name__ == "__main__":
    import sys

    for chunk, score in find(" ".join(sys.argv[1:]) or "what is this about?"):
        print(f"{score:.3f}  {chunk.source}#{chunk.index}")
        print(f"        {chunk.text[:150]}...\n")

import sys

from openai import OpenAI

from search import find

writer = OpenAI(base_url="http://127.0.0.1:8081/v1", api_key="not-needed")
MODEL = writer.models.list().data[0].id

JOB = """You answer questions using only the passages provided.

Rules, in order of importance:
1. If the passages do not contain the answer, say exactly:
   "That is not in these documents." Say nothing else.
2. Never use knowledge from outside the passages, even if you are confident
   it is correct. The user is asking about their documents, not the world.
3. Cite the source of each claim in square brackets, like [contract.txt].
4. If the passages disagree with each other, say so rather than choosing.

Answer in plain language, briefly."""


def ask(question: str, top: int = 5) -> str:
    hits = find(question, top=top)
    if not hits or hits[0][1] < 0.25:
        # Nothing came back that is even plausibly related. Do not waste a
        # generation on it, and do not give the model the chance to improvise.
        return "That is not in these documents."

    passages = "\n\n".join(
        f"[{chunk.source}]\n{chunk.text}" for chunk, _ in hits
    )
    reply = writer.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": JOB},
            {"role": "user", "content": f"Passages:\n\n{passages}\n\nQuestion: {question}"},
        ],
        temperature=0.1,
    )
    return reply.choices[0].message.content


if __name__ == "__main__":
    question = " ".join(sys.argv[1:])
    if not question:
        sys.exit('usage: python ask.py "your question"')
    print(ask(question))