scripts/rag_chat.py

from langchain.chains import RetrievalQA
from langchain_openai import ChatOpenAI
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain.prompts import PromptTemplate
from pathlib import Path

BASE_DIR = Path(__file__).resolve().parent.parent
DB_DIR = str(BASE_DIR / "db")

def build_general_qa_chain(model_name=None):
    embedding = OpenAIEmbeddings(model="text-embedding-3-small")
    vectorstore = Chroma(persist_directory=DB_DIR, embedding_function=embedding)
    
    # Custom prompt with source attribution
    template = """Use the following context to answer the question. 
    If the answer isn't found in the context, use your general knowledge but say so.
    Always cite your sources at the end with 'Source: <filename>' when using course materials.

    Context: {context}
    Question: {question}
    Helpful Answer:"""
    
    QA_PROMPT = PromptTemplate(
        template=template,
        input_variables=["context", "question"]
    )
    
    llm = ChatOpenAI(model_name=model_name or "gpt-4o-mini", temperature=0.0) 
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=vectorstore.as_retriever(search_kwargs={"k": 4}),
        chain_type_kwargs={"prompt": QA_PROMPT},
        return_source_documents=True
    )
    return qa_chain