rag_dspy.py

import os
import dspy
from dotenv import load_dotenv
from pydantic import BaseModel

load_dotenv()

# Using OpenRouter. Switch to another LLM provider as needed
lm = dspy.LM(
    model="openrouter/google/gemini-2.0-flash-001",
    api_base="https://openrouter.ai/api/v1",
    api_key=os.environ["OPENROUTER_API_KEY"],
)
dspy.configure(lm=lm)


class Response(BaseModel):
    question: str
    answer: str


class RAG(dspy.Signature):
    """
    Answer the question in full sentences using the provided context.
    Do not make up an answer. If the information is not provided in the context, say so clearly.
    """
    question: str = dspy.InputField()
    context: str = dspy.InputField()
    answer: Response = dspy.OutputField()


def main(question: str, context: str) -> None:
    if not question.strip():
        raise ValueError("Question cannot be empty.")
    rag = dspy.Predict(RAG)
    r = rag(question=question, context=context)
    print(r.answer)


if __name__ == "__main__":
    test_cases = [
        {
            "question": "How many people were in the study?",
            "context": "According to the study, eating a daily average of one-quarter or more of a serving of processed red meats (roughly two slices of bacon, one and half slices of bologna, or a hot dog), compared to those consuming a minimal amount (less than one-tenth a serving each day), had a 13 percent higher risk of developing dementia.",
        },
        {
            "question": "What is the association between processed red meat and dementia?",
            "context": "According to the study, eating a daily average of one-quarter or more of a serving of processed red meats (roughly two slices of bacon, one and half slices of bologna, or a hot dog), compared to those consuming a minimal amount (less than one-tenth a serving each day), had a 13 percent higher risk of developing dementia.",
        },
        {
            "question": "",
            "context": "According to the study, eating a daily average of one-quarter or more of a serving of processed red meats (roughly two slices of bacon, one and half slices of bologna, or a hot dog), compared to those consuming a minimal amount (less than one-tenth a serving each day), had a 13 percent higher risk of developing dementia.",
        },
    ]
    for case in test_cases:
        main(case["question"], case["context"])
        print("---")