Solution compiled successfully. Ready to publish.: update src/utils.py
This commit is contained in:
+22
-11
@@ -1,19 +1,20 @@
|
|||||||
import os
|
|
||||||
import pathlib
|
import pathlib
|
||||||
from typing import List
|
from typing import List
|
||||||
|
|
||||||
from langchain_community.document_loaders import TextLoader
|
from langchain_community.document_loaders import TextLoader
|
||||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||||
from langchain_chroma import Chroma
|
from langchain_qdrant import QdrantVectorStore
|
||||||
from langchain_ollama import OllamaEmbeddings
|
from langchain_ollama import OllamaEmbeddings
|
||||||
|
|
||||||
CHROMA_PATH = pathlib.Path("./chroma_faq")
|
# Qdrant collection name
|
||||||
CHROMA_PATH.mkdir(parents=True, exist_ok=True)
|
COLLECTION_NAME = "faq"
|
||||||
|
# Qdrant server URL (default local)
|
||||||
|
QDRANT_URL = "http://localhost:6333"
|
||||||
|
|
||||||
|
|
||||||
def load_faq_to_chroma(md_dir: str = "data") -> None:
|
def load_faq_to_qdrant(md_dir: str = "data") -> None:
|
||||||
"""Load all .md files from md_dir into a Chroma vector store.
|
"""Load all .md files from md_dir into a Qdrant vector store.
|
||||||
The store is persisted at CHROMA_PATH.
|
The store is created/updated at COLLECTION_NAME.
|
||||||
"""
|
"""
|
||||||
loader = TextLoader
|
loader = TextLoader
|
||||||
all_docs = []
|
all_docs = []
|
||||||
@@ -24,12 +25,22 @@ def load_faq_to_chroma(md_dir: str = "data") -> None:
|
|||||||
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||||
docs = splitter.split_documents(all_docs)
|
docs = splitter.split_documents(all_docs)
|
||||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||||
Chroma.from_documents(docs, embeddings, persist_directory=str(CHROMA_PATH))
|
# Create or update Qdrant collection
|
||||||
|
QdrantVectorStore.from_documents(
|
||||||
|
docs,
|
||||||
|
embeddings,
|
||||||
|
collection_name=COLLECTION_NAME,
|
||||||
|
url=QDRANT_URL,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def search_course_docs(query: str, k: int = 3) -> List[str]:
|
def search_course_docs(query: str, k: int = 3) -> List[str]:
|
||||||
"""Return top k document snippets from the persisted Chroma store."""
|
"""Return top k document snippets from the persisted Qdrant store."""
|
||||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||||
chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=embeddings)
|
store = QdrantVectorStore(
|
||||||
results = chroma.similarity_search(query, k=k)
|
collection_name=COLLECTION_NAME,
|
||||||
|
url=QDRANT_URL,
|
||||||
|
embedding_function=embeddings,
|
||||||
|
)
|
||||||
|
results = store.similarity_search(query, k=k)
|
||||||
return [doc.page_content for doc in results]
|
return [doc.page_content for doc in results]
|
||||||
|
|||||||
Reference in New Issue
Block a user