mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-09-15 04:00:33 +00:00
Merge pull request #15 from Hawksight-AI/feat/vector-store
feat: Implement vector_store module - comprehensive vector database i…
This commit is contained in:
@@ -8,9 +8,115 @@ Exports:
|
||||
- VectorIndexer: Vector indexing and search
|
||||
- VectorRetriever: Vector retrieval and similarity search
|
||||
- VectorManager: Vector store management and operations
|
||||
- FAISSAdapter: FAISS integration
|
||||
- PineconeAdapter: Pinecone integration
|
||||
- WeaviateAdapter: Weaviate integration
|
||||
- QdrantAdapter: Qdrant integration
|
||||
- MilvusAdapter: Milvus integration
|
||||
- HybridSearch: Hybrid vector and metadata search
|
||||
- MetadataStore: Metadata indexing and management
|
||||
- NamespaceManager: Namespace isolation and management
|
||||
"""
|
||||
|
||||
# from .vector_store import VectorStore
|
||||
# from .vector_indexer import VectorIndexer
|
||||
# from .vector_retriever import VectorRetriever
|
||||
# from .vector_manager import VectorManager
|
||||
from .vector_store import (
|
||||
VectorStore,
|
||||
VectorIndexer,
|
||||
VectorRetriever,
|
||||
VectorManager
|
||||
)
|
||||
|
||||
from .faiss_adapter import (
|
||||
FAISSAdapter,
|
||||
FAISSIndex,
|
||||
FAISSSearch,
|
||||
FAISSIndexBuilder
|
||||
)
|
||||
|
||||
from .pinecone_adapter import (
|
||||
PineconeAdapter,
|
||||
PineconeIndex,
|
||||
PineconeQuery,
|
||||
PineconeMetadata
|
||||
)
|
||||
|
||||
from .weaviate_adapter import (
|
||||
WeaviateAdapter,
|
||||
WeaviateClient,
|
||||
WeaviateSchema,
|
||||
WeaviateQuery
|
||||
)
|
||||
|
||||
from .qdrant_adapter import (
|
||||
QdrantAdapter,
|
||||
QdrantClient,
|
||||
QdrantCollection,
|
||||
QdrantSearch
|
||||
)
|
||||
|
||||
from .milvus_adapter import (
|
||||
MilvusAdapter,
|
||||
MilvusClient,
|
||||
MilvusCollection,
|
||||
MilvusSearch
|
||||
)
|
||||
|
||||
from .hybrid_search import (
|
||||
HybridSearch,
|
||||
MetadataFilter,
|
||||
SearchRanker
|
||||
)
|
||||
|
||||
from .metadata_store import (
|
||||
MetadataStore,
|
||||
MetadataIndex,
|
||||
MetadataSchema
|
||||
)
|
||||
|
||||
from .namespace_manager import (
|
||||
NamespaceManager,
|
||||
Namespace
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Core vector store
|
||||
"VectorStore",
|
||||
"VectorIndexer",
|
||||
"VectorRetriever",
|
||||
"VectorManager",
|
||||
# FAISS
|
||||
"FAISSAdapter",
|
||||
"FAISSIndex",
|
||||
"FAISSSearch",
|
||||
"FAISSIndexBuilder",
|
||||
# Pinecone
|
||||
"PineconeAdapter",
|
||||
"PineconeIndex",
|
||||
"PineconeQuery",
|
||||
"PineconeMetadata",
|
||||
# Weaviate
|
||||
"WeaviateAdapter",
|
||||
"WeaviateClient",
|
||||
"WeaviateSchema",
|
||||
"WeaviateQuery",
|
||||
# Qdrant
|
||||
"QdrantAdapter",
|
||||
"QdrantClient",
|
||||
"QdrantCollection",
|
||||
"QdrantSearch",
|
||||
# Milvus
|
||||
"MilvusAdapter",
|
||||
"MilvusClient",
|
||||
"MilvusCollection",
|
||||
"MilvusSearch",
|
||||
# Hybrid search
|
||||
"HybridSearch",
|
||||
"MetadataFilter",
|
||||
"SearchRanker",
|
||||
# Metadata store
|
||||
"MetadataStore",
|
||||
"MetadataIndex",
|
||||
"MetadataSchema",
|
||||
# Namespace manager
|
||||
"NamespaceManager",
|
||||
"Namespace",
|
||||
]
|
||||
|
||||
@@ -5,10 +5,369 @@ This module provides FAISS integration for vector storage
|
||||
and similarity search.
|
||||
"""
|
||||
|
||||
# TODO: Implement FAISS adapter
|
||||
# - FAISS index creation and management
|
||||
# - Vector storage and retrieval
|
||||
# - Similarity search and filtering
|
||||
# - Index optimization and training
|
||||
# - Performance optimization
|
||||
# - Error handling and recovery
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
# Optional FAISS import
|
||||
try:
|
||||
import faiss
|
||||
FAISS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FAISS_AVAILABLE = False
|
||||
faiss = None
|
||||
|
||||
|
||||
class FAISSIndex:
|
||||
"""FAISS index wrapper."""
|
||||
|
||||
def __init__(self, index: Any, dimension: int, index_type: str = "flat"):
|
||||
"""Initialize FAISS index wrapper."""
|
||||
self.index = index
|
||||
self.dimension = dimension
|
||||
self.index_type = index_type
|
||||
self.vector_ids: List[str] = []
|
||||
self.metadata: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
def add_vectors(self, vectors: np.ndarray, ids: Optional[List[str]] = None):
|
||||
"""Add vectors to index."""
|
||||
if ids is None:
|
||||
ids = [f"vec_{i}" for i in range(len(vectors))]
|
||||
|
||||
self.index.add(vectors.astype(np.float32))
|
||||
self.vector_ids.extend(ids)
|
||||
|
||||
def search(self, query_vectors: np.ndarray, k: int = 10) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""Search for similar vectors."""
|
||||
return self.index.search(query_vectors.astype(np.float32), k)
|
||||
|
||||
def get_vector(self, vector_id: str) -> Optional[np.ndarray]:
|
||||
"""Get vector by ID."""
|
||||
if vector_id not in self.vector_ids:
|
||||
return None
|
||||
|
||||
idx = self.vector_ids.index(vector_id)
|
||||
# Note: FAISS doesn't directly support retrieval by index in all cases
|
||||
# This is a simplified approach
|
||||
return None
|
||||
|
||||
def save(self, path: Union[str, Path]):
|
||||
"""Save index to disk."""
|
||||
if FAISS_AVAILABLE:
|
||||
faiss.write_index(self.index, str(path))
|
||||
else:
|
||||
raise ProcessingError("FAISS not available")
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: Union[str, Path], dimension: int, index_type: str = "flat"):
|
||||
"""Load index from disk."""
|
||||
if not FAISS_AVAILABLE:
|
||||
raise ProcessingError("FAISS not available")
|
||||
|
||||
index = faiss.read_index(str(path))
|
||||
return cls(index, dimension, index_type)
|
||||
|
||||
|
||||
class FAISSSearch:
|
||||
"""FAISS search operations."""
|
||||
|
||||
def __init__(self, index: FAISSIndex):
|
||||
"""Initialize FAISS search."""
|
||||
self.index = index
|
||||
self.logger = get_logger("faiss_search")
|
||||
|
||||
def search_similar(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
k: int = 10,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
k: Number of results
|
||||
**options: Search options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
if query_vector.ndim == 1:
|
||||
query_vector = query_vector.reshape(1, -1)
|
||||
|
||||
distances, indices = self.index.search(query_vector, k)
|
||||
|
||||
results = []
|
||||
for i, (dist, idx) in enumerate(zip(distances[0], indices[0])):
|
||||
if idx < len(self.index.vector_ids):
|
||||
vector_id = self.index.vector_ids[idx]
|
||||
results.append({
|
||||
"id": vector_id,
|
||||
"score": float(dist),
|
||||
"distance": float(dist),
|
||||
"metadata": self.index.metadata.get(vector_id, {})
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class FAISSIndexBuilder:
|
||||
"""FAISS index builder."""
|
||||
|
||||
def __init__(self, dimension: int = 768):
|
||||
"""Initialize FAISS index builder."""
|
||||
self.dimension = dimension
|
||||
self.logger = get_logger("faiss_builder")
|
||||
|
||||
def build_index(
|
||||
self,
|
||||
index_type: str = "flat",
|
||||
metric: str = "L2",
|
||||
**options
|
||||
) -> FAISSIndex:
|
||||
"""
|
||||
Build FAISS index.
|
||||
|
||||
Args:
|
||||
index_type: Index type ("flat", "ivf", "hnsw", "pq")
|
||||
metric: Distance metric ("L2", "inner_product")
|
||||
**options: Index options
|
||||
|
||||
Returns:
|
||||
FAISSIndex instance
|
||||
"""
|
||||
if not FAISS_AVAILABLE:
|
||||
raise ProcessingError(
|
||||
"FAISS is not available. Install it with: pip install faiss-cpu or faiss-gpu"
|
||||
)
|
||||
|
||||
# Create index based on type
|
||||
if index_type == "flat":
|
||||
if metric == "L2":
|
||||
index = faiss.IndexFlatL2(self.dimension)
|
||||
else:
|
||||
index = faiss.IndexFlatIP(self.dimension)
|
||||
|
||||
elif index_type == "ivf":
|
||||
nlist = options.get("nlist", 100)
|
||||
quantizer = faiss.IndexFlatL2(self.dimension)
|
||||
index = faiss.IndexIVFFlat(quantizer, self.dimension, nlist)
|
||||
|
||||
elif index_type == "hnsw":
|
||||
M = options.get("M", 32)
|
||||
index = faiss.IndexHNSWFlat(self.dimension, M)
|
||||
|
||||
elif index_type == "pq":
|
||||
m = options.get("m", 8) # Number of subquantizers
|
||||
bits = options.get("bits", 8)
|
||||
index = faiss.IndexPQ(self.dimension, m, bits)
|
||||
|
||||
else:
|
||||
raise ValidationError(f"Unsupported index type: {index_type}")
|
||||
|
||||
return FAISSIndex(index, self.dimension, index_type)
|
||||
|
||||
def train_index(self, index: FAISSIndex, training_vectors: np.ndarray):
|
||||
"""Train index on sample vectors."""
|
||||
if not isinstance(index.index, faiss.IndexIVFFlat):
|
||||
return # Only IVF indices need training
|
||||
|
||||
index.index.train(training_vectors.astype(np.float32))
|
||||
|
||||
|
||||
class FAISSAdapter:
|
||||
"""
|
||||
FAISS adapter for vector storage and similarity search.
|
||||
|
||||
• FAISS index creation and management
|
||||
• Vector storage and retrieval
|
||||
• Similarity search and filtering
|
||||
• Index optimization and training
|
||||
• Performance optimization
|
||||
• Error handling and recovery
|
||||
"""
|
||||
|
||||
def __init__(self, dimension: int = 768, **config):
|
||||
"""Initialize FAISS adapter."""
|
||||
self.logger = get_logger("faiss_adapter")
|
||||
self.config = config
|
||||
self.dimension = dimension
|
||||
|
||||
self.index: Optional[FAISSIndex] = None
|
||||
self.index_builder = FAISSIndexBuilder(dimension)
|
||||
self.search_engine: Optional[FAISSSearch] = None
|
||||
|
||||
# Check FAISS availability
|
||||
if not FAISS_AVAILABLE:
|
||||
self.logger.warning(
|
||||
"FAISS not available. Install with: pip install faiss-cpu or faiss-gpu"
|
||||
)
|
||||
|
||||
def create_index(
|
||||
self,
|
||||
index_type: str = "flat",
|
||||
metric: str = "L2",
|
||||
**options
|
||||
) -> FAISSIndex:
|
||||
"""
|
||||
Create FAISS index.
|
||||
|
||||
Args:
|
||||
index_type: Index type ("flat", "ivf", "hnsw", "pq")
|
||||
metric: Distance metric ("L2", "inner_product")
|
||||
**options: Index options
|
||||
|
||||
Returns:
|
||||
FAISSIndex instance
|
||||
"""
|
||||
self.index = self.index_builder.build_index(index_type, metric, **options)
|
||||
self.search_engine = FAISSSearch(self.index)
|
||||
|
||||
self.logger.info(f"Created FAISS index: {index_type} with metric {metric}")
|
||||
return self.index
|
||||
|
||||
def add_vectors(
|
||||
self,
|
||||
vectors: Union[List[np.ndarray], np.ndarray],
|
||||
ids: Optional[List[str]] = None,
|
||||
metadata: Optional[List[Dict[str, Any]]] = None,
|
||||
**options
|
||||
) -> List[str]:
|
||||
"""
|
||||
Add vectors to index.
|
||||
|
||||
Args:
|
||||
vectors: List of vectors or numpy array
|
||||
ids: Vector IDs
|
||||
metadata: Vector metadata
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of vector IDs
|
||||
"""
|
||||
if self.index is None:
|
||||
self.create_index(**options)
|
||||
|
||||
# Convert to numpy array
|
||||
if isinstance(vectors, list):
|
||||
vectors = np.array(vectors)
|
||||
|
||||
vectors = vectors.astype(np.float32)
|
||||
|
||||
# Generate IDs if not provided
|
||||
if ids is None:
|
||||
ids = [f"vec_{len(self.index.vector_ids) + i}" for i in range(len(vectors))]
|
||||
|
||||
# Store metadata
|
||||
if metadata:
|
||||
for vec_id, meta in zip(ids, metadata):
|
||||
self.index.metadata[vec_id] = meta
|
||||
|
||||
# Add vectors to index
|
||||
self.index.add_vectors(vectors, ids)
|
||||
|
||||
self.logger.info(f"Added {len(vectors)} vectors to FAISS index")
|
||||
return ids
|
||||
|
||||
def search_similar(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
k: int = 10,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
k: Number of results
|
||||
**options: Search options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
if self.search_engine is None:
|
||||
raise ProcessingError("Index not initialized. Call create_index() first.")
|
||||
|
||||
return self.search_engine.search_similar(query_vector, k, **options)
|
||||
|
||||
def save_index(self, path: Union[str, Path], **options) -> bool:
|
||||
"""
|
||||
Save index to disk.
|
||||
|
||||
Args:
|
||||
path: Path to save index
|
||||
**options: Save options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if self.index is None:
|
||||
raise ProcessingError("No index to save")
|
||||
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.index.save(path)
|
||||
self.logger.info(f"Saved FAISS index to {path}")
|
||||
return True
|
||||
|
||||
def load_index(
|
||||
self,
|
||||
path: Union[str, Path],
|
||||
index_type: str = "flat",
|
||||
**options
|
||||
) -> FAISSIndex:
|
||||
"""
|
||||
Load index from disk.
|
||||
|
||||
Args:
|
||||
path: Path to index file
|
||||
index_type: Index type
|
||||
**options: Load options
|
||||
|
||||
Returns:
|
||||
FAISSIndex instance
|
||||
"""
|
||||
if not FAISS_AVAILABLE:
|
||||
raise ProcessingError("FAISS not available")
|
||||
|
||||
self.index = FAISSIndex.load(path, self.dimension, index_type)
|
||||
self.search_engine = FAISSSearch(self.index)
|
||||
|
||||
self.logger.info(f"Loaded FAISS index from {path}")
|
||||
return self.index
|
||||
|
||||
def optimize_index(self, **options) -> bool:
|
||||
"""
|
||||
Optimize index for better performance.
|
||||
|
||||
Args:
|
||||
**options: Optimization options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if self.index is None:
|
||||
raise ProcessingError("No index to optimize")
|
||||
|
||||
# FAISS optimization is typically done during index creation
|
||||
# This method can be used for additional optimization
|
||||
self.logger.info("Index optimization completed")
|
||||
return True
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get index statistics."""
|
||||
if self.index is None:
|
||||
return {"status": "no_index"}
|
||||
|
||||
return {
|
||||
"index_type": self.index.index_type,
|
||||
"dimension": self.index.dimension,
|
||||
"vector_count": len(self.index.vector_ids),
|
||||
"faiss_available": FAISS_AVAILABLE
|
||||
}
|
||||
|
||||
@@ -5,10 +5,414 @@ This module provides combined vector and metadata search
|
||||
capabilities for enhanced retrieval.
|
||||
"""
|
||||
|
||||
# TODO: Implement hybrid search
|
||||
# - Vector similarity search
|
||||
# - Metadata filtering and querying
|
||||
# - Result fusion and ranking
|
||||
# - Performance optimization
|
||||
# - Error handling and recovery
|
||||
# - Advanced search strategies
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
import numpy as np
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
|
||||
class MetadataFilter:
|
||||
"""Metadata filter builder."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize metadata filter."""
|
||||
self.conditions: List[Dict[str, Any]] = []
|
||||
|
||||
def add_condition(
|
||||
self,
|
||||
field: str,
|
||||
operator: str,
|
||||
value: Any
|
||||
) -> "MetadataFilter":
|
||||
"""
|
||||
Add filter condition.
|
||||
|
||||
Args:
|
||||
field: Field name
|
||||
operator: Operator ("eq", "ne", "gt", "gte", "lt", "lte", "in", "contains")
|
||||
value: Value to filter by
|
||||
|
||||
Returns:
|
||||
Self for chaining
|
||||
"""
|
||||
self.conditions.append({
|
||||
"field": field,
|
||||
"operator": operator,
|
||||
"value": value
|
||||
})
|
||||
return self
|
||||
|
||||
def eq(self, field: str, value: Any) -> "MetadataFilter":
|
||||
"""Add equality condition."""
|
||||
return self.add_condition(field, "eq", value)
|
||||
|
||||
def ne(self, field: str, value: Any) -> "MetadataFilter":
|
||||
"""Add not-equal condition."""
|
||||
return self.add_condition(field, "ne", value)
|
||||
|
||||
def gt(self, field: str, value: Any) -> "MetadataFilter":
|
||||
"""Add greater-than condition."""
|
||||
return self.add_condition(field, "gt", value)
|
||||
|
||||
def gte(self, field: str, value: Any) -> "MetadataFilter":
|
||||
"""Add greater-than-or-equal condition."""
|
||||
return self.add_condition(field, "gte", value)
|
||||
|
||||
def lt(self, field: str, value: Any) -> "MetadataFilter":
|
||||
"""Add less-than condition."""
|
||||
return self.add_condition(field, "lt", value)
|
||||
|
||||
def lte(self, field: str, value: Any) -> "MetadataFilter":
|
||||
"""Add less-than-or-equal condition."""
|
||||
return self.add_condition(field, "lte", value)
|
||||
|
||||
def contains(self, field: str, value: Any) -> "MetadataFilter":
|
||||
"""Add contains condition."""
|
||||
return self.add_condition(field, "contains", value)
|
||||
|
||||
def in_list(self, field: str, values: List[Any]) -> "MetadataFilter":
|
||||
"""Add in-list condition."""
|
||||
return self.add_condition(field, "in", values)
|
||||
|
||||
def matches(self, metadata: Dict[str, Any]) -> bool:
|
||||
"""Check if metadata matches all conditions."""
|
||||
for condition in self.conditions:
|
||||
field = condition["field"]
|
||||
operator = condition["operator"]
|
||||
value = condition["value"]
|
||||
|
||||
if field not in metadata:
|
||||
return False
|
||||
|
||||
field_value = metadata[field]
|
||||
|
||||
if operator == "eq" and field_value != value:
|
||||
return False
|
||||
elif operator == "ne" and field_value == value:
|
||||
return False
|
||||
elif operator == "gt" and not (field_value > value):
|
||||
return False
|
||||
elif operator == "gte" and not (field_value >= value):
|
||||
return False
|
||||
elif operator == "lt" and not (field_value < value):
|
||||
return False
|
||||
elif operator == "lte" and not (field_value <= value):
|
||||
return False
|
||||
elif operator == "contains":
|
||||
if isinstance(field_value, str) and isinstance(value, str):
|
||||
if value not in field_value:
|
||||
return False
|
||||
elif isinstance(field_value, list):
|
||||
if value not in field_value:
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
elif operator == "in" and field_value not in value:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class SearchRanker:
|
||||
"""Search result ranker."""
|
||||
|
||||
def __init__(self, strategy: str = "reciprocal_rank_fusion"):
|
||||
"""Initialize search ranker."""
|
||||
self.strategy = strategy
|
||||
self.logger = get_logger("search_ranker")
|
||||
|
||||
def reciprocal_rank_fusion(
|
||||
self,
|
||||
results: List[List[Dict[str, Any]]],
|
||||
k: int = 60
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Reciprocal Rank Fusion (RRF) algorithm.
|
||||
|
||||
Args:
|
||||
results: List of result lists from different sources
|
||||
k: RRF constant
|
||||
|
||||
Returns:
|
||||
Fused and ranked results
|
||||
"""
|
||||
scores: Dict[str, float] = {}
|
||||
|
||||
for result_list in results:
|
||||
for rank, result in enumerate(result_list, start=1):
|
||||
result_id = result.get("id", str(id(result)))
|
||||
score = 1.0 / (k + rank)
|
||||
scores[result_id] = scores.get(result_id, 0.0) + score
|
||||
|
||||
# Sort by score
|
||||
ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Reconstruct results
|
||||
fused_results = []
|
||||
result_map = {}
|
||||
for result_list in results:
|
||||
for result in result_list:
|
||||
result_id = result.get("id", str(id(result)))
|
||||
result_map[result_id] = result
|
||||
|
||||
for result_id, score in ranked:
|
||||
if result_id in result_map:
|
||||
result = result_map[result_id].copy()
|
||||
result["score"] = score
|
||||
fused_results.append(result)
|
||||
|
||||
return fused_results
|
||||
|
||||
def weighted_average(
|
||||
self,
|
||||
results: List[List[Dict[str, Any]]],
|
||||
weights: List[float]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Weighted average fusion.
|
||||
|
||||
Args:
|
||||
results: List of result lists
|
||||
weights: Weights for each result list
|
||||
|
||||
Returns:
|
||||
Fused results
|
||||
"""
|
||||
if len(weights) != len(results):
|
||||
weights = [1.0 / len(results)] * len(results)
|
||||
|
||||
scores: Dict[str, Tuple[float, Dict[str, Any]]] = {}
|
||||
|
||||
for weight, result_list in zip(weights, results):
|
||||
for result in result_list:
|
||||
result_id = result.get("id", str(id(result)))
|
||||
score = result.get("score", 0.0) * weight
|
||||
|
||||
if result_id not in scores:
|
||||
scores[result_id] = (0.0, result)
|
||||
|
||||
scores[result_id] = (
|
||||
scores[result_id][0] + score,
|
||||
scores[result_id][1]
|
||||
)
|
||||
|
||||
# Sort by score
|
||||
ranked = sorted(scores.values(), key=lambda x: x[0], reverse=True)
|
||||
|
||||
fused_results = []
|
||||
for score, result in ranked:
|
||||
result_copy = result.copy()
|
||||
result_copy["score"] = score
|
||||
fused_results.append(result_copy)
|
||||
|
||||
return fused_results
|
||||
|
||||
def rank(
|
||||
self,
|
||||
results: List[List[Dict[str, Any]]],
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Rank and fuse results.
|
||||
|
||||
Args:
|
||||
results: List of result lists
|
||||
**options: Ranking options
|
||||
|
||||
Returns:
|
||||
Fused and ranked results
|
||||
"""
|
||||
if self.strategy == "reciprocal_rank_fusion":
|
||||
k = options.get("k", 60)
|
||||
return self.reciprocal_rank_fusion(results, k)
|
||||
elif self.strategy == "weighted_average":
|
||||
weights = options.get("weights", [1.0 / len(results)] * len(results))
|
||||
return self.weighted_average(results, weights)
|
||||
else:
|
||||
return self.reciprocal_rank_fusion(results)
|
||||
|
||||
|
||||
class HybridSearch:
|
||||
"""
|
||||
Hybrid search combining vector similarity and metadata filtering.
|
||||
|
||||
• Vector similarity search
|
||||
• Metadata filtering and querying
|
||||
• Result fusion and ranking
|
||||
• Performance optimization
|
||||
• Error handling and recovery
|
||||
• Advanced search strategies
|
||||
"""
|
||||
|
||||
def __init__(self, **config):
|
||||
"""Initialize hybrid search."""
|
||||
self.logger = get_logger("hybrid_search")
|
||||
self.config = config
|
||||
self.ranker = SearchRanker(config.get("ranking_strategy", "reciprocal_rank_fusion"))
|
||||
|
||||
def search(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
vectors: List[np.ndarray],
|
||||
metadata: List[Dict[str, Any]],
|
||||
vector_ids: List[str],
|
||||
k: int = 10,
|
||||
metadata_filter: Optional[MetadataFilter] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Perform hybrid search.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
vectors: List of vectors to search
|
||||
metadata: List of metadata dictionaries
|
||||
vector_ids: Vector IDs
|
||||
k: Number of results
|
||||
metadata_filter: Optional metadata filter
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
if not vectors or not metadata:
|
||||
return []
|
||||
|
||||
# Filter by metadata first
|
||||
if metadata_filter:
|
||||
filtered_indices = [
|
||||
i for i, meta in enumerate(metadata)
|
||||
if metadata_filter.matches(meta)
|
||||
]
|
||||
|
||||
filtered_vectors = [vectors[i] for i in filtered_indices]
|
||||
filtered_metadata = [metadata[i] for i in filtered_indices]
|
||||
filtered_ids = [vector_ids[i] for i in filtered_indices]
|
||||
else:
|
||||
filtered_vectors = vectors
|
||||
filtered_metadata = metadata
|
||||
filtered_ids = vector_ids
|
||||
|
||||
if not filtered_vectors:
|
||||
return []
|
||||
|
||||
# Perform vector similarity search
|
||||
vector_results = self._vector_search(
|
||||
query_vector,
|
||||
filtered_vectors,
|
||||
filtered_ids,
|
||||
k * 2, # Get more results for ranking
|
||||
**options
|
||||
)
|
||||
|
||||
# Add metadata to results
|
||||
for result in vector_results:
|
||||
result_id = result.get("id")
|
||||
if result_id in filtered_ids:
|
||||
idx = filtered_ids.index(result_id)
|
||||
result["metadata"] = filtered_metadata[idx]
|
||||
|
||||
# Rank and return top k
|
||||
return vector_results[:k]
|
||||
|
||||
def _vector_search(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
vectors: List[np.ndarray],
|
||||
vector_ids: List[str],
|
||||
k: int,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Perform vector similarity search."""
|
||||
if not vectors:
|
||||
return []
|
||||
|
||||
# Convert to numpy
|
||||
if isinstance(vectors[0], list):
|
||||
vectors = np.array(vectors)
|
||||
else:
|
||||
vectors = np.vstack(vectors)
|
||||
|
||||
if isinstance(query_vector, list):
|
||||
query_vector = np.array(query_vector)
|
||||
|
||||
# Calculate cosine similarity
|
||||
query_norm = np.linalg.norm(query_vector)
|
||||
vector_norms = np.linalg.norm(vectors, axis=1)
|
||||
|
||||
similarities = np.dot(vectors, query_vector) / (vector_norms * query_norm + 1e-8)
|
||||
|
||||
# Get top k
|
||||
top_indices = np.argsort(similarities)[::-1][:k]
|
||||
|
||||
results = []
|
||||
for idx in top_indices:
|
||||
results.append({
|
||||
"id": vector_ids[idx],
|
||||
"score": float(similarities[idx]),
|
||||
"distance": 1.0 - float(similarities[idx])
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
def multi_source_search(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
sources: List[Dict[str, Any]],
|
||||
k: int = 10,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search across multiple sources and fuse results.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
sources: List of source dictionaries with 'vectors', 'metadata', 'ids'
|
||||
k: Number of results
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
Fused search results
|
||||
"""
|
||||
all_results = []
|
||||
|
||||
for source in sources:
|
||||
source_results = self.search(
|
||||
query_vector,
|
||||
source.get("vectors", []),
|
||||
source.get("metadata", []),
|
||||
source.get("ids", []),
|
||||
k=k,
|
||||
metadata_filter=source.get("filter"),
|
||||
**options
|
||||
)
|
||||
all_results.append(source_results)
|
||||
|
||||
# Fuse results using ranker
|
||||
fused_results = self.ranker.rank(all_results, **options)
|
||||
|
||||
return fused_results[:k]
|
||||
|
||||
def filter_by_metadata(
|
||||
self,
|
||||
results: List[Dict[str, Any]],
|
||||
metadata_filter: MetadataFilter
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Filter results by metadata.
|
||||
|
||||
Args:
|
||||
results: Search results
|
||||
metadata_filter: Metadata filter
|
||||
|
||||
Returns:
|
||||
Filtered results
|
||||
"""
|
||||
filtered = []
|
||||
for result in results:
|
||||
metadata = result.get("metadata", {})
|
||||
if metadata_filter.matches(metadata):
|
||||
filtered.append(result)
|
||||
|
||||
return filtered
|
||||
|
||||
@@ -5,10 +5,387 @@ This module provides metadata indexing and management
|
||||
for vector store operations.
|
||||
"""
|
||||
|
||||
# TODO: Implement metadata storage
|
||||
# - Metadata indexing and storage
|
||||
# - Metadata querying and filtering
|
||||
# - Schema management and validation
|
||||
# - Performance optimization
|
||||
# - Error handling and recovery
|
||||
# - Multi-format metadata support
|
||||
from typing import Any, Dict, List, Optional, Set, Union
|
||||
from collections import defaultdict
|
||||
import json
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
|
||||
class MetadataIndex:
|
||||
"""Metadata index for fast lookups."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize metadata index."""
|
||||
self.field_indexes: Dict[str, Dict[Any, Set[str]]] = defaultdict(lambda: defaultdict(set))
|
||||
self.vector_metadata: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
def index_metadata(self, vector_id: str, metadata: Dict[str, Any]):
|
||||
"""Index metadata for a vector."""
|
||||
self.vector_metadata[vector_id] = metadata
|
||||
|
||||
for field, value in metadata.items():
|
||||
if isinstance(value, (str, int, float, bool)):
|
||||
self.field_indexes[field][value].add(vector_id)
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
if isinstance(item, (str, int, float, bool)):
|
||||
self.field_indexes[field][item].add(vector_id)
|
||||
|
||||
def remove_metadata(self, vector_id: str):
|
||||
"""Remove metadata from index."""
|
||||
if vector_id not in self.vector_metadata:
|
||||
return
|
||||
|
||||
metadata = self.vector_metadata[vector_id]
|
||||
|
||||
for field, value in metadata.items():
|
||||
if isinstance(value, (str, int, float, bool)):
|
||||
if vector_id in self.field_indexes[field].get(value, set()):
|
||||
self.field_indexes[field][value].remove(vector_id)
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
if isinstance(item, (str, int, float, bool)):
|
||||
if vector_id in self.field_indexes[field].get(item, set()):
|
||||
self.field_indexes[field][item].remove(vector_id)
|
||||
|
||||
del self.vector_metadata[vector_id]
|
||||
|
||||
def query(
|
||||
self,
|
||||
conditions: Dict[str, Any],
|
||||
operator: str = "AND"
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Query vectors by metadata conditions.
|
||||
|
||||
Args:
|
||||
conditions: Field-value conditions
|
||||
operator: "AND" or "OR"
|
||||
|
||||
Returns:
|
||||
Set of vector IDs matching conditions
|
||||
"""
|
||||
if not conditions:
|
||||
return set(self.vector_metadata.keys())
|
||||
|
||||
result_sets = []
|
||||
|
||||
for field, value in conditions.items():
|
||||
if field in self.field_indexes:
|
||||
if value in self.field_indexes[field]:
|
||||
result_sets.append(self.field_indexes[field][value])
|
||||
else:
|
||||
result_sets.append(set())
|
||||
else:
|
||||
result_sets.append(set())
|
||||
|
||||
if operator == "AND":
|
||||
result = set.intersection(*result_sets) if result_sets else set()
|
||||
else: # OR
|
||||
result = set.union(*result_sets) if result_sets else set()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class MetadataSchema:
|
||||
"""Metadata schema validator."""
|
||||
|
||||
def __init__(self, schema: Optional[Dict[str, Any]] = None):
|
||||
"""Initialize metadata schema."""
|
||||
self.schema = schema or {}
|
||||
self.logger = get_logger("metadata_schema")
|
||||
|
||||
def validate(self, metadata: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Validate metadata against schema.
|
||||
|
||||
Args:
|
||||
metadata: Metadata to validate
|
||||
|
||||
Returns:
|
||||
True if valid
|
||||
"""
|
||||
if not self.schema:
|
||||
return True
|
||||
|
||||
for field, field_schema in self.schema.items():
|
||||
if field_schema.get("required", False):
|
||||
if field not in metadata:
|
||||
raise ValidationError(f"Required field '{field}' is missing")
|
||||
|
||||
if field in metadata:
|
||||
value = metadata[field]
|
||||
field_type = field_schema.get("type")
|
||||
|
||||
if field_type and not isinstance(value, field_type):
|
||||
raise ValidationError(
|
||||
f"Field '{field}' must be of type {field_type}, got {type(value)}"
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
def add_field(
|
||||
self,
|
||||
field: str,
|
||||
field_type: type,
|
||||
required: bool = False,
|
||||
default: Any = None
|
||||
):
|
||||
"""Add field to schema."""
|
||||
self.schema[field] = {
|
||||
"type": field_type,
|
||||
"required": required,
|
||||
"default": default
|
||||
}
|
||||
|
||||
|
||||
class MetadataStore:
|
||||
"""
|
||||
Metadata store for vector store operations.
|
||||
|
||||
• Metadata indexing and storage
|
||||
• Metadata querying and filtering
|
||||
• Schema management and validation
|
||||
• Performance optimization
|
||||
• Error handling and recovery
|
||||
• Multi-format metadata support
|
||||
"""
|
||||
|
||||
def __init__(self, schema: Optional[Dict[str, Any]] = None, **config):
|
||||
"""Initialize metadata store."""
|
||||
self.logger = get_logger("metadata_store")
|
||||
self.config = config
|
||||
self.schema = MetadataSchema(schema)
|
||||
self.index = MetadataIndex()
|
||||
self.metadata: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
def store_metadata(
|
||||
self,
|
||||
vector_id: str,
|
||||
metadata: Dict[str, Any],
|
||||
**options
|
||||
) -> bool:
|
||||
"""
|
||||
Store metadata for a vector.
|
||||
|
||||
Args:
|
||||
vector_id: Vector ID
|
||||
metadata: Metadata dictionary
|
||||
**options: Storage options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
try:
|
||||
# Validate against schema
|
||||
self.schema.validate(metadata)
|
||||
|
||||
# Store metadata
|
||||
self.metadata[vector_id] = metadata
|
||||
self.index.index_metadata(vector_id, metadata)
|
||||
|
||||
self.logger.debug(f"Stored metadata for vector {vector_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to store metadata: {str(e)}")
|
||||
|
||||
def get_metadata(self, vector_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get metadata for a vector.
|
||||
|
||||
Args:
|
||||
vector_id: Vector ID
|
||||
|
||||
Returns:
|
||||
Metadata dictionary or None
|
||||
"""
|
||||
return self.metadata.get(vector_id)
|
||||
|
||||
def update_metadata(
|
||||
self,
|
||||
vector_id: str,
|
||||
metadata_updates: Dict[str, Any],
|
||||
**options
|
||||
) -> bool:
|
||||
"""
|
||||
Update metadata for a vector.
|
||||
|
||||
Args:
|
||||
vector_id: Vector ID
|
||||
metadata_updates: Metadata updates
|
||||
**options: Update options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if vector_id not in self.metadata:
|
||||
raise ProcessingError(f"Vector {vector_id} not found")
|
||||
|
||||
try:
|
||||
# Merge updates
|
||||
updated_metadata = {**self.metadata[vector_id], **metadata_updates}
|
||||
|
||||
# Validate
|
||||
self.schema.validate(updated_metadata)
|
||||
|
||||
# Remove old index entries
|
||||
self.index.remove_metadata(vector_id)
|
||||
|
||||
# Update
|
||||
self.metadata[vector_id] = updated_metadata
|
||||
self.index.index_metadata(vector_id, updated_metadata)
|
||||
|
||||
self.logger.debug(f"Updated metadata for vector {vector_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to update metadata: {str(e)}")
|
||||
|
||||
def delete_metadata(self, vector_id: str, **options) -> bool:
|
||||
"""
|
||||
Delete metadata for a vector.
|
||||
|
||||
Args:
|
||||
vector_id: Vector ID
|
||||
**options: Delete options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if vector_id in self.metadata:
|
||||
self.index.remove_metadata(vector_id)
|
||||
del self.metadata[vector_id]
|
||||
self.logger.debug(f"Deleted metadata for vector {vector_id}")
|
||||
|
||||
return True
|
||||
|
||||
def query_metadata(
|
||||
self,
|
||||
conditions: Dict[str, Any],
|
||||
operator: str = "AND",
|
||||
**options
|
||||
) -> List[str]:
|
||||
"""
|
||||
Query vectors by metadata.
|
||||
|
||||
Args:
|
||||
conditions: Field-value conditions
|
||||
operator: "AND" or "OR"
|
||||
**options: Query options
|
||||
|
||||
Returns:
|
||||
List of vector IDs matching conditions
|
||||
"""
|
||||
matching_ids = self.index.query(conditions, operator)
|
||||
return list(matching_ids)
|
||||
|
||||
def filter_metadata(
|
||||
self,
|
||||
vector_ids: List[str],
|
||||
conditions: Dict[str, Any],
|
||||
operator: str = "AND",
|
||||
**options
|
||||
) -> List[str]:
|
||||
"""
|
||||
Filter vector IDs by metadata conditions.
|
||||
|
||||
Args:
|
||||
vector_ids: List of vector IDs to filter
|
||||
conditions: Field-value conditions
|
||||
operator: "AND" or "OR"
|
||||
**options: Filter options
|
||||
|
||||
Returns:
|
||||
Filtered list of vector IDs
|
||||
"""
|
||||
matching_ids = self.query_metadata(conditions, operator, **options)
|
||||
return [vid for vid in vector_ids if vid in matching_ids]
|
||||
|
||||
def get_all_metadata(self, vector_ids: Optional[List[str]] = None) -> Dict[str, Dict[str, Any]]:
|
||||
"""
|
||||
Get all metadata.
|
||||
|
||||
Args:
|
||||
vector_ids: Optional list of vector IDs to retrieve
|
||||
|
||||
Returns:
|
||||
Dictionary mapping vector IDs to metadata
|
||||
"""
|
||||
if vector_ids:
|
||||
return {vid: self.metadata.get(vid, {}) for vid in vector_ids}
|
||||
else:
|
||||
return self.metadata.copy()
|
||||
|
||||
def get_field_values(self, field: str) -> List[Any]:
|
||||
"""
|
||||
Get all unique values for a field.
|
||||
|
||||
Args:
|
||||
field: Field name
|
||||
|
||||
Returns:
|
||||
List of unique values
|
||||
"""
|
||||
if field in self.index.field_indexes:
|
||||
return list(self.index.field_indexes[field].keys())
|
||||
return []
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get metadata store statistics."""
|
||||
return {
|
||||
"total_vectors": len(self.metadata),
|
||||
"indexed_fields": len(self.index.field_indexes),
|
||||
"field_counts": {
|
||||
field: len(values)
|
||||
for field, values in self.index.field_indexes.items()
|
||||
}
|
||||
}
|
||||
|
||||
def export_metadata(self, format: str = "json", **options) -> Union[str, Dict[str, Any]]:
|
||||
"""
|
||||
Export metadata.
|
||||
|
||||
Args:
|
||||
format: Export format ("json", "dict")
|
||||
**options: Export options
|
||||
|
||||
Returns:
|
||||
Exported metadata
|
||||
"""
|
||||
if format == "json":
|
||||
return json.dumps(self.metadata, indent=2, default=str)
|
||||
else:
|
||||
return self.metadata.copy()
|
||||
|
||||
def import_metadata(
|
||||
self,
|
||||
data: Union[str, Dict[str, Any]],
|
||||
format: str = "json",
|
||||
**options
|
||||
) -> bool:
|
||||
"""
|
||||
Import metadata.
|
||||
|
||||
Args:
|
||||
data: Metadata data
|
||||
format: Data format ("json", "dict")
|
||||
**options: Import options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if format == "json":
|
||||
if isinstance(data, str):
|
||||
data = json.loads(data)
|
||||
|
||||
if not isinstance(data, dict):
|
||||
raise ValidationError("Metadata must be a dictionary")
|
||||
|
||||
for vector_id, metadata in data.items():
|
||||
self.store_metadata(vector_id, metadata, **options)
|
||||
|
||||
self.logger.info(f"Imported metadata for {len(data)} vectors")
|
||||
return True
|
||||
|
||||
@@ -5,10 +5,444 @@ This module provides Milvus integration for vector storage
|
||||
and similarity search.
|
||||
"""
|
||||
|
||||
# TODO: Implement Milvus adapter
|
||||
# - Milvus connection and authentication
|
||||
# - Collection and partition management
|
||||
# - Vector storage and retrieval
|
||||
# - Similarity search and filtering
|
||||
# - Performance optimization
|
||||
# - Error handling and recovery
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
import numpy as np
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
# Optional Milvus import
|
||||
try:
|
||||
from pymilvus import (
|
||||
connections, Collection, FieldSchema, CollectionSchema, DataType,
|
||||
utility, MilvusException
|
||||
)
|
||||
MILVUS_AVAILABLE = True
|
||||
except ImportError:
|
||||
MILVUS_AVAILABLE = False
|
||||
connections = None
|
||||
Collection = None
|
||||
FieldSchema = None
|
||||
CollectionSchema = None
|
||||
DataType = None
|
||||
utility = None
|
||||
MilvusException = None
|
||||
|
||||
|
||||
class MilvusClient:
|
||||
"""Milvus client wrapper."""
|
||||
|
||||
def __init__(self, alias: str = "default"):
|
||||
"""Initialize Milvus client wrapper."""
|
||||
self.alias = alias
|
||||
self.logger = get_logger("milvus_client")
|
||||
|
||||
def connect(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 19530,
|
||||
user: Optional[str] = None,
|
||||
password: Optional[str] = None,
|
||||
**options
|
||||
) -> bool:
|
||||
"""Connect to Milvus server."""
|
||||
if not MILVUS_AVAILABLE:
|
||||
raise ProcessingError("Milvus not available")
|
||||
|
||||
try:
|
||||
connections.connect(
|
||||
alias=self.alias,
|
||||
host=host,
|
||||
port=port,
|
||||
user=user,
|
||||
password=password,
|
||||
**options
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to connect to Milvus: {str(e)}")
|
||||
|
||||
def disconnect(self):
|
||||
"""Disconnect from Milvus server."""
|
||||
if not MILVUS_AVAILABLE:
|
||||
return
|
||||
|
||||
try:
|
||||
connections.disconnect(self.alias)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to disconnect: {str(e)}")
|
||||
|
||||
|
||||
class MilvusCollection:
|
||||
"""Milvus collection wrapper."""
|
||||
|
||||
def __init__(self, collection: Any, collection_name: str):
|
||||
"""Initialize Milvus collection wrapper."""
|
||||
self.collection = collection
|
||||
self.collection_name = collection_name
|
||||
self.logger = get_logger("milvus_collection")
|
||||
|
||||
def insert(
|
||||
self,
|
||||
data: List[List[Any]],
|
||||
**options
|
||||
) -> Any:
|
||||
"""Insert data into collection."""
|
||||
if not MILVUS_AVAILABLE:
|
||||
raise ProcessingError("Milvus not available")
|
||||
|
||||
try:
|
||||
insert_result = self.collection.insert(data, **options)
|
||||
return insert_result
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to insert data: {str(e)}")
|
||||
|
||||
def search(
|
||||
self,
|
||||
vectors: List[np.ndarray],
|
||||
anns_field: str,
|
||||
param: Dict[str, Any],
|
||||
limit: int = 10,
|
||||
expr: Optional[str] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Search vectors in collection."""
|
||||
if not MILVUS_AVAILABLE:
|
||||
raise ProcessingError("Milvus not available")
|
||||
|
||||
try:
|
||||
search_results = self.collection.search(
|
||||
data=[v.tolist() for v in vectors],
|
||||
anns_field=anns_field,
|
||||
param=param,
|
||||
limit=limit,
|
||||
expr=expr,
|
||||
**options
|
||||
)
|
||||
|
||||
results = []
|
||||
for hits in search_results:
|
||||
batch_results = []
|
||||
for hit in hits:
|
||||
batch_results.append({
|
||||
"id": hit.id,
|
||||
"distance": hit.distance,
|
||||
"score": 1.0 - hit.distance if hit.distance <= 1.0 else 1.0 / (1.0 + hit.distance)
|
||||
})
|
||||
results.append(batch_results)
|
||||
|
||||
return results[0] if len(results) == 1 else results
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to search: {str(e)}")
|
||||
|
||||
def load(self, **options) -> bool:
|
||||
"""Load collection into memory."""
|
||||
if not MILVUS_AVAILABLE:
|
||||
raise ProcessingError("Milvus not available")
|
||||
|
||||
try:
|
||||
self.collection.load(**options)
|
||||
return True
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to load collection: {str(e)}")
|
||||
|
||||
def release(self) -> bool:
|
||||
"""Release collection from memory."""
|
||||
if not MILVUS_AVAILABLE:
|
||||
raise ProcessingError("Milvus not available")
|
||||
|
||||
try:
|
||||
self.collection.release()
|
||||
return True
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to release collection: {str(e)}")
|
||||
|
||||
|
||||
class MilvusSearch:
|
||||
"""Milvus search operations."""
|
||||
|
||||
def __init__(self, collection: MilvusCollection):
|
||||
"""Initialize Milvus search."""
|
||||
self.collection = collection
|
||||
self.logger = get_logger("milvus_search")
|
||||
|
||||
def similarity_search(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
limit: int = 10,
|
||||
metric_type: str = "L2",
|
||||
expr: Optional[str] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Perform similarity search.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
limit: Number of results
|
||||
metric_type: Distance metric ("L2", "IP", "COSINE")
|
||||
expr: Filter expression
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
search_params = {
|
||||
"metric_type": metric_type,
|
||||
"params": options.get("params", {"nprobe": 10})
|
||||
}
|
||||
|
||||
return self.collection.search(
|
||||
vectors=[query_vector],
|
||||
anns_field="vector",
|
||||
param=search_params,
|
||||
limit=limit,
|
||||
expr=expr,
|
||||
**options
|
||||
)
|
||||
|
||||
|
||||
class MilvusAdapter:
|
||||
"""
|
||||
Milvus adapter for vector storage and similarity search.
|
||||
|
||||
• Milvus connection and authentication
|
||||
• Collection and partition management
|
||||
• Vector storage and retrieval
|
||||
• Similarity search and filtering
|
||||
• Performance optimization
|
||||
• Error handling and recovery
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 19530,
|
||||
user: Optional[str] = None,
|
||||
password: Optional[str] = None,
|
||||
**config
|
||||
):
|
||||
"""Initialize Milvus adapter."""
|
||||
self.logger = get_logger("milvus_adapter")
|
||||
self.config = config
|
||||
self.host = host or config.get("host", "localhost")
|
||||
self.port = port or config.get("port", 19530)
|
||||
self.user = user or config.get("user")
|
||||
self.password = password or config.get("password")
|
||||
|
||||
self.client: Optional[MilvusClient] = None
|
||||
self.collection: Optional[MilvusCollection] = None
|
||||
self.search_engine: Optional[MilvusSearch] = None
|
||||
|
||||
# Check Milvus availability
|
||||
if not MILVUS_AVAILABLE:
|
||||
self.logger.warning(
|
||||
"Milvus not available. Install with: pip install pymilvus"
|
||||
)
|
||||
|
||||
def connect(self, **options) -> bool:
|
||||
"""
|
||||
Connect to Milvus service.
|
||||
|
||||
Args:
|
||||
**options: Connection options
|
||||
|
||||
Returns:
|
||||
True if connected successfully
|
||||
"""
|
||||
if not MILVUS_AVAILABLE:
|
||||
raise ProcessingError(
|
||||
"Milvus is not available. Install it with: pip install pymilvus"
|
||||
)
|
||||
|
||||
try:
|
||||
self.client = MilvusClient()
|
||||
self.client.connect(
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
user=self.user,
|
||||
password=self.password,
|
||||
**options
|
||||
)
|
||||
|
||||
self.logger.info(f"Connected to Milvus at {self.host}:{self.port}")
|
||||
return True
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to connect to Milvus: {str(e)}")
|
||||
|
||||
def create_collection(
|
||||
self,
|
||||
collection_name: str,
|
||||
dimension: int,
|
||||
metric_type: str = "L2",
|
||||
**options
|
||||
) -> MilvusCollection:
|
||||
"""
|
||||
Create Milvus collection.
|
||||
|
||||
Args:
|
||||
collection_name: Name of the collection
|
||||
dimension: Vector dimension
|
||||
metric_type: Distance metric ("L2", "IP", "COSINE")
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
MilvusCollection instance
|
||||
"""
|
||||
if self.client is None:
|
||||
self.connect()
|
||||
|
||||
if not MILVUS_AVAILABLE:
|
||||
raise ProcessingError("Milvus not available")
|
||||
|
||||
try:
|
||||
# Check if collection exists
|
||||
if utility.has_collection(collection_name):
|
||||
self.logger.info(f"Collection {collection_name} already exists")
|
||||
return self.get_collection(collection_name)
|
||||
|
||||
# Define schema
|
||||
fields = [
|
||||
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
|
||||
FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dimension)
|
||||
]
|
||||
|
||||
schema = CollectionSchema(
|
||||
fields=fields,
|
||||
description=f"Vector collection for {collection_name}"
|
||||
)
|
||||
|
||||
# Create collection
|
||||
collection = Collection(
|
||||
name=collection_name,
|
||||
schema=schema
|
||||
)
|
||||
|
||||
# Create index
|
||||
index_params = {
|
||||
"metric_type": metric_type,
|
||||
"index_type": "IVF_FLAT",
|
||||
"params": {"nlist": 1024}
|
||||
}
|
||||
collection.create_index(field_name="vector", index_params=index_params)
|
||||
|
||||
self.collection = MilvusCollection(collection, collection_name)
|
||||
self.search_engine = MilvusSearch(self.collection)
|
||||
|
||||
self.logger.info(f"Created Milvus collection: {collection_name}")
|
||||
return self.collection
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to create collection: {str(e)}")
|
||||
|
||||
def get_collection(self, collection_name: str) -> MilvusCollection:
|
||||
"""
|
||||
Get existing collection.
|
||||
|
||||
Args:
|
||||
collection_name: Name of the collection
|
||||
|
||||
Returns:
|
||||
MilvusCollection instance
|
||||
"""
|
||||
if self.client is None:
|
||||
self.connect()
|
||||
|
||||
if not MILVUS_AVAILABLE:
|
||||
raise ProcessingError("Milvus not available")
|
||||
|
||||
try:
|
||||
if not utility.has_collection(collection_name):
|
||||
raise ProcessingError(f"Collection {collection_name} does not exist")
|
||||
|
||||
collection = Collection(collection_name)
|
||||
self.collection = MilvusCollection(collection, collection_name)
|
||||
self.search_engine = MilvusSearch(self.collection)
|
||||
return self.collection
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to get collection: {str(e)}")
|
||||
|
||||
def insert_vectors(
|
||||
self,
|
||||
vectors: List[Union[np.ndarray, List[float]]],
|
||||
**options
|
||||
) -> Any:
|
||||
"""
|
||||
Insert vectors into collection.
|
||||
|
||||
Args:
|
||||
vectors: List of vectors
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
Insert result
|
||||
"""
|
||||
if self.collection is None:
|
||||
raise ProcessingError("Collection not initialized. Call create_collection() or get_collection() first.")
|
||||
|
||||
if not MILVUS_AVAILABLE:
|
||||
raise ProcessingError("Milvus not available")
|
||||
|
||||
try:
|
||||
# Convert vectors to list format
|
||||
vector_data = []
|
||||
for vector in vectors:
|
||||
if isinstance(vector, np.ndarray):
|
||||
vector = vector.tolist()
|
||||
vector_data.append(vector)
|
||||
|
||||
data = [vector_data]
|
||||
return self.collection.insert(data, **options)
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to insert vectors: {str(e)}")
|
||||
|
||||
def search_vectors(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
limit: int = 10,
|
||||
metric_type: str = "L2",
|
||||
expr: Optional[str] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search vectors in collection.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
limit: Number of results
|
||||
metric_type: Distance metric
|
||||
expr: Filter expression
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
if self.search_engine is None:
|
||||
raise ProcessingError("Collection not initialized. Call create_collection() or get_collection() first.")
|
||||
|
||||
# Load collection if not loaded
|
||||
if not self.collection.collection.has_index():
|
||||
self.collection.load()
|
||||
|
||||
return self.search_engine.similarity_search(
|
||||
query_vector, limit, metric_type, expr, **options
|
||||
)
|
||||
|
||||
def get_stats(self, collection_name: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Get collection statistics."""
|
||||
if self.collection is None and collection_name:
|
||||
self.get_collection(collection_name)
|
||||
|
||||
if self.collection is None:
|
||||
raise ProcessingError("Collection not initialized. Call create_collection() or get_collection() first.")
|
||||
|
||||
try:
|
||||
stats = self.collection.collection.num_entities
|
||||
return {
|
||||
"entity_count": stats,
|
||||
"collection_name": self.collection.collection_name
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to get stats: {str(e)}")
|
||||
return {"status": "unknown"}
|
||||
|
||||
@@ -5,10 +5,364 @@ This module provides namespace isolation and management
|
||||
for vector store operations.
|
||||
"""
|
||||
|
||||
# TODO: Implement namespace management
|
||||
# - Namespace creation and management
|
||||
# - Isolation and access control
|
||||
# - Namespace metadata and configuration
|
||||
# - Performance optimization
|
||||
# - Error handling and recovery
|
||||
# - Multi-tenant support
|
||||
from typing import Any, Dict, List, Optional, Set
|
||||
from datetime import datetime
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
|
||||
class Namespace:
|
||||
"""Namespace container."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
description: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
**config
|
||||
):
|
||||
"""Initialize namespace."""
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.metadata = metadata or {}
|
||||
self.config = config
|
||||
self.created_at = datetime.now()
|
||||
self.updated_at = datetime.now()
|
||||
self.vector_ids: Set[str] = set()
|
||||
self.access_control: Dict[str, List[str]] = {} # user/role -> permissions
|
||||
|
||||
def add_vector(self, vector_id: str):
|
||||
"""Add vector to namespace."""
|
||||
self.vector_ids.add(vector_id)
|
||||
self.updated_at = datetime.now()
|
||||
|
||||
def remove_vector(self, vector_id: str):
|
||||
"""Remove vector from namespace."""
|
||||
self.vector_ids.discard(vector_id)
|
||||
self.updated_at = datetime.now()
|
||||
|
||||
def has_vector(self, vector_id: str) -> bool:
|
||||
"""Check if namespace contains vector."""
|
||||
return vector_id in self.vector_ids
|
||||
|
||||
def get_vector_count(self) -> int:
|
||||
"""Get number of vectors in namespace."""
|
||||
return len(self.vector_ids)
|
||||
|
||||
def update_metadata(self, metadata: Dict[str, Any]):
|
||||
"""Update namespace metadata."""
|
||||
self.metadata.update(metadata)
|
||||
self.updated_at = datetime.now()
|
||||
|
||||
def set_access_control(self, entity: str, permissions: List[str]):
|
||||
"""Set access control for entity."""
|
||||
self.access_control[entity] = permissions
|
||||
self.updated_at = datetime.now()
|
||||
|
||||
def has_permission(self, entity: str, permission: str) -> bool:
|
||||
"""Check if entity has permission."""
|
||||
if entity in self.access_control:
|
||||
return permission in self.access_control[entity]
|
||||
return False
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert namespace to dictionary."""
|
||||
return {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"metadata": self.metadata,
|
||||
"config": self.config,
|
||||
"created_at": self.created_at.isoformat(),
|
||||
"updated_at": self.updated_at.isoformat(),
|
||||
"vector_count": len(self.vector_ids),
|
||||
"access_control": self.access_control
|
||||
}
|
||||
|
||||
|
||||
class NamespaceManager:
|
||||
"""
|
||||
Namespace manager for vector store operations.
|
||||
|
||||
• Namespace creation and management
|
||||
• Isolation and access control
|
||||
• Namespace metadata and configuration
|
||||
• Performance optimization
|
||||
• Error handling and recovery
|
||||
• Multi-tenant support
|
||||
"""
|
||||
|
||||
def __init__(self, **config):
|
||||
"""Initialize namespace manager."""
|
||||
self.logger = get_logger("namespace_manager")
|
||||
self.config = config
|
||||
self.namespaces: Dict[str, Namespace] = {}
|
||||
self.default_namespace = config.get("default_namespace", "default")
|
||||
self.vector_namespace_map: Dict[str, str] = {} # vector_id -> namespace
|
||||
|
||||
def create_namespace(
|
||||
self,
|
||||
name: str,
|
||||
description: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
**options
|
||||
) -> Namespace:
|
||||
"""
|
||||
Create a new namespace.
|
||||
|
||||
Args:
|
||||
name: Namespace name
|
||||
description: Namespace description
|
||||
metadata: Namespace metadata
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
Namespace instance
|
||||
"""
|
||||
if name in self.namespaces:
|
||||
raise ValidationError(f"Namespace '{name}' already exists")
|
||||
|
||||
if not self._validate_namespace_name(name):
|
||||
raise ValidationError(f"Invalid namespace name: {name}")
|
||||
|
||||
namespace = Namespace(name, description, metadata, **options)
|
||||
self.namespaces[name] = namespace
|
||||
|
||||
self.logger.info(f"Created namespace: {name}")
|
||||
return namespace
|
||||
|
||||
def get_namespace(self, name: str) -> Optional[Namespace]:
|
||||
"""
|
||||
Get namespace by name.
|
||||
|
||||
Args:
|
||||
name: Namespace name
|
||||
|
||||
Returns:
|
||||
Namespace instance or None
|
||||
"""
|
||||
return self.namespaces.get(name)
|
||||
|
||||
def delete_namespace(self, name: str, **options) -> bool:
|
||||
"""
|
||||
Delete namespace.
|
||||
|
||||
Args:
|
||||
name: Namespace name
|
||||
**options: Delete options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if name not in self.namespaces:
|
||||
raise ProcessingError(f"Namespace '{name}' does not exist")
|
||||
|
||||
if name == self.default_namespace:
|
||||
raise ProcessingError("Cannot delete default namespace")
|
||||
|
||||
# Remove namespace vectors from mapping
|
||||
namespace = self.namespaces[name]
|
||||
for vector_id in list(namespace.vector_ids):
|
||||
self.vector_namespace_map.pop(vector_id, None)
|
||||
|
||||
del self.namespaces[name]
|
||||
self.logger.info(f"Deleted namespace: {name}")
|
||||
return True
|
||||
|
||||
def list_namespaces(self, **options) -> List[str]:
|
||||
"""
|
||||
List all namespace names.
|
||||
|
||||
Args:
|
||||
**options: List options
|
||||
|
||||
Returns:
|
||||
List of namespace names
|
||||
"""
|
||||
return list(self.namespaces.keys())
|
||||
|
||||
def add_vector_to_namespace(
|
||||
self,
|
||||
vector_id: str,
|
||||
namespace: str,
|
||||
**options
|
||||
) -> bool:
|
||||
"""
|
||||
Add vector to namespace.
|
||||
|
||||
Args:
|
||||
vector_id: Vector ID
|
||||
namespace: Namespace name
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if namespace not in self.namespaces:
|
||||
raise ProcessingError(f"Namespace '{namespace}' does not exist")
|
||||
|
||||
# Remove from old namespace if exists
|
||||
old_namespace = self.vector_namespace_map.get(vector_id)
|
||||
if old_namespace and old_namespace in self.namespaces:
|
||||
self.namespaces[old_namespace].remove_vector(vector_id)
|
||||
|
||||
# Add to new namespace
|
||||
self.namespaces[namespace].add_vector(vector_id)
|
||||
self.vector_namespace_map[vector_id] = namespace
|
||||
|
||||
self.logger.debug(f"Added vector {vector_id} to namespace {namespace}")
|
||||
return True
|
||||
|
||||
def remove_vector_from_namespace(
|
||||
self,
|
||||
vector_id: str,
|
||||
namespace: Optional[str] = None,
|
||||
**options
|
||||
) -> bool:
|
||||
"""
|
||||
Remove vector from namespace.
|
||||
|
||||
Args:
|
||||
vector_id: Vector ID
|
||||
namespace: Namespace name (if None, uses vector's current namespace)
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if namespace is None:
|
||||
namespace = self.vector_namespace_map.get(vector_id)
|
||||
|
||||
if namespace and namespace in self.namespaces:
|
||||
self.namespaces[namespace].remove_vector(vector_id)
|
||||
self.vector_namespace_map.pop(vector_id, None)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def get_vector_namespace(self, vector_id: str) -> Optional[str]:
|
||||
"""
|
||||
Get namespace for a vector.
|
||||
|
||||
Args:
|
||||
vector_id: Vector ID
|
||||
|
||||
Returns:
|
||||
Namespace name or None
|
||||
"""
|
||||
return self.vector_namespace_map.get(vector_id)
|
||||
|
||||
def get_namespace_vectors(self, namespace: str) -> List[str]:
|
||||
"""
|
||||
Get all vectors in a namespace.
|
||||
|
||||
Args:
|
||||
namespace: Namespace name
|
||||
|
||||
Returns:
|
||||
List of vector IDs
|
||||
"""
|
||||
if namespace not in self.namespaces:
|
||||
return []
|
||||
|
||||
return list(self.namespaces[namespace].vector_ids)
|
||||
|
||||
def set_namespace_access_control(
|
||||
self,
|
||||
namespace: str,
|
||||
entity: str,
|
||||
permissions: List[str],
|
||||
**options
|
||||
) -> bool:
|
||||
"""
|
||||
Set access control for namespace.
|
||||
|
||||
Args:
|
||||
namespace: Namespace name
|
||||
entity: Entity (user/role) name
|
||||
permissions: List of permissions
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if namespace not in self.namespaces:
|
||||
raise ProcessingError(f"Namespace '{namespace}' does not exist")
|
||||
|
||||
self.namespaces[namespace].set_access_control(entity, permissions)
|
||||
return True
|
||||
|
||||
def check_namespace_access(
|
||||
self,
|
||||
namespace: str,
|
||||
entity: str,
|
||||
permission: str
|
||||
) -> bool:
|
||||
"""
|
||||
Check if entity has access to namespace.
|
||||
|
||||
Args:
|
||||
namespace: Namespace name
|
||||
entity: Entity name
|
||||
permission: Permission to check
|
||||
|
||||
Returns:
|
||||
True if entity has permission
|
||||
"""
|
||||
if namespace not in self.namespaces:
|
||||
return False
|
||||
|
||||
return self.namespaces[namespace].has_permission(entity, permission)
|
||||
|
||||
def get_namespace_stats(self, namespace: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get namespace statistics.
|
||||
|
||||
Args:
|
||||
namespace: Namespace name
|
||||
|
||||
Returns:
|
||||
Namespace statistics
|
||||
"""
|
||||
if namespace not in self.namespaces:
|
||||
raise ProcessingError(f"Namespace '{namespace}' does not exist")
|
||||
|
||||
ns = self.namespaces[namespace]
|
||||
return {
|
||||
"name": ns.name,
|
||||
"description": ns.description,
|
||||
"vector_count": ns.get_vector_count(),
|
||||
"created_at": ns.created_at.isoformat(),
|
||||
"updated_at": ns.updated_at.isoformat(),
|
||||
"metadata": ns.metadata
|
||||
}
|
||||
|
||||
def get_all_stats(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get statistics for all namespaces."""
|
||||
return {
|
||||
name: self.get_namespace_stats(name)
|
||||
for name in self.namespaces.keys()
|
||||
}
|
||||
|
||||
def _validate_namespace_name(self, name: str) -> bool:
|
||||
"""Validate namespace name."""
|
||||
if not name or not isinstance(name, str):
|
||||
return False
|
||||
|
||||
# Basic validation: alphanumeric, underscore, hyphen
|
||||
return all(c.isalnum() or c in ['_', '-'] for c in name)
|
||||
|
||||
def ensure_namespace(self, name: str, **options) -> Namespace:
|
||||
"""
|
||||
Ensure namespace exists, create if not.
|
||||
|
||||
Args:
|
||||
name: Namespace name
|
||||
**options: Creation options
|
||||
|
||||
Returns:
|
||||
Namespace instance
|
||||
"""
|
||||
if name not in self.namespaces:
|
||||
return self.create_namespace(name, **options)
|
||||
return self.namespaces[name]
|
||||
|
||||
@@ -5,10 +5,418 @@ This module provides Pinecone integration for vector storage
|
||||
and similarity search.
|
||||
"""
|
||||
|
||||
# TODO: Implement Pinecone adapter
|
||||
# - Pinecone connection and authentication
|
||||
# - Vector storage and retrieval
|
||||
# - Similarity search and filtering
|
||||
# - Namespace and index management
|
||||
# - Performance optimization
|
||||
# - Error handling and recovery
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
import numpy as np
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
# Optional Pinecone import
|
||||
try:
|
||||
import pinecone
|
||||
from pinecone import Pinecone, ServerlessSpec, PodSpec
|
||||
PINECONE_AVAILABLE = True
|
||||
except ImportError:
|
||||
PINECONE_AVAILABLE = False
|
||||
pinecone = None
|
||||
Pinecone = None
|
||||
ServerlessSpec = None
|
||||
PodSpec = None
|
||||
|
||||
|
||||
class PineconeIndex:
|
||||
"""Pinecone index wrapper."""
|
||||
|
||||
def __init__(self, index: Any, index_name: str):
|
||||
"""Initialize Pinecone index wrapper."""
|
||||
self.index = index
|
||||
self.index_name = index_name
|
||||
self.logger = get_logger("pinecone_index")
|
||||
|
||||
def upsert_vectors(
|
||||
self,
|
||||
vectors: List[Dict[str, Any]],
|
||||
namespace: Optional[str] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""Upsert vectors to index."""
|
||||
if not PINECONE_AVAILABLE:
|
||||
raise ProcessingError("Pinecone not available")
|
||||
|
||||
try:
|
||||
response = self.index.upsert(
|
||||
vectors=vectors,
|
||||
namespace=namespace,
|
||||
**options
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to upsert vectors: {str(e)}")
|
||||
|
||||
def query_vectors(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
top_k: int = 10,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""Query similar vectors."""
|
||||
if not PINECONE_AVAILABLE:
|
||||
raise ProcessingError("Pinecone not available")
|
||||
|
||||
try:
|
||||
response = self.index.query(
|
||||
vector=query_vector.tolist(),
|
||||
top_k=top_k,
|
||||
namespace=namespace,
|
||||
filter=filter,
|
||||
include_metadata=True,
|
||||
**options
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to query vectors: {str(e)}")
|
||||
|
||||
def delete_vectors(
|
||||
self,
|
||||
ids: List[str],
|
||||
namespace: Optional[str] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""Delete vectors from index."""
|
||||
if not PINECONE_AVAILABLE:
|
||||
raise ProcessingError("Pinecone not available")
|
||||
|
||||
try:
|
||||
response = self.index.delete(ids=ids, namespace=namespace, **options)
|
||||
return response
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to delete vectors: {str(e)}")
|
||||
|
||||
def fetch_vectors(
|
||||
self,
|
||||
ids: List[str],
|
||||
namespace: Optional[str] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""Fetch vectors by IDs."""
|
||||
if not PINECONE_AVAILABLE:
|
||||
raise ProcessingError("Pinecone not available")
|
||||
|
||||
try:
|
||||
response = self.index.fetch(ids=ids, namespace=namespace, **options)
|
||||
return response
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to fetch vectors: {str(e)}")
|
||||
|
||||
def describe_index_stats(self, namespace: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Get index statistics."""
|
||||
if not PINECONE_AVAILABLE:
|
||||
raise ProcessingError("Pinecone not available")
|
||||
|
||||
try:
|
||||
stats = self.index.describe_index_stats(namespace=namespace)
|
||||
return stats
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to get index stats: {str(e)}")
|
||||
|
||||
|
||||
class PineconeQuery:
|
||||
"""Pinecone query builder."""
|
||||
|
||||
def __init__(self, index: PineconeIndex):
|
||||
"""Initialize Pinecone query builder."""
|
||||
self.index = index
|
||||
self.logger = get_logger("pinecone_query")
|
||||
|
||||
def build_query(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
top_k: int = 10,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""Build query parameters."""
|
||||
return {
|
||||
"vector": query_vector.tolist(),
|
||||
"top_k": top_k,
|
||||
"namespace": namespace,
|
||||
"filter": filter,
|
||||
**options
|
||||
}
|
||||
|
||||
def execute(self, query_params: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
"""Execute query and format results."""
|
||||
response = self.index.query_vectors(**query_params)
|
||||
|
||||
results = []
|
||||
for match in response.get("matches", []):
|
||||
results.append({
|
||||
"id": match.get("id"),
|
||||
"score": match.get("score", 0.0),
|
||||
"metadata": match.get("metadata", {})
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class PineconeMetadata:
|
||||
"""Pinecone metadata handler."""
|
||||
|
||||
@staticmethod
|
||||
def validate_metadata(metadata: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Validate and sanitize metadata."""
|
||||
# Pinecone metadata restrictions
|
||||
validated = {}
|
||||
|
||||
for key, value in metadata.items():
|
||||
# Convert to allowed types
|
||||
if isinstance(value, (str, int, float, bool, list)):
|
||||
validated[key] = value
|
||||
elif isinstance(value, dict):
|
||||
# Nested dicts not directly supported
|
||||
validated[key] = str(value)
|
||||
else:
|
||||
validated[key] = str(value)
|
||||
|
||||
return validated
|
||||
|
||||
|
||||
class PineconeAdapter:
|
||||
"""
|
||||
Pinecone adapter for vector storage and similarity search.
|
||||
|
||||
• Pinecone connection and authentication
|
||||
• Vector storage and retrieval
|
||||
• Similarity search and filtering
|
||||
• Namespace and index management
|
||||
• Performance optimization
|
||||
• Error handling and recovery
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None, environment: Optional[str] = None, **config):
|
||||
"""Initialize Pinecone adapter."""
|
||||
self.logger = get_logger("pinecone_adapter")
|
||||
self.config = config
|
||||
self.api_key = api_key or config.get("api_key")
|
||||
self.environment = environment or config.get("environment")
|
||||
|
||||
self.client: Optional[Any] = None
|
||||
self.index: Optional[PineconeIndex] = None
|
||||
self.query_builder: Optional[PineconeQuery] = None
|
||||
|
||||
# Check Pinecone availability
|
||||
if not PINECONE_AVAILABLE:
|
||||
self.logger.warning(
|
||||
"Pinecone not available. Install with: pip install pinecone-client"
|
||||
)
|
||||
|
||||
def connect(self, api_key: Optional[str] = None, **options) -> bool:
|
||||
"""
|
||||
Connect to Pinecone service.
|
||||
|
||||
Args:
|
||||
api_key: Pinecone API key
|
||||
**options: Connection options
|
||||
|
||||
Returns:
|
||||
True if connected successfully
|
||||
"""
|
||||
if not PINECONE_AVAILABLE:
|
||||
raise ProcessingError(
|
||||
"Pinecone is not available. Install it with: pip install pinecone-client"
|
||||
)
|
||||
|
||||
api_key = api_key or self.api_key
|
||||
if not api_key:
|
||||
raise ValidationError("Pinecone API key is required")
|
||||
|
||||
try:
|
||||
self.client = Pinecone(api_key=api_key)
|
||||
self.logger.info("Connected to Pinecone")
|
||||
return True
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to connect to Pinecone: {str(e)}")
|
||||
|
||||
def create_index(
|
||||
self,
|
||||
index_name: str,
|
||||
dimension: int,
|
||||
metric: str = "cosine",
|
||||
spec: Optional[Dict[str, Any]] = None,
|
||||
**options
|
||||
) -> PineconeIndex:
|
||||
"""
|
||||
Create new vector index.
|
||||
|
||||
Args:
|
||||
index_name: Name of the index
|
||||
dimension: Vector dimension
|
||||
metric: Distance metric ("cosine", "euclidean", "dotproduct")
|
||||
spec: Index specification (serverless or pod)
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
PineconeIndex instance
|
||||
"""
|
||||
if self.client is None:
|
||||
self.connect()
|
||||
|
||||
if not PINECONE_AVAILABLE:
|
||||
raise ProcessingError("Pinecone not available")
|
||||
|
||||
try:
|
||||
# Check if index exists
|
||||
existing_indexes = [idx.name for idx in self.client.list_indexes()]
|
||||
if index_name in existing_indexes:
|
||||
self.logger.info(f"Index {index_name} already exists")
|
||||
return self.get_index(index_name)
|
||||
|
||||
# Create index specification
|
||||
if spec is None:
|
||||
spec = ServerlessSpec(cloud="aws", region="us-east-1")
|
||||
|
||||
# Create index
|
||||
self.client.create_index(
|
||||
name=index_name,
|
||||
dimension=dimension,
|
||||
metric=metric,
|
||||
spec=spec,
|
||||
**options
|
||||
)
|
||||
|
||||
self.logger.info(f"Created Pinecone index: {index_name}")
|
||||
return self.get_index(index_name)
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to create index: {str(e)}")
|
||||
|
||||
def get_index(self, index_name: str) -> PineconeIndex:
|
||||
"""
|
||||
Get existing index.
|
||||
|
||||
Args:
|
||||
index_name: Name of the index
|
||||
|
||||
Returns:
|
||||
PineconeIndex instance
|
||||
"""
|
||||
if self.client is None:
|
||||
self.connect()
|
||||
|
||||
if not PINECONE_AVAILABLE:
|
||||
raise ProcessingError("Pinecone not available")
|
||||
|
||||
try:
|
||||
index = self.client.Index(index_name)
|
||||
self.index = PineconeIndex(index, index_name)
|
||||
self.query_builder = PineconeQuery(self.index)
|
||||
return self.index
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to get index: {str(e)}")
|
||||
|
||||
def upsert_vectors(
|
||||
self,
|
||||
vectors: List[Union[np.ndarray, List[float]]],
|
||||
ids: List[str],
|
||||
metadata: Optional[List[Dict[str, Any]]] = None,
|
||||
namespace: Optional[str] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Insert or update vectors.
|
||||
|
||||
Args:
|
||||
vectors: List of vectors
|
||||
ids: Vector IDs
|
||||
metadata: Vector metadata
|
||||
namespace: Namespace name
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
Upsert response
|
||||
"""
|
||||
if self.index is None:
|
||||
raise ProcessingError("Index not initialized. Call create_index() or get_index() first.")
|
||||
|
||||
# Format vectors
|
||||
formatted_vectors = []
|
||||
for i, vector in enumerate(vectors):
|
||||
if isinstance(vector, np.ndarray):
|
||||
vector = vector.tolist()
|
||||
|
||||
vector_data = {"id": ids[i], "values": vector}
|
||||
|
||||
if metadata and i < len(metadata):
|
||||
vector_data["metadata"] = PineconeMetadata.validate_metadata(metadata[i])
|
||||
|
||||
formatted_vectors.append(vector_data)
|
||||
|
||||
return self.index.upsert_vectors(formatted_vectors, namespace, **options)
|
||||
|
||||
def query_vectors(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
top_k: int = 10,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Query similar vectors.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
top_k: Number of results
|
||||
namespace: Namespace name
|
||||
filter: Metadata filter
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
if self.query_builder is None:
|
||||
raise ProcessingError("Index not initialized. Call create_index() or get_index() first.")
|
||||
|
||||
query_params = self.query_builder.build_query(
|
||||
query_vector, top_k, namespace, filter, **options
|
||||
)
|
||||
|
||||
return self.query_builder.execute(query_params)
|
||||
|
||||
def delete_vectors(
|
||||
self,
|
||||
ids: List[str],
|
||||
namespace: Optional[str] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Delete vectors from index.
|
||||
|
||||
Args:
|
||||
ids: Vector IDs to delete
|
||||
namespace: Namespace name
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
Delete response
|
||||
"""
|
||||
if self.index is None:
|
||||
raise ProcessingError("Index not initialized. Call create_index() or get_index() first.")
|
||||
|
||||
return self.index.delete_vectors(ids, namespace, **options)
|
||||
|
||||
def get_stats(self, namespace: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Get index statistics."""
|
||||
if self.index is None:
|
||||
raise ProcessingError("Index not initialized. Call create_index() or get_index() first.")
|
||||
|
||||
stats = self.index.describe_index_stats(namespace)
|
||||
return {
|
||||
"total_vector_count": stats.get("total_vector_count", 0),
|
||||
"dimension": stats.get("dimension", 0),
|
||||
"index_fullness": stats.get("index_fullness", 0.0),
|
||||
"namespaces": stats.get("namespaces", {})
|
||||
}
|
||||
|
||||
@@ -5,10 +5,436 @@ This module provides Qdrant integration for vector storage
|
||||
and similarity search.
|
||||
"""
|
||||
|
||||
# TODO: Implement Qdrant adapter
|
||||
# - Qdrant connection and authentication
|
||||
# - Collection and point management
|
||||
# - Vector storage and retrieval
|
||||
# - Similarity search and filtering
|
||||
# - Performance optimization
|
||||
# - Error handling and recovery
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
import numpy as np
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
# Optional Qdrant import
|
||||
try:
|
||||
from qdrant_client import QdrantClient as QdrantClientLib
|
||||
from qdrant_client.models import (
|
||||
Distance, VectorParams, PointStruct, Filter, FieldCondition,
|
||||
MatchValue, CollectionStatus
|
||||
)
|
||||
QDRANT_AVAILABLE = True
|
||||
except ImportError:
|
||||
QDRANT_AVAILABLE = False
|
||||
QdrantClientLib = None
|
||||
Distance = None
|
||||
VectorParams = None
|
||||
PointStruct = None
|
||||
Filter = None
|
||||
FieldCondition = None
|
||||
MatchValue = None
|
||||
CollectionStatus = None
|
||||
|
||||
|
||||
class QdrantClient:
|
||||
"""Qdrant client wrapper."""
|
||||
|
||||
def __init__(self, client: Any):
|
||||
"""Initialize Qdrant client wrapper."""
|
||||
self.client = client
|
||||
self.logger = get_logger("qdrant_client")
|
||||
|
||||
def create_collection(
|
||||
self,
|
||||
collection_name: str,
|
||||
vector_size: int,
|
||||
distance: str = "Cosine",
|
||||
**options
|
||||
) -> bool:
|
||||
"""Create a collection in Qdrant."""
|
||||
if not QDRANT_AVAILABLE:
|
||||
raise ProcessingError("Qdrant not available")
|
||||
|
||||
try:
|
||||
distance_map = {
|
||||
"Cosine": Distance.COSINE,
|
||||
"Euclidean": Distance.EUCLID,
|
||||
"Dot": Distance.DOT
|
||||
}
|
||||
|
||||
distance_enum = distance_map.get(distance, Distance.COSINE)
|
||||
|
||||
self.client.create_collection(
|
||||
collection_name=collection_name,
|
||||
vectors_config=VectorParams(
|
||||
size=vector_size,
|
||||
distance=distance_enum
|
||||
),
|
||||
**options
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to create collection: {str(e)}")
|
||||
|
||||
def get_collection(self, collection_name: str) -> Any:
|
||||
"""Get collection info."""
|
||||
if not QDRANT_AVAILABLE:
|
||||
raise ProcessingError("Qdrant not available")
|
||||
|
||||
try:
|
||||
return self.client.get_collection(collection_name)
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to get collection: {str(e)}")
|
||||
|
||||
|
||||
class QdrantCollection:
|
||||
"""Qdrant collection wrapper."""
|
||||
|
||||
def __init__(self, client: Any, collection_name: str):
|
||||
"""Initialize Qdrant collection wrapper."""
|
||||
self.client = client
|
||||
self.collection_name = collection_name
|
||||
self.logger = get_logger("qdrant_collection")
|
||||
|
||||
def upsert_points(
|
||||
self,
|
||||
points: List[PointStruct],
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""Upsert points to collection."""
|
||||
if not QDRANT_AVAILABLE:
|
||||
raise ProcessingError("Qdrant not available")
|
||||
|
||||
try:
|
||||
response = self.client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=points,
|
||||
**options
|
||||
)
|
||||
return {"status": response.status}
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to upsert points: {str(e)}")
|
||||
|
||||
def search_points(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
limit: int = 10,
|
||||
query_filter: Optional[Filter] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Search for similar points."""
|
||||
if not QDRANT_AVAILABLE:
|
||||
raise ProcessingError("Qdrant not available")
|
||||
|
||||
try:
|
||||
search_results = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_vector.tolist(),
|
||||
limit=limit,
|
||||
query_filter=query_filter,
|
||||
with_payload=True,
|
||||
with_vectors=False,
|
||||
**options
|
||||
)
|
||||
|
||||
results = []
|
||||
for result in search_results:
|
||||
results.append({
|
||||
"id": result.id,
|
||||
"score": result.score,
|
||||
"payload": result.payload or {}
|
||||
})
|
||||
|
||||
return results
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to search points: {str(e)}")
|
||||
|
||||
def delete_points(self, point_ids: List[Union[str, int]], **options) -> Dict[str, Any]:
|
||||
"""Delete points from collection."""
|
||||
if not QDRANT_AVAILABLE:
|
||||
raise ProcessingError("Qdrant not available")
|
||||
|
||||
try:
|
||||
response = self.client.delete(
|
||||
collection_name=self.collection_name,
|
||||
points_selector=point_ids,
|
||||
**options
|
||||
)
|
||||
return {"status": response.status}
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to delete points: {str(e)}")
|
||||
|
||||
|
||||
class QdrantSearch:
|
||||
"""Qdrant search operations."""
|
||||
|
||||
def __init__(self, collection: QdrantCollection):
|
||||
"""Initialize Qdrant search."""
|
||||
self.collection = collection
|
||||
self.logger = get_logger("qdrant_search")
|
||||
|
||||
def similarity_search(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
limit: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Perform similarity search.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
limit: Number of results
|
||||
filter: Metadata filter
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
query_filter = None
|
||||
if filter and QDRANT_AVAILABLE:
|
||||
# Build Qdrant filter from dict
|
||||
conditions = []
|
||||
for key, value in filter.items():
|
||||
conditions.append(
|
||||
FieldCondition(key=key, match=MatchValue(value=value))
|
||||
)
|
||||
if conditions:
|
||||
query_filter = Filter(must=conditions)
|
||||
|
||||
return self.collection.search_points(query_vector, limit, query_filter, **options)
|
||||
|
||||
|
||||
class QdrantAdapter:
|
||||
"""
|
||||
Qdrant adapter for vector storage and similarity search.
|
||||
|
||||
• Qdrant connection and authentication
|
||||
• Collection and point management
|
||||
• Vector storage and retrieval
|
||||
• Similarity search and filtering
|
||||
• Performance optimization
|
||||
• Error handling and recovery
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
**config
|
||||
):
|
||||
"""Initialize Qdrant adapter."""
|
||||
self.logger = get_logger("qdrant_adapter")
|
||||
self.config = config
|
||||
self.url = url or config.get("url", "http://localhost:6333")
|
||||
self.api_key = api_key or config.get("api_key")
|
||||
|
||||
self.client: Optional[Any] = None
|
||||
self.collection: Optional[QdrantCollection] = None
|
||||
self.search_engine: Optional[QdrantSearch] = None
|
||||
|
||||
# Check Qdrant availability
|
||||
if not QDRANT_AVAILABLE:
|
||||
self.logger.warning(
|
||||
"Qdrant not available. Install with: pip install qdrant-client"
|
||||
)
|
||||
|
||||
def connect(self, url: Optional[str] = None, api_key: Optional[str] = None, **options) -> bool:
|
||||
"""
|
||||
Connect to Qdrant service.
|
||||
|
||||
Args:
|
||||
url: Qdrant URL
|
||||
api_key: API key for authentication
|
||||
**options: Connection options
|
||||
|
||||
Returns:
|
||||
True if connected successfully
|
||||
"""
|
||||
if not QDRANT_AVAILABLE:
|
||||
raise ProcessingError(
|
||||
"Qdrant is not available. Install it with: pip install qdrant-client"
|
||||
)
|
||||
|
||||
url = url or self.url
|
||||
api_key = api_key or self.api_key
|
||||
|
||||
try:
|
||||
if api_key:
|
||||
self.client = QdrantClientLib(url=url, api_key=api_key, **options)
|
||||
else:
|
||||
self.client = QdrantClientLib(url=url, **options)
|
||||
|
||||
self.logger.info(f"Connected to Qdrant at {url}")
|
||||
return True
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to connect to Qdrant: {str(e)}")
|
||||
|
||||
def create_collection(
|
||||
self,
|
||||
collection_name: str,
|
||||
vector_size: int,
|
||||
distance: str = "Cosine",
|
||||
**options
|
||||
) -> QdrantCollection:
|
||||
"""
|
||||
Create Qdrant collection.
|
||||
|
||||
Args:
|
||||
collection_name: Name of the collection
|
||||
vector_size: Vector dimension
|
||||
distance: Distance metric ("Cosine", "Euclidean", "Dot")
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
QdrantCollection instance
|
||||
"""
|
||||
if self.client is None:
|
||||
self.connect()
|
||||
|
||||
if not QDRANT_AVAILABLE:
|
||||
raise ProcessingError("Qdrant not available")
|
||||
|
||||
try:
|
||||
client_wrapper = QdrantClient(self.client)
|
||||
client_wrapper.create_collection(collection_name, vector_size, distance, **options)
|
||||
|
||||
self.collection = QdrantCollection(self.client, collection_name)
|
||||
self.search_engine = QdrantSearch(self.collection)
|
||||
|
||||
self.logger.info(f"Created Qdrant collection: {collection_name}")
|
||||
return self.collection
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to create collection: {str(e)}")
|
||||
|
||||
def get_collection(self, collection_name: str) -> QdrantCollection:
|
||||
"""
|
||||
Get existing collection.
|
||||
|
||||
Args:
|
||||
collection_name: Name of the collection
|
||||
|
||||
Returns:
|
||||
QdrantCollection instance
|
||||
"""
|
||||
if self.client is None:
|
||||
self.connect()
|
||||
|
||||
if not QDRANT_AVAILABLE:
|
||||
raise ProcessingError("Qdrant not available")
|
||||
|
||||
try:
|
||||
client_wrapper = QdrantClient(self.client)
|
||||
client_wrapper.get_collection(collection_name)
|
||||
|
||||
self.collection = QdrantCollection(self.client, collection_name)
|
||||
self.search_engine = QdrantSearch(self.collection)
|
||||
return self.collection
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to get collection: {str(e)}")
|
||||
|
||||
def insert_vectors(
|
||||
self,
|
||||
vectors: List[Union[np.ndarray, List[float]]],
|
||||
ids: List[Union[str, int]],
|
||||
payloads: Optional[List[Dict[str, Any]]] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Insert vectors into collection.
|
||||
|
||||
Args:
|
||||
vectors: List of vectors
|
||||
ids: Point IDs
|
||||
payloads: Optional metadata payloads
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
Insert response
|
||||
"""
|
||||
if self.collection is None:
|
||||
raise ProcessingError("Collection not initialized. Call create_collection() or get_collection() first.")
|
||||
|
||||
if not QDRANT_AVAILABLE:
|
||||
raise ProcessingError("Qdrant not available")
|
||||
|
||||
try:
|
||||
points = []
|
||||
for i, (vector, point_id) in enumerate(zip(vectors, ids)):
|
||||
if isinstance(vector, np.ndarray):
|
||||
vector = vector.tolist()
|
||||
|
||||
payload = payloads[i] if payloads and i < len(payloads) else {}
|
||||
|
||||
points.append(
|
||||
PointStruct(
|
||||
id=point_id,
|
||||
vector=vector,
|
||||
payload=payload
|
||||
)
|
||||
)
|
||||
|
||||
return self.collection.upsert_points(points, **options)
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to insert vectors: {str(e)}")
|
||||
|
||||
def search_vectors(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
limit: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search vectors in collection.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
limit: Number of results
|
||||
filter: Metadata filter
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
if self.search_engine is None:
|
||||
raise ProcessingError("Collection not initialized. Call create_collection() or get_collection() first.")
|
||||
|
||||
return self.search_engine.similarity_search(query_vector, limit, filter, **options)
|
||||
|
||||
def delete_vectors(
|
||||
self,
|
||||
point_ids: List[Union[str, int]],
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Delete vectors from collection.
|
||||
|
||||
Args:
|
||||
point_ids: Point IDs to delete
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
Delete response
|
||||
"""
|
||||
if self.collection is None:
|
||||
raise ProcessingError("Collection not initialized. Call create_collection() or get_collection() first.")
|
||||
|
||||
return self.collection.delete_points(point_ids, **options)
|
||||
|
||||
def get_stats(self, collection_name: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Get collection statistics."""
|
||||
if self.collection is None and collection_name:
|
||||
self.get_collection(collection_name)
|
||||
|
||||
if self.collection is None:
|
||||
raise ProcessingError("Collection not initialized. Call create_collection() or get_collection() first.")
|
||||
|
||||
try:
|
||||
collection_info = self.client.get_collection(self.collection.collection_name)
|
||||
return {
|
||||
"points_count": collection_info.points_count,
|
||||
"vectors_count": collection_info.vectors_count,
|
||||
"status": str(collection_info.status) if hasattr(collection_info, 'status') else "unknown"
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to get stats: {str(e)}")
|
||||
return {"status": "unknown"}
|
||||
|
||||
@@ -5,10 +5,425 @@ This module provides Weaviate integration for vector storage
|
||||
and similarity search.
|
||||
"""
|
||||
|
||||
# TODO: Implement Weaviate adapter
|
||||
# - Weaviate connection and authentication
|
||||
# - Schema and class management
|
||||
# - Vector storage and retrieval
|
||||
# - GraphQL query support
|
||||
# - Performance optimization
|
||||
# - Error handling and recovery
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
import numpy as np
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
# Optional Weaviate import
|
||||
try:
|
||||
import weaviate
|
||||
from weaviate.classes.query import MetadataQuery, QueryReturn
|
||||
WEAVIATE_AVAILABLE = True
|
||||
except ImportError:
|
||||
WEAVIATE_AVAILABLE = False
|
||||
weaviate = None
|
||||
MetadataQuery = None
|
||||
QueryReturn = None
|
||||
|
||||
|
||||
class WeaviateClient:
|
||||
"""Weaviate client wrapper."""
|
||||
|
||||
def __init__(self, client: Any):
|
||||
"""Initialize Weaviate client wrapper."""
|
||||
self.client = client
|
||||
self.logger = get_logger("weaviate_client")
|
||||
|
||||
def create_class(self, class_name: str, schema: Dict[str, Any], **options) -> bool:
|
||||
"""Create a class (collection) in Weaviate."""
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError("Weaviate not available")
|
||||
|
||||
try:
|
||||
# Convert schema to Weaviate format
|
||||
class_obj = self.client.collections.create(
|
||||
name=class_name,
|
||||
vectorizer_config=weaviate.classes.config.Configure.vectorizer.none(),
|
||||
**schema,
|
||||
**options
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to create class: {str(e)}")
|
||||
|
||||
def get_collection(self, class_name: str) -> Any:
|
||||
"""Get collection by class name."""
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError("Weaviate not available")
|
||||
|
||||
try:
|
||||
return self.client.collections.get(class_name)
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to get collection: {str(e)}")
|
||||
|
||||
def query_graphql(self, query: str, **options) -> Dict[str, Any]:
|
||||
"""Execute GraphQL query."""
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError("Weaviate not available")
|
||||
|
||||
try:
|
||||
result = self.client.query.raw(query)
|
||||
return result
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to execute GraphQL query: {str(e)}")
|
||||
|
||||
|
||||
class WeaviateSchema:
|
||||
"""Weaviate schema builder."""
|
||||
|
||||
@staticmethod
|
||||
def build_schema(
|
||||
class_name: str,
|
||||
properties: List[Dict[str, Any]],
|
||||
vectorizer: Optional[str] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""Build Weaviate schema."""
|
||||
schema = {
|
||||
"class": class_name,
|
||||
"properties": properties
|
||||
}
|
||||
|
||||
if vectorizer:
|
||||
schema["vectorizer"] = vectorizer
|
||||
|
||||
return schema
|
||||
|
||||
@staticmethod
|
||||
def build_property(
|
||||
name: str,
|
||||
data_type: str = "text",
|
||||
description: Optional[str] = None,
|
||||
**options
|
||||
) -> Dict[str, Any]:
|
||||
"""Build property definition."""
|
||||
prop = {
|
||||
"name": name,
|
||||
"dataType": [data_type]
|
||||
}
|
||||
|
||||
if description:
|
||||
prop["description"] = description
|
||||
|
||||
return {**prop, **options}
|
||||
|
||||
|
||||
class WeaviateQuery:
|
||||
"""Weaviate query builder."""
|
||||
|
||||
def __init__(self, collection: Any):
|
||||
"""Initialize Weaviate query builder."""
|
||||
self.collection = collection
|
||||
self.logger = get_logger("weaviate_query")
|
||||
|
||||
def similarity_search(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
limit: int = 10,
|
||||
where: Optional[Dict[str, Any]] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Perform similarity search.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
limit: Number of results
|
||||
where: Filter conditions
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError("Weaviate not available")
|
||||
|
||||
try:
|
||||
response = self.collection.query.near_vector(
|
||||
near_vector=query_vector.tolist(),
|
||||
limit=limit,
|
||||
where=where,
|
||||
return_metadata=MetadataQuery(distance=True),
|
||||
**options
|
||||
)
|
||||
|
||||
results = []
|
||||
for obj in response.objects:
|
||||
results.append({
|
||||
"id": str(obj.uuid),
|
||||
"properties": obj.properties,
|
||||
"distance": obj.metadata.distance if obj.metadata else None,
|
||||
"score": 1.0 - (obj.metadata.distance if obj.metadata and obj.metadata.distance else 0.0)
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to execute similarity search: {str(e)}")
|
||||
|
||||
def get_all(self, limit: int = 100, **options) -> List[Dict[str, Any]]:
|
||||
"""Get all objects from collection."""
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError("Weaviate not available")
|
||||
|
||||
try:
|
||||
response = self.collection.query.fetch_objects(limit=limit, **options)
|
||||
|
||||
results = []
|
||||
for obj in response.objects:
|
||||
results.append({
|
||||
"id": str(obj.uuid),
|
||||
"properties": obj.properties
|
||||
})
|
||||
|
||||
return results
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to get objects: {str(e)}")
|
||||
|
||||
|
||||
class WeaviateAdapter:
|
||||
"""
|
||||
Weaviate adapter for vector storage and similarity search.
|
||||
|
||||
• Weaviate connection and authentication
|
||||
• Schema and class management
|
||||
• Vector storage and retrieval
|
||||
• GraphQL query support
|
||||
• Performance optimization
|
||||
• Error handling and recovery
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
**config
|
||||
):
|
||||
"""Initialize Weaviate adapter."""
|
||||
self.logger = get_logger("weaviate_adapter")
|
||||
self.config = config
|
||||
self.url = url or config.get("url", "http://localhost:8080")
|
||||
self.api_key = api_key or config.get("api_key")
|
||||
|
||||
self.client: Optional[Any] = None
|
||||
self.collection: Optional[Any] = None
|
||||
self.query_builder: Optional[WeaviateQuery] = None
|
||||
|
||||
# Check Weaviate availability
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
self.logger.warning(
|
||||
"Weaviate not available. Install with: pip install weaviate-client"
|
||||
)
|
||||
|
||||
def connect(self, url: Optional[str] = None, api_key: Optional[str] = None, **options) -> bool:
|
||||
"""
|
||||
Connect to Weaviate service.
|
||||
|
||||
Args:
|
||||
url: Weaviate URL
|
||||
api_key: API key for authentication
|
||||
**options: Connection options
|
||||
|
||||
Returns:
|
||||
True if connected successfully
|
||||
"""
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError(
|
||||
"Weaviate is not available. Install it with: pip install weaviate-client"
|
||||
)
|
||||
|
||||
url = url or self.url
|
||||
api_key = api_key or self.api_key
|
||||
|
||||
try:
|
||||
auth_config = None
|
||||
if api_key:
|
||||
auth_config = weaviate.auth.AuthApiKey(api_key=api_key)
|
||||
|
||||
self.client = weaviate.connect_to_local(
|
||||
host=url.replace("http://", "").replace("https://", ""),
|
||||
auth_credentials=auth_config,
|
||||
**options
|
||||
)
|
||||
|
||||
self.logger.info(f"Connected to Weaviate at {url}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to connect to Weaviate: {str(e)}")
|
||||
|
||||
def create_schema(
|
||||
self,
|
||||
class_name: str,
|
||||
properties: List[Dict[str, Any]],
|
||||
vectorizer: Optional[str] = None,
|
||||
**options
|
||||
) -> bool:
|
||||
"""
|
||||
Create Weaviate schema.
|
||||
|
||||
Args:
|
||||
class_name: Name of the class
|
||||
properties: List of property definitions
|
||||
vectorizer: Vectorizer configuration
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if self.client is None:
|
||||
self.connect()
|
||||
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError("Weaviate not available")
|
||||
|
||||
try:
|
||||
client_wrapper = WeaviateClient(self.client)
|
||||
schema = WeaviateSchema.build_schema(class_name, properties, vectorizer, **options)
|
||||
client_wrapper.create_class(class_name, schema, **options)
|
||||
|
||||
self.logger.info(f"Created Weaviate schema for class: {class_name}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to create schema: {str(e)}")
|
||||
|
||||
def get_collection(self, class_name: str) -> Any:
|
||||
"""
|
||||
Get collection by class name.
|
||||
|
||||
Args:
|
||||
class_name: Name of the class
|
||||
|
||||
Returns:
|
||||
Collection instance
|
||||
"""
|
||||
if self.client is None:
|
||||
self.connect()
|
||||
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError("Weaviate not available")
|
||||
|
||||
try:
|
||||
self.collection = self.client.collections.get(class_name)
|
||||
self.query_builder = WeaviateQuery(self.collection)
|
||||
return self.collection
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to get collection: {str(e)}")
|
||||
|
||||
def add_objects(
|
||||
self,
|
||||
objects: List[Dict[str, Any]],
|
||||
vectors: Optional[List[np.ndarray]] = None,
|
||||
class_name: Optional[str] = None,
|
||||
**options
|
||||
) -> List[str]:
|
||||
"""
|
||||
Add objects to Weaviate.
|
||||
|
||||
Args:
|
||||
objects: List of objects with properties
|
||||
vectors: Optional list of vectors
|
||||
class_name: Class name (if not using default collection)
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of object IDs
|
||||
"""
|
||||
if self.collection is None and class_name:
|
||||
self.get_collection(class_name)
|
||||
|
||||
if self.collection is None:
|
||||
raise ProcessingError("Collection not initialized. Call get_collection() first.")
|
||||
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError("Weaviate not available")
|
||||
|
||||
try:
|
||||
object_ids = []
|
||||
|
||||
with self.collection.batch.dynamic() as batch:
|
||||
for i, obj in enumerate(objects):
|
||||
vector = vectors[i].tolist() if vectors and i < len(vectors) else None
|
||||
uuid = batch.add_object(
|
||||
properties=obj,
|
||||
vector=vector,
|
||||
**options
|
||||
)
|
||||
object_ids.append(str(uuid))
|
||||
|
||||
self.logger.info(f"Added {len(objects)} objects to Weaviate")
|
||||
return object_ids
|
||||
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to add objects: {str(e)}")
|
||||
|
||||
def query_vectors(
|
||||
self,
|
||||
query_vector: np.ndarray,
|
||||
limit: int = 10,
|
||||
where: Optional[Dict[str, Any]] = None,
|
||||
class_name: Optional[str] = None,
|
||||
**options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Query similar vectors.
|
||||
|
||||
Args:
|
||||
query_vector: Query vector
|
||||
limit: Number of results
|
||||
where: Filter conditions
|
||||
class_name: Class name (if not using default collection)
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
if self.collection is None and class_name:
|
||||
self.get_collection(class_name)
|
||||
|
||||
if self.query_builder is None:
|
||||
raise ProcessingError("Collection not initialized. Call get_collection() first.")
|
||||
|
||||
return self.query_builder.similarity_search(query_vector, limit, where, **options)
|
||||
|
||||
def graphql_query(self, query: str, **options) -> Dict[str, Any]:
|
||||
"""
|
||||
Execute GraphQL query.
|
||||
|
||||
Args:
|
||||
query: GraphQL query string
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
Query results
|
||||
"""
|
||||
if self.client is None:
|
||||
self.connect()
|
||||
|
||||
if not WEAVIATE_AVAILABLE:
|
||||
raise ProcessingError("Weaviate not available")
|
||||
|
||||
client_wrapper = WeaviateClient(self.client)
|
||||
return client_wrapper.query_graphql(query, **options)
|
||||
|
||||
def get_stats(self, class_name: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Get collection statistics."""
|
||||
if self.collection is None and class_name:
|
||||
self.get_collection(class_name)
|
||||
|
||||
if self.collection is None:
|
||||
raise ProcessingError("Collection not initialized. Call get_collection() first.")
|
||||
|
||||
try:
|
||||
# Get approximate count
|
||||
count = len(self.query_builder.get_all(limit=10000))
|
||||
return {
|
||||
"object_count": count,
|
||||
"class_name": class_name or "default"
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to get stats: {str(e)}")
|
||||
return {"status": "unknown"}
|
||||
|
||||
Reference in New Issue
Block a user