72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
"""Vector store implementation using Qdrant and Ollama embeddings.
|
|
|
|
This module provides a simple wrapper around QdrantVectorStore that handles
|
|
- Initialization of the Qdrant client and collection.
|
|
- Chunking of documents using RecursiveCharacterTextSplitter.
|
|
- Adding documents with embeddings from Ollama.
|
|
- Semantic search.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any
|
|
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_qdrant import QdrantVectorStore
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain.schema import Document
|
|
|
|
# Global configuration
|
|
QDRANT_HOST = "localhost"
|
|
QDRANT_PORT = 6333
|
|
COLLECTION_NAME = "knowledge_base"
|
|
EMBEDDING_MODEL = "nomic-embed-text"
|
|
|
|
# Initialize embeddings and splitter
|
|
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
|
|
# Create or connect to Qdrant collection
|
|
vector_store = QdrantVectorStore(
|
|
client_kwargs={"host": QDRANT_HOST, "port": QDRANT_PORT},
|
|
collection_name=COLLECTION_NAME,
|
|
embeddings=embeddings,
|
|
)
|
|
|
|
# Ensure collection exists
|
|
if not vector_store.client.has_collection(COLLECTION_NAME):
|
|
vector_store.client.create_collection(COLLECTION_NAME)
|
|
|
|
def add_documents(docs: List[Dict[str, str]]) -> None:
|
|
"""Add a list of documents to the vector store.
|
|
|
|
Each document dict must contain ``title`` and ``content`` keys.
|
|
The content is split into chunks before being stored.
|
|
"""
|
|
documents: List[Document] = []
|
|
for doc in docs:
|
|
title = doc.get("title", "")
|
|
content = doc.get("content", "")
|
|
# Split content into chunks
|
|
chunks = text_splitter.split_text(content)
|
|
for i, chunk in enumerate(chunks):
|
|
meta = {"title": title, "chunk_index": i}
|
|
documents.append(Document(page_content=chunk, metadata=meta))
|
|
vector_store.add_documents(documents)
|
|
|
|
def search(query: str, k: int = 5) -> List[Document]:
|
|
"""Semantic search in the vector store.
|
|
|
|
Returns a list of Documents ordered by relevance.
|
|
"""
|
|
return vector_store.similarity_search(query, k=k)
|
|
|
|
# Convenience: add a single document
|
|
|
|
def add_document(title: str, content: str) -> None:
|
|
add_documents([{"title": title, "content": content}])
|
|
|
|
# Convenience: search and return plain strings
|
|
|
|
def search_text(query: str, k: int = 5) -> List[str]:
|
|
docs = search(query, k)
|
|
return [f"{doc.metadata.get('title', 'Untitled')} (chunk {doc.metadata.get('chunk_index', 0)}): {doc.page_content[:200]}..." for doc in docs] |