generate_answer.py

from __future__ import annotations

from pathlib import Path

import ollama

from query_rag import search


DB_DIR = Path("vector_db")


def build_context(results: list[tuple[float, dict]]) -> str:
    parts = []

    for rank, (score, record) in enumerate(results, start=1):
        parts.append(
            f"""
[Result {rank}]
Chapter: {record['chapter']}
Subtopic: {record['subtopic']}
Pages: {record['page_start']} - {record['page_end']}

Content:
{record['text']}
"""
        )

    return "\n\n".join(parts)


def generate(question: str, top_k: int = 3) -> str:
    results = search(DB_DIR, question, top_k=top_k)

    context = build_context(results)

    prompt = f"""
You are an Operating Systems tutor.

Answer the user's question ONLY using the retrieved textbook context.

If the answer is not present in the context, say:
"Information not found in retrieved context."

Retrieved Context:
{context}

Question:
{question}

Answer:
"""

    response = ollama.chat(
        model="phi3",
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
    )

    return response["message"]["content"]


if __name__ == "__main__":
    question = input("Ask a question: ")

    answer = generate(question)

    print("\n")
    print("=" * 80)
    print(answer)