| import os | |
| from pathlib import Path | |
| from langchain.chains import RetrievalQA | |
| from langchain_openai import ChatOpenAI, OpenAIEmbeddings | |
| from langchain_chroma import Chroma | |
| from langchain.prompts import PromptTemplate | |
| BASE_DIR = Path(__file__).resolve().parent.parent | |
| DB_DIR = BASE_DIR / "db" | |
| def build_general_qa_chain(model_name=None): | |
| if not DB_DIR.exists(): | |
| print("📦 No DB found. Building vectorstore...") | |
| from scripts import load_documents, chunk_and_embed, setup_vectorstore | |
| load_documents.main() | |
| chunk_and_embed.main() | |
| setup_vectorstore.main() | |
| embedding = OpenAIEmbeddings(model="text-embedding-3-small") | |
| vectorstore = Chroma(persist_directory=str(DB_DIR), embedding_function=embedding) | |
| 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. | |
| Do not answer questions that are completely irrelevant to the main points of the context. | |
| 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 | |