diff --git a/src/rag_tools.py b/src/rag_tools.py new file mode 100644 index 0000000..fec20e3 --- /dev/null +++ b/src/rag_tools.py @@ -0,0 +1,24 @@ +"""Search tool for the Chroma vector store. + +The function `search_course_docs` performs a similarity search on the +persisted Chroma collection and returns the top *k* documents. +""" + +from typing import List + +from langchain_chroma import Chroma +from langchain_core.documents import Document + +from .config import CHROMA_DIR, CHROMA_TOP_K + +# Load the persistent store once – this is cheap because it just reads +# the SQLite file. +_chroma = Chroma(persist_directory=CHROMA_DIR) + +def search_course_docs(query: str, k: int = CHROMA_TOP_K) -> List[Document]: + """Return the top *k* documents that match *query*. + + The function is intentionally simple; it does not filter by metadata + because the FAQ files are small. + """ + return _chroma.similarity_search(query, k=k) \ No newline at end of file