Add src/rag_tools.py
This commit is contained in:
@@ -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)
|
||||||
Reference in New Issue
Block a user