Solution published: update src/utils.py
This commit is contained in:
+8
-12
@@ -3,17 +3,15 @@ from typing import List
|
||||
|
||||
from langchain_community.document_loaders import TextLoader
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from langchain_chroma import Chroma
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
|
||||
# Qdrant collection name
|
||||
# Chroma collection name
|
||||
COLLECTION_NAME = "faq"
|
||||
# Qdrant server URL (default local)
|
||||
QDRANT_URL = "http://localhost:6333"
|
||||
|
||||
|
||||
def load_faq_to_qdrant(md_dir: str = "data") -> None:
|
||||
"""Load all .md files from md_dir into a Qdrant vector store.
|
||||
def load_faq_to_chroma(md_dir: str = "data") -> None:
|
||||
"""Load all .md files from md_dir into a Chroma vector store.
|
||||
The store is created/updated at COLLECTION_NAME.
|
||||
"""
|
||||
loader = TextLoader
|
||||
@@ -25,21 +23,19 @@ def load_faq_to_qdrant(md_dir: str = "data") -> None:
|
||||
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||
docs = splitter.split_documents(all_docs)
|
||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
# Create or update Qdrant collection
|
||||
QdrantVectorStore.from_documents(
|
||||
# Create or update Chroma collection
|
||||
Chroma.from_documents(
|
||||
docs,
|
||||
embeddings,
|
||||
collection_name=COLLECTION_NAME,
|
||||
url=QDRANT_URL,
|
||||
)
|
||||
|
||||
|
||||
def search_course_docs(query: str, k: int = 3) -> List[str]:
|
||||
"""Return top k document snippets from the persisted Qdrant store."""
|
||||
"""Return top k document snippets from the persisted Chroma store."""
|
||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
store = QdrantVectorStore(
|
||||
store = Chroma(
|
||||
collection_name=COLLECTION_NAME,
|
||||
url=QDRANT_URL,
|
||||
embedding_function=embeddings,
|
||||
)
|
||||
results = store.similarity_search(query, k=k)
|
||||
|
||||
Reference in New Issue
Block a user