RAG Crash Course: build a company chatbot in 6 levels

def make_chunks(documents):
    """One paragraph = one chunk. Every chunk keeps its document's labels."""
    chunks = []
    for doc in documents:
        for paragraph in doc["text"].split("\n\n"):
            chunks.append({
                "text": paragraph.strip(),
                "title": doc["title"],
                "doc_type": doc["doc_type"],
                "department": doc["department"],
                "access": doc["access"],
                "year": doc["year"],
            })
    return chunks


def build_index():
    chunks = make_chunks(DOCUMENTS)
    print(f"Cut {len(DOCUMENTS)} documents into {len(chunks)} chunks")

    # ------------------------------------------------------ 2. EMBED
    vectors = embed([c["text"] for c in chunks])
    print(f"Each chunk is now a list of {len(vectors[0])} numbers")

    # ------------------------------------------------------ 3. STORE
    client = connect()

    # A collection is like a table. We tell it: 384 numbers per vector,
    # and compare vectors with cosine similarity.
    client.collections.recreate(
        COLLECTION,
        vectors_config=VectorParams(size=EMBED_SIZE, distance=Distance.Cosine),
    )

    # A point = an id + the vector + a payload (the original text and labels)
    points = [
        PointStruct(id=i, vector=vectors[i], payload=chunks[i])
        for i in range(len(chunks))
    ]
    client.upload_points(COLLECTION, points)
    print(f"Stored {client.points.count(COLLECTION)} points in VectorAI DB")
    return client

def ask(client, question):
    results = dense_search(client, question, limit=3)
    show(results, f"Top 3 chunks for: {question}")
    print("\nAnswer:", answer_from(question, results))

def tokenize(text):
    """Split text into lowercase words. Keeps codes like 'brk-2207' as ONE word."""
    return re.findall(r"[a-z0-9]+(?:-[a-z0-9]+)*", text.lower())


class KeywordSearch:
    """BM25 keyword search over every chunk stored in VectorAI DB."""

    def __init__(self, client):
        # Read every point back from the database (its payload holds the text)
        self.points = []
        for page in client.points.scroll_all(COLLECTION, with_payload=True):
            self.points.extend(page)
        self.bm25 = BM25Okapi([tokenize(p.payload["text"]) for p in self.points])

    def search(self, question, limit=5):
        scores = self.bm25.get_scores(tokenize(question))
        best = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:limit]
        # Same result shape as VectorAI DB, so the fusion helper can merge both lists
        return [
            ScoredPoint(id=self.points[i].id, score=float(scores[i]), payload=self.points[i].payload)
            for i in best
            if scores[i] > 0
        ]


def hybrid_search(client, keyword_search, question, limit=5):
    meaning = dense_search(client, question, limit=20)    # search 1: meaning (VectorAI DB)
    keywords = keyword_search.search(question, limit=20)  # search 2: exact words (BM25)
    merged = reciprocal_rank_fusion([meaning, keywords], limit=limit)  # merge (SDK helper)
    return meaning, keywords, merged

def rewrite(question):
    """Ask the LLM to turn a messy question into a clear search query."""
    return ask_llm(
        "Rewrite this user question as a short, clear search query for a motorcycle "
        "company's documents. Write it in English, the language of the documents. "
        "Use precise technical words. "
        f"Reply with the query only.\n\nQuestion: {question}"
    ).strip()


def search_both(client, original, rewritten, limit=5):
    """Search with both questions and let VectorAI DB merge the two lists."""
    return client.points.query(
        COLLECTION,
        prefetch=[
            PrefetchQuery(query=embed(original), limit=20),   # search 1: what the user typed
            PrefetchQuery(query=embed(rewritten), limit=20),  # search 2: the clean version
        ],
        query={"fusion": Fusion.RRF},                          # merge both on the server
        limit=limit,
    )

reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")  # small, runs locally

def rerank(question, results, top=3):
    """Score every (question, chunk) pair, then sort by the new score."""
    pairs = [(question, r.payload["text"]) for r in results]
    scores = reranker.predict(pairs)
    for r, s in zip(results, scores):
        r.score = float(s)   # replace the first-pass score with the re-ranker's score
    return sorted(results, key=lambda r: r.score, reverse=True)[:top]

def access_filter(role, year=None):
    """Build the search rules for a user."""
    rules = FilterBuilder()
    if role != "hr_manager":
        rules = rules.must_not(Field("access").eq("hr_only"))  # hide confidential HR files
    if year is not None:
        rules = rules.must(Field("year").eq(year))             # only this year's documents
    return None if rules.is_empty() else rules.build()

MAX_STEPS = 4        # safety limit: the agent can't loop forever
WEAK_SCORE = 0.45    # if the best match scores below this, we tell the agent "weak results"

AGENT_RULES = """You are the Star Motors assistant. You can search company documents.

Each turn, reply with ONE JSON object and nothing else:
  {"action": "search", "query": "<search words>", "year": <2025, 2026 or null>}
  {"action": "answer", "text": "<your final answer>"}

Rules:
- Small talk or thanks: answer directly, do NOT search.
- Company questions: search first, then answer only from what you found.
- Comparing years: search once per year, with the year filled in.
- If results are marked WEAK, search again with different, more specific words.
- If you still can't find it, say so."""

def search_tool(client, query, year=None):
    """The agent's only tool: a VectorAI DB search, optionally filtered by year."""
    rules = FilterBuilder().must(Field("year").eq(year)).build() if year else None
    results = dense_search(client, query, limit=3, filter=rules)
    if not results:
        return "No results."
    quality = "WEAK" if results[0].score < WEAK_SCORE else "GOOD"
    found = "\n".join(f"- ({r.score:.2f}) {r.payload['title']}: {r.payload['text']}" for r in results)
    return f"Results ({quality}):\n{found}"


def parse_decision(reply):
    """Pull the FIRST JSON object out of the LLM's reply (it sometimes sends two)."""
    match = re.search(r"\{.*?\}", reply, re.DOTALL)
    return json.loads(match.group(0)) if match else {"action": "answer", "text": reply}


def run_agent(client, question):
    print(f"\n========== User: {question}")
    history = f"User question: {question}"

    for step in range(1, MAX_STEPS + 1):
        decision = parse_decision(ask_llm(history, system=AGENT_RULES))

        if decision["action"] == "answer":
            print(f"Step {step}: ANSWER -> {decision['text']}")
            return decision["text"]

        query = decision.get("query") or question
        year = int(decision["year"]) if str(decision.get("year")).isdigit() else None
        print(f"Step {step}: SEARCH '{query}' (year={year})")
        observation = search_tool(client, query, year)
        print("   " + observation.splitlines()[0])
        history += f"\n\nYou searched: {query} (year={year})\n{observation}"

    print("Stopped: reached the step limit.")