Удалить 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