Re-publish after adding utils: update src/utils.py
This commit is contained in:
@@ -6,37 +6,23 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
|
||||
CHROMA_PATH = Path("./chroma_faq")
|
||||
|
||||
|
||||
def load_faq_to_chroma(md_dir: str = "data"):
|
||||
"""Load all .md files from md_dir into a persistent Chroma store.
|
||||
The function will create or update the store at CHROMA_PATH.
|
||||
"""
|
||||
# Ensure directory exists
|
||||
Path(md_dir).mkdir(parents=True, exist_ok=True)
|
||||
# Gather all markdown files
|
||||
md_files = list(Path(md_dir).glob("*.md"))
|
||||
if not md_files:
|
||||
raise FileNotFoundError(f"No .md files found in {md_dir}")
|
||||
|
||||
# Read and split documents
|
||||
texts = []
|
||||
for md_file in md_files:
|
||||
text = md_file.read_text(encoding="utf-8")
|
||||
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||
texts.extend(splitter.split_text(text))
|
||||
|
||||
# Create embeddings
|
||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
|
||||
# Persist to Chroma
|
||||
chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=embeddings)
|
||||
chroma.add_texts(texts)
|
||||
chroma.persist()
|
||||
return chroma
|
||||
|
||||
|
||||
def search_course_docs(query: str, k: int = 3):
|
||||
"""Search the persistent Chroma store for the top k documents matching query."""
|
||||
chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=OllamaEmbeddings(model="nomic-embed-text"))
|
||||
retriever = chroma.as_retriever(search_kwargs={"k": k})
|
||||
return retriever.get_relevant_documents(query)
|
||||
|
||||
Reference in New Issue
Block a user