37 lines
1.5 KiB
Python
37 lines
1.5 KiB
Python
"""Wrapper around QdrantVectorStore using Ollama embeddings.
|
|
|
|
This module is intentionally lightweight and can be imported from any
|
|
context without requiring the caller to set up a package structure.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
# Ensure the current directory is in sys.path so that relative imports work
|
|
sys.path.append(os.path.dirname(__file__))
|
|
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_qdrant import QdrantVectorStore
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from typing import List, Dict
|
|
|
|
class QdrantStore:
|
|
"""Wrapper around QdrantVectorStore using Ollama embeddings."""
|
|
|
|
def __init__(self, host: str = "localhost", port: int = 6333, collection_name: str = "knowledge_base"):
|
|
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
self.store = QdrantVectorStore(
|
|
url=f"http://{host}:{port}",
|
|
collection_name=collection_name,
|
|
embeddings=self.embeddings,
|
|
)
|
|
|
|
def add_document(self, content: str, title: str) -> None:
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
chunks = splitter.split_text(content)
|
|
metadatas = [{"title": title, "source": title} for _ in chunks]
|
|
self.store.add_texts(chunks, metadatas=metadatas)
|
|
|
|
def search(self, query: str, k: int = 5) -> List[Dict]:
|
|
results = self.store.similarity_search(query, k=k)
|
|
return [{"content": doc.page_content, "metadata": doc.metadata} for doc in results] |