Удалить vector_store.py
This commit is contained in:
@@ -1,64 +0,0 @@
|
||||
"""
|
||||
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)]
|
||||
Reference in New Issue
Block a user