65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
"""
|
|
Vector store implementation using ChromaDB.
|
|
|
|
Provides methods to add documents and perform similarity search.
|
|
"""
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List, Dict
|
|
|
|
from chromadb import Client as ChromadbClient
|
|
from chromadb.config import Settings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_ollama import OllamaEmbeddings
|
|
|
|
# Ensure persistent directory exists
|
|
CHROMA_DIR = Path("./chroma_db")
|
|
CHROMA_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Embedding model via Ollama
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
|
|
# Chroma client with persistence
|
|
client = ChromadbClient(Settings(persist_directory=str(CHROMA_DIR)))
|
|
collection_name = "knowledge"
|
|
if collection_name not in client.list_collections():
|
|
client.create_collection(name=collection_name)
|
|
col = client.get_or_create_collection(name=collection_name)
|
|
|
|
# Chunker from chunker.py
|
|
from chunker import chunker
|
|
|
|
class VectorStore:
|
|
def add_documents(self, docs: List[str], metadatas: List[Dict]):
|
|
"""Add documents to the collection.
|
|
|
|
Parameters
|
|
----------
|
|
docs: list of raw text strings.
|
|
metadatas: list of metadata dicts corresponding to each doc.
|
|
"""
|
|
# Split into chunks and embed
|
|
all_chunks = []
|
|
all_metadatas = []
|
|
for doc, meta in zip(docs, metadatas):
|
|
chunks = chunker.split_text(doc)
|
|
all_chunks.extend(chunks)
|
|
all_metadatas.extend([meta] * len(chunks))
|
|
|
|
# Generate embeddings via Ollama
|
|
embeds = embeddings.embed_documents(all_chunks)
|
|
ids = [f"chunk_{i}" for i in range(len(all_chunks))]
|
|
col.add(ids=ids, documents=all_chunks, metadatas=all_metadatas, embeddings=embeds)
|
|
|
|
def similarity_search(self, query: str, k: int = 5) -> List[Dict]:
|
|
"""Return top-k similar chunks with metadata."""
|
|
results = col.query(
|
|
query_texts=[query],
|
|
n_results=k,
|
|
include=['documents', 'metadatas'],
|
|
)
|
|
# results is dict with keys documents, metadatas
|
|
docs = results["documents"][0]
|
|
metas = results["metadatas"][0]
|
|
return [{"content": d, "metadata": m} for d, m in zip(docs, metas)]
|