Merge pull request #15 from Hawksight-AI/feat/vector-store

feat: Implement vector_store module - comprehensive vector database i…
This commit is contained in:
Mohd Kaif
2025-11-04 19:57:48 +05:30
committed by GitHub
9 changed files with 3343 additions and 60 deletions
+110 -4
View File
@@ -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",
]
+366 -7
View File
@@ -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
}
+411 -7
View File
@@ -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
+384 -7
View File
@@ -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
+441 -7
View File
@@ -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]
+415 -7
View File
@@ -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", {})
}
+433 -7
View File
@@ -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"}
+422 -7
View File
@@ -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"}