880 lines
44 KiB
Python
880 lines
44 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
|
||
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
|
||
|
||
|
||
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
|
||
logger.info(
|
||
f"[Embedding进度] ✓ 单个生成完成: 成功写入 {total_added}/{total_count} 个文档 "
|
||
f"(耗时: {total_elapsed:.1f}秒, 平均: {total_elapsed/total_added*1000:.1f}ms/个, 跳过: {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):
|
||
"""
|
||
Get a retriever for querying the vector store
|
||
|
||
Args:
|
||
top_k: Number of documents to retrieve (defaults to settings.TOP_K)
|
||
|
||
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 ''
|