1204 lines
57 KiB
Python
1204 lines
57 KiB
Python
"""
|
||
Vector store management using ChromaDB
|
||
"""
|
||
import os
|
||
import time
|
||
from datetime import datetime, date
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from typing import List, Dict, Any, Tuple
|
||
import chromadb
|
||
from chromadb.config import Settings as ChromaSettings
|
||
from llama_index.vector_stores.chroma import ChromaVectorStore
|
||
from llama_index.core import VectorStoreIndex, StorageContext
|
||
from llama_index.embeddings.ollama import OllamaEmbedding
|
||
from loguru import logger
|
||
from config import settings
|
||
import numpy as np
|
||
import re
|
||
import string
|
||
from rank_bm25 import BM25Okapi
|
||
from utils.query_processor import MetadataFilter
|
||
|
||
|
||
class VectorStoreManager:
|
||
"""Manage ChromaDB vector store and LlamaIndex integration"""
|
||
|
||
def __init__(self):
|
||
self.chroma_client = None
|
||
self.collection = None
|
||
self.vector_store = None
|
||
self.index = None
|
||
self.embed_model = None
|
||
self._initialize()
|
||
|
||
def _cleanup_lock_files(self):
|
||
"""Clean up ChromaDB lock files if they exist (only for PersistentClient mode)"""
|
||
# Skip lock file cleanup if using HttpClient mode
|
||
if settings.CHROMA_SERVER_HOST:
|
||
return
|
||
|
||
try:
|
||
db_path = settings.CHROMA_DB_PATH
|
||
journal_file = os.path.join(db_path, "chroma.sqlite3-journal")
|
||
wal_file = os.path.join(db_path, "chroma.sqlite3-wal")
|
||
shm_file = os.path.join(db_path, "chroma.sqlite3-shm")
|
||
|
||
# Check if journal file exists and is old (more than 1 minute)
|
||
if os.path.exists(journal_file):
|
||
file_age = time.time() - os.path.getmtime(journal_file)
|
||
if file_age > 60: # Older than 1 minute
|
||
logger.warning(f"Found stale journal file (age: {file_age:.0f}s), attempting cleanup...")
|
||
try:
|
||
os.remove(journal_file)
|
||
logger.info("Removed stale journal file")
|
||
except Exception as e:
|
||
logger.warning(f"Could not remove journal file: {e}")
|
||
|
||
# Clean up WAL and SHM files if journal is gone
|
||
if not os.path.exists(journal_file):
|
||
for lock_file in [wal_file, shm_file]:
|
||
if os.path.exists(lock_file):
|
||
try:
|
||
os.remove(lock_file)
|
||
logger.debug(f"Removed lock file: {os.path.basename(lock_file)}")
|
||
except Exception as e:
|
||
logger.debug(f"Could not remove lock file {os.path.basename(lock_file)}: {e}")
|
||
except Exception as e:
|
||
logger.debug(f"Error during lock file cleanup: {e}")
|
||
|
||
def _initialize(self, retry_count: int = 3):
|
||
"""
|
||
Initialize ChromaDB client and collection
|
||
|
||
Args:
|
||
retry_count: Number of retry attempts if database is locked (only for PersistentClient)
|
||
"""
|
||
for attempt in range(retry_count):
|
||
try:
|
||
# Determine client mode: HttpClient or PersistentClient
|
||
if settings.CHROMA_SERVER_HOST:
|
||
# Use HttpClient mode (for Docker/production)
|
||
logger.info(f"Initializing ChromaDB HttpClient mode: {settings.CHROMA_SERVER_HOST}:{settings.CHROMA_SERVER_PORT}")
|
||
self.chroma_client = chromadb.HttpClient(
|
||
host=settings.CHROMA_SERVER_HOST,
|
||
port=settings.CHROMA_SERVER_PORT,
|
||
settings=ChromaSettings(
|
||
anonymized_telemetry=False,
|
||
allow_reset=False
|
||
)
|
||
)
|
||
else:
|
||
# Use PersistentClient mode (for local development)
|
||
if attempt > 0:
|
||
logger.info(f"Retrying ChromaDB initialization (attempt {attempt + 1}/{retry_count})...")
|
||
self._cleanup_lock_files()
|
||
time.sleep(1) # Wait a bit before retrying
|
||
else:
|
||
self._cleanup_lock_files()
|
||
|
||
logger.info(f"Initializing ChromaDB PersistentClient mode: {settings.CHROMA_DB_PATH}")
|
||
self.chroma_client = chromadb.PersistentClient(
|
||
path=settings.CHROMA_DB_PATH,
|
||
settings=ChromaSettings(
|
||
anonymized_telemetry=False,
|
||
allow_reset=False
|
||
)
|
||
)
|
||
|
||
# Get or create collection
|
||
self.collection = self.chroma_client.get_or_create_collection(
|
||
name=settings.CHROMA_COLLECTION_NAME,
|
||
metadata={"hnsw:space": "cosine"}
|
||
)
|
||
|
||
# Initialize embedding model
|
||
self.embed_model = OllamaEmbedding(
|
||
model_name=settings.OLLAMA_EMBEDDING_MODEL,
|
||
base_url=settings.OLLAMA_BASE_URL
|
||
)
|
||
|
||
# Create ChromaVectorStore
|
||
self.vector_store = ChromaVectorStore(chroma_collection=self.collection)
|
||
|
||
# Create storage context and index
|
||
storage_context = StorageContext.from_defaults(
|
||
vector_store=self.vector_store
|
||
)
|
||
|
||
# Load existing index or create new one
|
||
try:
|
||
self.index = VectorStoreIndex.from_vector_store(
|
||
vector_store=self.vector_store,
|
||
embed_model=self.embed_model
|
||
)
|
||
logger.info("Loaded existing vector store index")
|
||
except Exception:
|
||
# Create new index if none exists
|
||
from llama_index.core import Document
|
||
self.index = VectorStoreIndex.from_documents(
|
||
documents=[],
|
||
storage_context=storage_context,
|
||
embed_model=self.embed_model
|
||
)
|
||
logger.info("Created new vector store index")
|
||
|
||
logger.info(f"Vector store initialized: {settings.CHROMA_COLLECTION_NAME}")
|
||
return # Success, exit retry loop
|
||
|
||
except Exception as e:
|
||
error_msg = str(e).lower()
|
||
is_locked_error = "locked" in error_msg or "database is locked" in error_msg
|
||
|
||
# Only retry for PersistentClient mode with lock errors
|
||
if not settings.CHROMA_SERVER_HOST and is_locked_error and attempt < retry_count - 1:
|
||
logger.warning(f"ChromaDB database is locked (attempt {attempt + 1}/{retry_count}): {e}")
|
||
continue # Retry
|
||
else:
|
||
# Last attempt or non-lock error
|
||
logger.error(f"Failed to initialize vector store: {e}")
|
||
if not settings.CHROMA_SERVER_HOST and is_locked_error:
|
||
logger.error(
|
||
"ChromaDB database is locked. Possible causes:\n"
|
||
" 1. Another process is using the database\n"
|
||
" 2. Previous process did not close properly\n"
|
||
" 3. Stale lock files exist\n"
|
||
"Solution: Check for other running processes or manually remove lock files:\n"
|
||
f" rm -f {settings.CHROMA_DB_PATH}/chroma.sqlite3-journal\n"
|
||
f" rm -f {settings.CHROMA_DB_PATH}/chroma.sqlite3-wal\n"
|
||
f" rm -f {settings.CHROMA_DB_PATH}/chroma.sqlite3-shm"
|
||
)
|
||
elif settings.CHROMA_SERVER_HOST:
|
||
logger.error(
|
||
f"Failed to connect to ChromaDB server at {settings.CHROMA_SERVER_HOST}:{settings.CHROMA_SERVER_PORT}\n"
|
||
"Please ensure:\n"
|
||
" 1. ChromaDB server is running\n"
|
||
" 2. Server address and port are correct\n"
|
||
" 3. Network connectivity is available"
|
||
)
|
||
raise
|
||
|
||
def _sanitize_metadata(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""
|
||
Sanitize metadata to ensure all values are ChromaDB-compatible types.
|
||
ChromaDB only accepts: str, int, float, bool, None, or SparseVector.
|
||
|
||
Args:
|
||
metadata: Metadata dictionary that may contain unsupported types
|
||
|
||
Returns:
|
||
Sanitized metadata dictionary with all values converted to supported types
|
||
"""
|
||
sanitized = {}
|
||
for key, value in metadata.items():
|
||
if value is None:
|
||
sanitized[key] = None
|
||
elif isinstance(value, (str, int, float, bool)):
|
||
sanitized[key] = value
|
||
elif isinstance(value, (datetime, date)):
|
||
# Convert datetime/date to ISO format string
|
||
sanitized[key] = value.isoformat() if isinstance(value, datetime) else value.isoformat()
|
||
elif isinstance(value, (list, tuple)):
|
||
# Convert lists/tuples to comma-separated string
|
||
sanitized[key] = ", ".join(str(v) for v in value)
|
||
elif isinstance(value, dict):
|
||
# Convert dict to JSON string (or skip if too complex)
|
||
try:
|
||
import json
|
||
sanitized[key] = json.dumps(value)
|
||
except (TypeError, ValueError):
|
||
sanitized[key] = str(value)
|
||
else:
|
||
# Fallback: convert to string
|
||
sanitized[key] = str(value)
|
||
|
||
return sanitized
|
||
|
||
def get_existing_doc_ids(self) -> set:
|
||
"""
|
||
Get set of existing document IDs from ChromaDB
|
||
|
||
Returns:
|
||
Set of existing document IDs, including:
|
||
- Original document IDs (unique_doc_id from metadata['doc_id'])
|
||
- Chunk IDs (doc.id_ format: {unique_doc_id}_chunk_{index})
|
||
Format: {db_source}_{table_name}_{id} or {db_source}_{table_name}_{id}_chunk_{index}
|
||
"""
|
||
try:
|
||
# Get all IDs and metadata from ChromaDB collection
|
||
# Note: 'ids' is automatically returned, don't include it in the include parameter
|
||
results = self.collection.get(include=['metadatas'])
|
||
|
||
existing_ids = set()
|
||
|
||
# Add all ChromaDB IDs (these are the actual chunk/document IDs stored)
|
||
# IDs are always returned by get(), even if not in include parameter
|
||
existing_ids.update(results.get('ids', []))
|
||
|
||
# Also extract original doc_ids from metadata for backward compatibility
|
||
# This helps identify if a document (not just chunks) has been processed
|
||
metadatas = results.get('metadatas', [])
|
||
for metadata in metadatas:
|
||
if metadata:
|
||
# Use doc_id (unique_doc_id) if available, format: {db_source}_{table_name}_{id}
|
||
if 'doc_id' in metadata:
|
||
existing_ids.add(metadata['doc_id'])
|
||
# Fallback: construct unique_doc_id from original_doc_id for older documents
|
||
elif 'original_doc_id' in metadata:
|
||
original_doc_id = metadata['original_doc_id']
|
||
db_source = metadata.get('db_source', '')
|
||
db_table = metadata.get('db_table', '')
|
||
if db_source and db_table:
|
||
# Construct unique_doc_id: {db_source}_{table_name}_{id}
|
||
unique_doc_id = f"{db_source}_{db_table}_{original_doc_id}"
|
||
existing_ids.add(unique_doc_id)
|
||
elif db_source:
|
||
# Fallback to old format: {db_source}_{id}
|
||
unique_doc_id = f"{db_source}_{original_doc_id}"
|
||
existing_ids.add(unique_doc_id)
|
||
else:
|
||
# Last fallback: use original_doc_id as-is
|
||
existing_ids.add(original_doc_id)
|
||
|
||
return existing_ids
|
||
except Exception as e:
|
||
logger.warning(f"Error getting existing document IDs: {e}")
|
||
return set()
|
||
|
||
def document_exists(self, doc_id: str) -> bool:
|
||
"""
|
||
Check if a document exists in the vector store
|
||
|
||
Args:
|
||
doc_id: Document ID to check
|
||
|
||
Returns:
|
||
True if document exists, False otherwise
|
||
"""
|
||
try:
|
||
# Try to get the document from ChromaDB
|
||
results = self.collection.get(ids=[doc_id])
|
||
return len(results.get('ids', [])) > 0
|
||
except Exception as e:
|
||
logger.debug(f"Error checking document existence: {e}")
|
||
return False
|
||
|
||
def get_document_by_id(self, doc_id: str) -> List[Dict[str, Any]]:
|
||
"""
|
||
Get all chunks/content for a document by its doc_id
|
||
|
||
This method retrieves all chunks that belong to a document, even if the document
|
||
was split into multiple chunks. It searches by metadata['doc_id'] field.
|
||
|
||
Args:
|
||
doc_id: Document ID to retrieve
|
||
|
||
Returns:
|
||
List of dictionaries containing document chunks with:
|
||
- id: Chunk ID
|
||
- text: Chunk content
|
||
- metadata: Chunk metadata
|
||
- chunk_index: Chunk index (if available)
|
||
"""
|
||
try:
|
||
# Query ChromaDB by metadata filter
|
||
# ChromaDB supports filtering by metadata using where clause
|
||
# Simple format: {"metadata_field": value} is equivalent to {"metadata_field": {"$eq": value}}
|
||
# Try simple format first (more compatible)
|
||
try:
|
||
results = self.collection.get(
|
||
where={"doc_id": doc_id},
|
||
include=['documents', 'metadatas']
|
||
)
|
||
except Exception:
|
||
# Fallback to explicit $eq operator if simple format doesn't work
|
||
results = self.collection.get(
|
||
where={"doc_id": {"$eq": doc_id}},
|
||
include=['documents', 'metadatas']
|
||
)
|
||
|
||
if not results or not results.get('ids'):
|
||
# Also try searching by id_ directly (for non-chunked documents)
|
||
try:
|
||
direct_results = self.collection.get(
|
||
ids=[doc_id],
|
||
include=['documents', 'metadatas']
|
||
)
|
||
if direct_results and direct_results.get('ids'):
|
||
results = direct_results
|
||
except Exception:
|
||
pass
|
||
|
||
if not results or not results.get('ids'):
|
||
return []
|
||
|
||
# Combine results into list of dictionaries
|
||
documents = []
|
||
ids = results.get('ids', [])
|
||
texts = results.get('documents', [])
|
||
metadatas = results.get('metadatas', [])
|
||
|
||
for i, chunk_id in enumerate(ids):
|
||
doc_dict = {
|
||
'id': chunk_id,
|
||
'text': texts[i] if i < len(texts) else '',
|
||
'metadata': metadatas[i] if i < len(metadatas) else {}
|
||
}
|
||
# Add chunk_index if available
|
||
if 'chunk_index' in doc_dict['metadata']:
|
||
doc_dict['chunk_index'] = doc_dict['metadata']['chunk_index']
|
||
elif 'chunk_id' in doc_dict['metadata']:
|
||
# Try to extract chunk index from chunk_id
|
||
chunk_id_str = doc_dict['metadata']['chunk_id']
|
||
if '_chunk_' in chunk_id_str:
|
||
try:
|
||
doc_dict['chunk_index'] = int(chunk_id_str.rsplit('_chunk_', 1)[1])
|
||
except (ValueError, IndexError):
|
||
pass
|
||
|
||
documents.append(doc_dict)
|
||
|
||
# Sort by chunk_index if available
|
||
documents.sort(key=lambda x: x.get('chunk_index', 0))
|
||
|
||
logger.info(f"Retrieved {len(documents)} chunk(s) for document ID: {doc_id}")
|
||
return documents
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error retrieving document by ID {doc_id}: {e}")
|
||
return []
|
||
|
||
def add_documents(self, documents: list, skip_existing: bool = True):
|
||
"""
|
||
Add documents to the vector store
|
||
|
||
Args:
|
||
documents: List of Document objects from LlamaIndex (may be chunked documents)
|
||
skip_existing:
|
||
- If True (default): Skip documents that already exist. Uses collection.add()
|
||
which will skip if ID exists (for initial sync/restart).
|
||
- If False: Update existing documents. Uses collection.upsert() which will
|
||
update if ID exists, insert if not (for incremental sync).
|
||
"""
|
||
try:
|
||
if not documents:
|
||
logger.warning("No documents to add")
|
||
return
|
||
|
||
# Filter out existing documents if skip_existing is True
|
||
# Note: For chunked documents, doc.id_ format is {unique_doc_id}_chunk_{index}
|
||
# We check doc.id_ directly to see if the chunk/document already exists
|
||
new_documents = []
|
||
skipped_count = 0
|
||
|
||
if skip_existing:
|
||
for doc in documents:
|
||
# Get document id_ (for chunks, this is {unique_doc_id}_chunk_{index})
|
||
# For non-chunk documents, this is {unique_doc_id}
|
||
doc_id_attr = getattr(doc, 'id_', None)
|
||
|
||
if not doc_id_attr:
|
||
# Document has no id_, add it (will be processed)
|
||
new_documents.append(doc)
|
||
elif self.document_exists(doc_id_attr):
|
||
# This specific chunk/document already exists, skip it
|
||
skipped_count += 1
|
||
else:
|
||
# Document/chunk doesn't exist, add it for processing
|
||
new_documents.append(doc)
|
||
else:
|
||
new_documents = documents
|
||
|
||
if not new_documents:
|
||
logger.info(f"All {len(documents)} documents already exist, skipping")
|
||
return
|
||
|
||
# Optimized batch insertion: directly use ChromaDB API for better performance
|
||
# This bypasses LlamaIndex's insert_nodes which may process nodes individually
|
||
# We batch generate embeddings and directly add to ChromaDB collection
|
||
try:
|
||
# Prepare data for batch insertion
|
||
texts = []
|
||
ids = []
|
||
metadatas = []
|
||
|
||
for doc in new_documents:
|
||
doc_id_attr = getattr(doc, 'id_', None)
|
||
if not doc_id_attr:
|
||
logger.warning(f"Document has no id_, skipping")
|
||
continue
|
||
|
||
# Skip empty or whitespace-only text
|
||
doc_text = doc.text if doc.text else ""
|
||
doc_text = doc_text.strip()
|
||
if not doc_text:
|
||
logger.warning(f"Document {doc_id_attr} has empty text, skipping")
|
||
skipped_count += 1
|
||
continue
|
||
|
||
texts.append(doc_text)
|
||
ids.append(doc_id_attr)
|
||
# Prepare metadata for ChromaDB
|
||
metadata = doc.metadata.copy() if doc.metadata else {}
|
||
# Ensure doc_id is in metadata for retrieval
|
||
if 'doc_id' not in metadata:
|
||
# Extract doc_id from chunk_id if available
|
||
if 'chunk_id' in metadata:
|
||
chunk_id = metadata['chunk_id']
|
||
# Extract original doc_id by removing _chunk_{index} suffix
|
||
if '_chunk_' in chunk_id:
|
||
metadata['doc_id'] = chunk_id.rsplit('_chunk_', 1)[0]
|
||
else:
|
||
# Use the id_ itself as doc_id if no chunk_id
|
||
metadata['doc_id'] = doc_id_attr
|
||
|
||
# Sanitize metadata to ensure all values are ChromaDB-compatible
|
||
# This converts datetime objects, lists, etc. to strings
|
||
sanitized_metadata = self._sanitize_metadata(metadata)
|
||
metadatas.append(sanitized_metadata)
|
||
|
||
if not texts:
|
||
logger.warning("No valid documents to insert (all documents have empty text or no id_)")
|
||
return
|
||
|
||
# Batch generate embeddings and write to ChromaDB incrementally
|
||
# Split into smaller batches to avoid overwhelming Ollama and show progress
|
||
# Write each batch to ChromaDB immediately after processing to reduce memory usage
|
||
total_count = len(texts)
|
||
batch_size = 100 # Process 100 documents per batch
|
||
logger.info(f"[Embedding进度] 开始批量生成 embeddings 并增量写入: {total_count} 个文档 (每批 {batch_size} 个)")
|
||
start_time = time.time()
|
||
|
||
total_added = 0 # Track total documents successfully added
|
||
|
||
try:
|
||
# Process in batches and write immediately
|
||
num_batches = (total_count + batch_size - 1) // batch_size
|
||
for batch_idx in range(num_batches):
|
||
batch_start = batch_idx * batch_size
|
||
batch_end = min(batch_start + batch_size, total_count)
|
||
batch_texts = texts[batch_start:batch_end]
|
||
batch_ids = ids[batch_start:batch_end]
|
||
batch_metadatas = metadatas[batch_start:batch_end]
|
||
|
||
batch_start_time = time.time()
|
||
logger.info(f"[Embedding进度] 处理批次 [{batch_idx+1}/{num_batches}]: {len(batch_texts)} 个文档...")
|
||
|
||
# Generate embeddings for this batch
|
||
# Catch IndexError and other exceptions from Ollama embedding generation
|
||
try:
|
||
batch_embeddings = self.embed_model.get_text_embedding_batch(batch_texts)
|
||
except (IndexError, ValueError, Exception) as emb_error:
|
||
# Ollama may return empty embeddings or raise IndexError
|
||
logger.warning(f"批量生成 embeddings 失败 (批次 {batch_idx+1}): {emb_error}, 切换到单个生成模式")
|
||
# Fallback to individual generation for this batch
|
||
batch_embeddings = []
|
||
for text in batch_texts:
|
||
try:
|
||
emb = self.embed_model.get_text_embedding(text)
|
||
if emb and len(emb) > 0:
|
||
batch_embeddings.append(emb)
|
||
else:
|
||
batch_embeddings.append(None)
|
||
except Exception as e:
|
||
logger.warning(f"单个 embedding 生成失败: {e}")
|
||
batch_embeddings.append(None)
|
||
|
||
embedding_elapsed = time.time() - batch_start_time
|
||
|
||
# Validate batch embeddings
|
||
if not batch_embeddings:
|
||
logger.warning(f"批次 {batch_idx+1} 返回空 embeddings 列表，跳过此批次")
|
||
skipped_count += len(batch_texts)
|
||
continue
|
||
|
||
# Filter out None embeddings and handle count mismatch
|
||
valid_batch_embeddings = []
|
||
valid_batch_texts = []
|
||
valid_batch_ids = []
|
||
valid_batch_metadatas = []
|
||
|
||
for i, emb in enumerate(batch_embeddings):
|
||
if emb is not None and len(emb) > 0:
|
||
valid_batch_embeddings.append(emb)
|
||
valid_batch_texts.append(batch_texts[i])
|
||
valid_batch_ids.append(batch_ids[i])
|
||
valid_batch_metadatas.append(batch_metadatas[i])
|
||
else:
|
||
logger.warning(f"批次 {batch_idx+1} 中文档 {batch_ids[i]} 的 embedding 为空，跳过")
|
||
skipped_count += 1
|
||
|
||
if not valid_batch_embeddings:
|
||
logger.warning(f"批次 {batch_idx+1} 没有有效的 embeddings，跳过")
|
||
continue
|
||
|
||
# Write this batch to ChromaDB immediately
|
||
write_start_time = time.time()
|
||
try:
|
||
# Use upsert if skip_existing is False (for updates), otherwise use add
|
||
if skip_existing:
|
||
# Skip existing: use add() which will skip if ID exists
|
||
self.collection.add(
|
||
ids=valid_batch_ids,
|
||
embeddings=valid_batch_embeddings,
|
||
documents=valid_batch_texts,
|
||
metadatas=valid_batch_metadatas
|
||
)
|
||
else:
|
||
# Update existing: use upsert() which will update if ID exists, insert if not
|
||
self.collection.upsert(
|
||
ids=valid_batch_ids,
|
||
embeddings=valid_batch_embeddings,
|
||
documents=valid_batch_texts,
|
||
metadatas=valid_batch_metadatas
|
||
)
|
||
write_elapsed = time.time() - write_start_time
|
||
total_added += len(valid_batch_ids)
|
||
|
||
batch_elapsed = time.time() - batch_start_time
|
||
processed = batch_end
|
||
progress = processed / total_count * 100
|
||
total_elapsed = time.time() - start_time
|
||
avg_time = total_elapsed / processed * 1000 if processed > 0 else 0
|
||
remaining = (total_count - processed) * avg_time / 1000 if processed > 0 else 0
|
||
|
||
logger.info(
|
||
f"[Embedding进度] 批次 [{batch_idx+1}/{num_batches}] 完成并已写入: "
|
||
f"{len(valid_batch_ids)} 个文档 (生成: {embedding_elapsed:.1f}秒, "
|
||
f"写入: {write_elapsed:.1f}秒, 总计: {batch_elapsed:.1f}秒, "
|
||
f"总进度: {processed}/{total_count} ({progress:.1f}%), "
|
||
f"已用: {total_elapsed:.1f}秒, 预计剩余: {remaining:.1f}秒, "
|
||
f"累计写入: {total_added} 个)"
|
||
)
|
||
except Exception as write_error:
|
||
logger.error(f"批次 {batch_idx+1} 写入 ChromaDB 失败: {write_error}")
|
||
skipped_count += len(valid_batch_ids)
|
||
# Continue processing next batch even if this one failed
|
||
continue
|
||
|
||
total_elapsed = time.time() - start_time
|
||
logger.info(
|
||
f"[Embedding进度] ✓ 批量处理完成: 成功写入 {total_added}/{total_count} 个文档 "
|
||
f"(总耗时: {total_elapsed:.1f}秒, 平均: {total_elapsed/total_count*1000:.1f}ms/个, 跳过: {skipped_count}个)"
|
||
)
|
||
|
||
# Validate final result
|
||
if total_added == 0:
|
||
raise ValueError("No valid documents were added to ChromaDB (all embeddings were empty or invalid)")
|
||
|
||
# Return early since we've already written to ChromaDB
|
||
return
|
||
|
||
except Exception as e:
|
||
logger.error(f"批量生成 embeddings 失败: {e}")
|
||
# Fallback: try individual embedding generation with incremental writes
|
||
logger.info(f"[Embedding进度] 切换到单个生成模式并增量写入: {len(texts)} 个文档")
|
||
|
||
start_time = time.time()
|
||
last_log_time = start_time
|
||
log_interval = 5.0 # 每5秒打印一次进度
|
||
batch_for_write = [] # Accumulate documents for batch write
|
||
batch_size_for_write = 50 # Write every 50 documents
|
||
total_added = 0
|
||
|
||
for i, text in enumerate(texts):
|
||
try:
|
||
# Skip empty text
|
||
if not text or not text.strip():
|
||
logger.warning(f"Skipping empty text for document {ids[i]}")
|
||
skipped_count += 1
|
||
continue
|
||
|
||
# Generate embedding individually
|
||
emb_start = time.time()
|
||
emb = self.embed_model.get_text_embedding(text)
|
||
emb_elapsed = time.time() - emb_start
|
||
|
||
if emb and len(emb) > 0:
|
||
batch_for_write.append({
|
||
'id': ids[i],
|
||
'text': text,
|
||
'metadata': metadatas[i],
|
||
'embedding': emb
|
||
})
|
||
|
||
# Write batch when it reaches the size limit
|
||
if len(batch_for_write) >= batch_size_for_write:
|
||
try:
|
||
batch_ids = [item['id'] for item in batch_for_write]
|
||
batch_texts = [item['text'] for item in batch_for_write]
|
||
batch_metadatas = [item['metadata'] for item in batch_for_write]
|
||
batch_embeddings = [item['embedding'] for item in batch_for_write]
|
||
|
||
# Use upsert if skip_existing is False (for updates), otherwise use add
|
||
if skip_existing:
|
||
self.collection.add(
|
||
ids=batch_ids,
|
||
embeddings=batch_embeddings,
|
||
documents=batch_texts,
|
||
metadatas=batch_metadatas
|
||
)
|
||
else:
|
||
self.collection.upsert(
|
||
ids=batch_ids,
|
||
embeddings=batch_embeddings,
|
||
documents=batch_texts,
|
||
metadatas=batch_metadatas
|
||
)
|
||
total_added += len(batch_ids)
|
||
batch_for_write = [] # Clear batch
|
||
except Exception as write_error:
|
||
logger.error(f"写入 ChromaDB 失败: {write_error}")
|
||
batch_for_write = [] # Clear batch to continue
|
||
else:
|
||
logger.warning(f"Empty embedding for document {ids[i]}, skipping")
|
||
skipped_count += 1
|
||
|
||
# Print progress periodically
|
||
current_time = time.time()
|
||
if current_time - last_log_time >= log_interval:
|
||
progress = (i + 1) / total_count * 100
|
||
elapsed = current_time - start_time
|
||
avg_time = elapsed / (i + 1) * 1000 if (i + 1) > 0 else 0
|
||
remaining = (total_count - i - 1) * avg_time / 1000
|
||
logger.info(
|
||
f"[Embedding进度] 单个生成模式: {i+1}/{total_count} ({progress:.1f}%), "
|
||
f"已用: {elapsed:.1f}秒, 预计剩余: {remaining:.1f}秒, "
|
||
f"累计写入: {total_added} 个"
|
||
)
|
||
last_log_time = current_time
|
||
|
||
except Exception as emb_error:
|
||
logger.warning(f"Error generating embedding for document {ids[i]}: {emb_error}, skipping")
|
||
skipped_count += 1
|
||
continue
|
||
|
||
# Write remaining documents in batch
|
||
if batch_for_write:
|
||
try:
|
||
batch_ids = [item['id'] for item in batch_for_write]
|
||
batch_texts = [item['text'] for item in batch_for_write]
|
||
batch_metadatas = [item['metadata'] for item in batch_for_write]
|
||
batch_embeddings = [item['embedding'] for item in batch_for_write]
|
||
|
||
# Use upsert if skip_existing is False (for updates), otherwise use add
|
||
if skip_existing:
|
||
self.collection.add(
|
||
ids=batch_ids,
|
||
embeddings=batch_embeddings,
|
||
documents=batch_texts,
|
||
metadatas=batch_metadatas
|
||
)
|
||
else:
|
||
self.collection.upsert(
|
||
ids=batch_ids,
|
||
embeddings=batch_embeddings,
|
||
documents=batch_texts,
|
||
metadatas=batch_metadatas
|
||
)
|
||
total_added += len(batch_ids)
|
||
except Exception as write_error:
|
||
logger.error(f"写入剩余文档到 ChromaDB 失败: {write_error}")
|
||
|
||
total_elapsed = time.time() - start_time
|
||
if total_added > 0:
|
||
avg_time = total_elapsed/total_added*1000
|
||
logger.info(
|
||
f"[Embedding进度] ✓ 单个生成完成: 成功写入 {total_added}/{total_count} 个文档 "
|
||
f"(耗时: {total_elapsed:.1f}秒, 平均: {avg_time:.1f}ms/个, 跳过: {skipped_count}个)"
|
||
)
|
||
else:
|
||
logger.info(
|
||
f"[Embedding进度] ✓ 单个生成完成: 成功写入 {total_added}/{total_count} 个文档 "
|
||
f"(耗时: {total_elapsed:.1f}秒, 跳过: {skipped_count}个)"
|
||
)
|
||
|
||
if total_added == 0:
|
||
raise ValueError(f"Failed to add any valid documents to ChromaDB: {e}")
|
||
|
||
# Return early since we've already written to ChromaDB
|
||
return
|
||
|
||
# Note: We need to refresh the index to make the new documents searchable
|
||
# However, since we're directly adding to ChromaDB, the index should automatically reflect changes
|
||
# If needed, we can reload the index, but it's usually not necessary for ChromaDB
|
||
|
||
except AttributeError:
|
||
# If refresh_ref_docs doesn't exist, fall back to insert_nodes
|
||
logger.warning("refresh_ref_docs method not available, falling back to insert_nodes")
|
||
from llama_index.core.schema import TextNode
|
||
|
||
try:
|
||
# Convert documents to nodes for batch insertion
|
||
all_nodes = []
|
||
for doc in new_documents:
|
||
try:
|
||
# Use id_ as per LlamaIndex documentation
|
||
doc_id_attr = getattr(doc, 'id_', None)
|
||
node = TextNode(
|
||
text=doc.text,
|
||
node_id=doc_id_attr,
|
||
metadata=doc.metadata if doc.metadata else {}
|
||
)
|
||
all_nodes.append(node)
|
||
except Exception as e:
|
||
# Get doc_id from metadata or id_ attribute
|
||
doc_id = doc.metadata.get('doc_id') if doc.metadata else getattr(doc, 'id_', 'unknown')
|
||
logger.warning(f"Error converting document {doc_id} to node: {e}")
|
||
continue
|
||
|
||
if all_nodes:
|
||
self.index.insert_nodes(all_nodes)
|
||
logger.info(f"Added {len(all_nodes)} nodes (from {len(new_documents)} documents) to vector store (skipped {skipped_count} existing)")
|
||
else:
|
||
logger.warning("No nodes to insert after conversion")
|
||
except AttributeError:
|
||
# Final fallback: individual inserts
|
||
logger.warning("insert_nodes method not available, falling back to individual inserts")
|
||
total_added = 0
|
||
total_failed = 0
|
||
for doc in new_documents:
|
||
try:
|
||
self.index.insert(doc)
|
||
total_added += 1
|
||
except Exception as e:
|
||
total_failed += 1
|
||
# Get doc_id from metadata or id_ attribute
|
||
doc_id = doc.metadata.get('doc_id') if doc.metadata else getattr(doc, 'id_', 'unknown')
|
||
logger.warning(f"Error inserting document {doc_id}: {e}")
|
||
continue
|
||
|
||
if total_failed > 0:
|
||
logger.warning(f"Added {total_added}/{len(new_documents)} documents to vector store (failed {total_failed}, skipped {skipped_count} existing)")
|
||
else:
|
||
logger.info(f"Added {total_added} new documents to vector store (skipped {skipped_count} existing)")
|
||
except Exception as e:
|
||
logger.error(f"Error during batch insert/update: {e}")
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Error adding documents to vector store: {e}")
|
||
raise
|
||
|
||
def delete_documents(self, doc_ids: list):
|
||
"""
|
||
Delete documents from the vector store
|
||
|
||
Args:
|
||
doc_ids: List of document IDs to delete
|
||
"""
|
||
try:
|
||
if not doc_ids:
|
||
return
|
||
|
||
# Delete from ChromaDB collection
|
||
self.collection.delete(ids=doc_ids)
|
||
logger.info(f"Deleted {len(doc_ids)} documents from vector store")
|
||
except Exception as e:
|
||
logger.error(f"Error deleting documents from vector store: {e}")
|
||
raise
|
||
|
||
def delete_documents_by_source(self, target_db_source: str):
|
||
"""
|
||
Delete documents from the vector store by db_source
|
||
Args:
|
||
db_source: Source string to filter documents (e.g., 'ruoyi-gitlink_pms_product_requirement')
|
||
"""
|
||
try:
|
||
if not target_db_source:
|
||
return
|
||
|
||
# Delete from ChromaDB collection using where clause
|
||
self.collection.delete(where={"db_source": target_db_source})
|
||
logger.info(f"Deleted documents from vector store with db_source: {target_db_source}")
|
||
except Exception as e:
|
||
logger.error(f"Error deleting documents from vector store by source: {e}")
|
||
raise
|
||
|
||
def has_documents(self) -> bool:
|
||
"""
|
||
Check if the collection has any documents
|
||
|
||
Returns:
|
||
True if collection has documents, False otherwise
|
||
"""
|
||
try:
|
||
count = self.collection.count()
|
||
return count > 0
|
||
except Exception as e:
|
||
logger.warning(f"Error checking document count: {e}")
|
||
return False
|
||
|
||
def get_retriever(self, top_k: int = None, filters: dict = None):
|
||
"""
|
||
Get a retriever for querying the vector store
|
||
|
||
Args:
|
||
top_k: Number of documents to retrieve (defaults to settings.TOP_K)
|
||
filters: Metadata filters to apply before vector search
|
||
|
||
Returns:
|
||
VectorStoreRetriever instance
|
||
|
||
Raises:
|
||
RuntimeError: If vector store is not initialized
|
||
"""
|
||
if self.index is None:
|
||
raise RuntimeError("Vector store index not initialized")
|
||
|
||
if top_k is None:
|
||
top_k = settings.TOP_K
|
||
|
||
retriever = self.index.as_retriever(similarity_top_k=top_k)
|
||
return retriever
|
||
|
||
def reset(self):
|
||
"""Reset the vector store (delete all data)"""
|
||
try:
|
||
self.chroma_client.delete_collection(name=settings.CHROMA_COLLECTION_NAME)
|
||
self._initialize()
|
||
logger.info("Vector store reset successfully")
|
||
except Exception as e:
|
||
logger.error(f"Error resetting vector store: {e}")
|
||
raise
|
||
|
||
def get_specific_db_source_metadata(self, target_db_source: str) -> str:
|
||
"""
|
||
获取指定db_source值的metadata信息
|
||
|
||
Args:
|
||
target_db_source: 目标db_source值，如'ruoyi-gitlink_pms_product_requirement'
|
||
|
||
Returns:
|
||
指定db_source的metadata中的content_column
|
||
"""
|
||
try:
|
||
# 使用ChromaDB的where过滤条件
|
||
origin_res = self.collection.get(
|
||
where={"db_source": {"$eq": target_db_source}},
|
||
include=['metadatas']
|
||
)
|
||
|
||
metadatas = origin_res.get('metadatas', [])
|
||
|
||
if metadatas:
|
||
# 获取 metadata中存储的 content_column 字段
|
||
content_column = metadatas[0].get('content_column', '')
|
||
logger.info(f"成功获取db_source({target_db_source}) metadata中content_column({content_column})")
|
||
return content_column
|
||
else:
|
||
return ''
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取db_source({target_db_source}) metadata中content_column失败: {e}")
|
||
return ''
|
||
|
||
def _preprocess_text(self, text: str) -> List[str]:
|
||
"""
|
||
预处理文本用于BM25搜索
|
||
|
||
Args:
|
||
text: 原始文本
|
||
|
||
Returns:
|
||
分词后的文本列表
|
||
"""
|
||
text = text.lower() # 转换为小写
|
||
text = text.translate(str.maketrans('', '', string.punctuation)) # 移除标点符号
|
||
tokens = re.findall(r'\b\w+\b', text) # 分词
|
||
return tokens
|
||
|
||
def keyword_search(self, query: str, top_k: int = 5, filters: dict = None) -> List[Tuple[str, float, Dict[str, Any]]]:
|
||
"""
|
||
使用BM25算法进行关键词搜索
|
||
|
||
Args:
|
||
query: 搜索查询
|
||
top_k: 返回结果数量
|
||
filters: metadata过滤条件
|
||
|
||
Returns:
|
||
排序后的结果列表，每个元素包含(doc_id, score, metadata)
|
||
"""
|
||
try:
|
||
# 获取文档
|
||
results = self.collection.get(
|
||
include=['documents', 'metadatas']
|
||
)
|
||
|
||
documents = results.get('documents', [])
|
||
metadatas = results.get('metadatas', [])
|
||
ids = results.get('ids', [])
|
||
|
||
if not documents:
|
||
return []
|
||
|
||
# 应用过滤条件
|
||
filtered_documents = []
|
||
filtered_metadatas = []
|
||
filtered_ids = []
|
||
|
||
for doc, meta, doc_id in zip(documents, metadatas, ids):
|
||
if not filters:
|
||
# 没有过滤条件，直接添加
|
||
filtered_documents.append(doc)
|
||
filtered_metadatas.append(meta)
|
||
filtered_ids.append(doc_id)
|
||
else:
|
||
# 应用过滤条件
|
||
match = True
|
||
for key, value in filters.items():
|
||
if key not in meta:
|
||
match = False
|
||
break
|
||
meta_value = meta[key]
|
||
if isinstance(meta_value, str) and isinstance(value, str):
|
||
# 对于字符串类型，使用大小写不敏感的模糊匹配
|
||
if value.lower() not in meta_value.lower():
|
||
match = False
|
||
break
|
||
else:
|
||
# 对于其他类型，使用精确匹配
|
||
if meta_value != value:
|
||
match = False
|
||
break
|
||
if match:
|
||
filtered_documents.append(doc)
|
||
filtered_metadatas.append(meta)
|
||
filtered_ids.append(doc_id)
|
||
|
||
if not filtered_documents:
|
||
return []
|
||
|
||
tokenized_docs = [self._preprocess_text(doc) for doc in filtered_documents] # 预处理文档
|
||
bm25 = BM25Okapi(tokenized_docs) # 初始化BM25
|
||
tokenized_query = self._preprocess_text(query) # 预处理查询
|
||
scores = bm25.get_scores(tokenized_query) # 计算BM25得分
|
||
sorted_indices = np.argsort(scores)[::-1][:top_k] # 排序并获取top_k结果
|
||
|
||
# 构建结果列表
|
||
search_results = []
|
||
for idx in sorted_indices:
|
||
if scores[idx] > 0: # 只返回得分大于0的结果
|
||
doc_id = filtered_ids[idx]
|
||
score = float(scores[idx])
|
||
metadata = filtered_metadatas[idx]
|
||
|
||
# 由于已经在获取文档时应用了过滤条件，这里不需要再次应用
|
||
search_results.append((doc_id, score, metadata))
|
||
return search_results
|
||
|
||
except Exception as e:
|
||
logger.error(f"关键词搜索失败: {e}")
|
||
return []
|
||
|
||
def hybrid_search(self, query: str, top_k: int = 5, vector_weight: float = 0.6, keyword_weight: float = 0.4, filters: dict = None) -> List[Tuple[str, float, Dict[str, Any]]]:
|
||
"""
|
||
融合检索：结合向量搜索和关键词搜索
|
||
|
||
Args:
|
||
query: 搜索查询
|
||
top_k: 返回结果数量
|
||
vector_weight: 向量搜索权重
|
||
keyword_weight: 关键词搜索权重
|
||
filters: metadata过滤条件
|
||
|
||
Returns:
|
||
排序后的结果列表，每个元素包含(doc_id, score, metadata)
|
||
"""
|
||
try:
|
||
# 1. 执行向量搜索
|
||
# 获取更多结果以确保有足够的候选
|
||
vector_retriever = self.get_retriever(top_k=top_k * 4) # 获取更多结果
|
||
vector_nodes = vector_retriever.retrieve(query)
|
||
|
||
vector_results = {}
|
||
for node in vector_nodes:
|
||
if hasattr(node, 'id_'):
|
||
doc_id = node.id_
|
||
elif hasattr(node, 'node_id'):
|
||
doc_id = node.node_id
|
||
else:
|
||
continue
|
||
|
||
vector_results[doc_id] = {
|
||
'score': node.score if hasattr(node, 'score') else 0.5,
|
||
'metadata': node.metadata if hasattr(node, 'metadata') else {},
|
||
'text': node.text if hasattr(node, 'text') else ''
|
||
}
|
||
|
||
# 2. 执行关键词搜索
|
||
keyword_results = self.keyword_search(query, top_k=top_k * 4, filters=filters)
|
||
|
||
keyword_scores = {}
|
||
keyword_metadata = {}
|
||
for doc_id, score, metadata in keyword_results:
|
||
keyword_scores[doc_id] = score
|
||
keyword_metadata[doc_id] = metadata
|
||
|
||
# 3. 归一化得分
|
||
# 归一化向量得分
|
||
if vector_results:
|
||
vector_scores = list(vector_results.values())
|
||
vector_min = min(item['score'] for item in vector_scores)
|
||
vector_max = max(item['score'] for item in vector_scores)
|
||
vector_range = vector_max - vector_min if vector_max > vector_min else 1
|
||
|
||
for doc_id in vector_results:
|
||
vector_results[doc_id]['normalized_score'] = (vector_results[doc_id]['score'] - vector_min) / vector_range
|
||
|
||
# 归一化关键词得分
|
||
if keyword_scores:
|
||
keyword_min = min(keyword_scores.values())
|
||
keyword_max = max(keyword_scores.values())
|
||
keyword_range = keyword_max - keyword_min if keyword_max > keyword_min else 1
|
||
|
||
for doc_id in keyword_scores:
|
||
keyword_scores[doc_id] = (keyword_scores[doc_id] - keyword_min) / keyword_range
|
||
|
||
# 4. 融合得分
|
||
hybrid_results = {}
|
||
|
||
# 合并向量搜索结果
|
||
for doc_id, info in vector_results.items():
|
||
# 应用过滤条件
|
||
if filters:
|
||
metadata = info['metadata']
|
||
match = True
|
||
for key, value in filters.items():
|
||
if key not in metadata:
|
||
match = False
|
||
break
|
||
meta_value = metadata[key]
|
||
if isinstance(meta_value, str) and isinstance(value, str):
|
||
# 对于字符串类型，使用大小写不敏感的模糊匹配
|
||
if value.lower() not in meta_value.lower():
|
||
match = False
|
||
break
|
||
else:
|
||
# 对于其他类型，使用精确匹配
|
||
if meta_value != value:
|
||
match = False
|
||
break
|
||
if not match:
|
||
continue
|
||
|
||
vector_score = info.get('normalized_score', 0)
|
||
keyword_score = keyword_scores.get(doc_id, 0)
|
||
|
||
# 计算融合得分
|
||
hybrid_score = vector_weight * vector_score + keyword_weight * keyword_score
|
||
|
||
hybrid_results[doc_id] = {
|
||
'score': hybrid_score,
|
||
'metadata': info['metadata'],
|
||
'text': info['text']
|
||
}
|
||
|
||
# 合并关键词搜索结果（不在向量搜索结果中的）
|
||
for doc_id, score in keyword_scores.items():
|
||
if doc_id not in hybrid_results:
|
||
# 应用过滤条件
|
||
if filters:
|
||
metadata = keyword_metadata.get(doc_id, {})
|
||
match = True
|
||
for key, value in filters.items():
|
||
if key not in metadata:
|
||
match = False
|
||
break
|
||
meta_value = metadata[key]
|
||
if isinstance(meta_value, str) and isinstance(value, str):
|
||
# 对于字符串类型，使用大小写不敏感的模糊匹配
|
||
if value.lower() not in meta_value.lower():
|
||
match = False
|
||
break
|
||
else:
|
||
# 对于其他类型，使用精确匹配
|
||
if meta_value != value:
|
||
match = False
|
||
break
|
||
if not match:
|
||
continue
|
||
|
||
hybrid_score = keyword_weight * score
|
||
hybrid_results[doc_id] = {
|
||
'score': hybrid_score,
|
||
'metadata': keyword_metadata.get(doc_id, {}),
|
||
'text': ''
|
||
}
|
||
|
||
# 5. 排序并获取top_k结果
|
||
sorted_results = sorted(
|
||
hybrid_results.items(),
|
||
key=lambda x: x[1]['score'],
|
||
reverse=True
|
||
)[:top_k]
|
||
|
||
# 6. 构建最终结果
|
||
final_results = []
|
||
for doc_id, info in sorted_results:
|
||
metadata = info['metadata']
|
||
final_results.append((doc_id, info['score'], metadata))
|
||
|
||
return final_results
|
||
|
||
except Exception as e:
|
||
logger.error(f"融合检索失败: {e}")
|
||
# 失败时回退到向量搜索
|
||
vector_retriever = self.get_retriever(top_k=top_k)
|
||
vector_nodes = vector_retriever.retrieve(query)
|
||
|
||
fallback_results = []
|
||
for node in vector_nodes:
|
||
if hasattr(node, 'id_'):
|
||
doc_id = node.id_
|
||
elif hasattr(node, 'node_id'):
|
||
doc_id = node.node_id
|
||
else:
|
||
continue
|
||
|
||
# 应用过滤条件
|
||
if filters:
|
||
metadata = node.metadata if hasattr(node, 'metadata') else {}
|
||
match = True
|
||
for key, value in filters.items():
|
||
if key not in metadata:
|
||
match = False
|
||
break
|
||
meta_value = metadata[key]
|
||
if isinstance(meta_value, str) and isinstance(value, str):
|
||
# 对于字符串类型，使用大小写不敏感的模糊匹配
|
||
if value.lower() not in meta_value.lower():
|
||
match = False
|
||
break
|
||
else:
|
||
# 对于其他类型，使用精确匹配
|
||
if meta_value != value:
|
||
match = False
|
||
break
|
||
if not match:
|
||
continue
|
||
|
||
score = node.score if hasattr(node, 'score') else 0.5
|
||
metadata = node.metadata if hasattr(node, 'metadata') else {}
|
||
|
||
fallback_results.append((doc_id, score, metadata))
|
||
|
||
return fallback_results
|
||
|
||
async def ahybrid_search(self, query: str, top_k: int = 5, vector_weight: float = 0.6, keyword_weight: float = 0.4, filters: dict = None) -> List[Tuple[str, float, Dict[str, Any]]]:
|
||
"""
|
||
异步融合检索：结合向量搜索和关键词搜索
|
||
|
||
Args:
|
||
query: 搜索查询
|
||
top_k: 返回结果数量
|
||
vector_weight: 向量搜索权重
|
||
keyword_weight: 关键词搜索权重
|
||
filters: metadata过滤条件
|
||
|
||
Returns:
|
||
排序后的结果列表，每个元素包含(doc_id, score, metadata)
|
||
"""
|
||
# 由于BM25搜索是CPU密集型的，这里使用同步方法
|
||
# 在实际生产环境中，可以使用线程池来异步执行
|
||
return self.hybrid_search(query, top_k, vector_weight, keyword_weight, filters)
|