mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-06 20:02:11 +02:00
Fix hard coded vector size (#555)
* Fixed hard-coded embeddings store size * Vector store lazy-creates collections, different collections for different dimension lengths. * Added tech spec for vector store lifecycle * Fixed some tests for the new spec
This commit is contained in:
parent
05b9063fea
commit
6129bb68c1
22 changed files with 793 additions and 572 deletions
299
docs/tech-specs/vector-store-lifecycle.md
Normal file
299
docs/tech-specs/vector-store-lifecycle.md
Normal file
|
|
@ -0,0 +1,299 @@
|
||||||
|
# Vector Store Lifecycle Management
|
||||||
|
|
||||||
|
## Overview
|
||||||
|
|
||||||
|
This document describes how TrustGraph manages vector store collections across different backend implementations (Qdrant, Pinecone, Milvus). The design addresses the challenge of supporting embeddings with different dimensions without hardcoding dimension values.
|
||||||
|
|
||||||
|
## Problem Statement
|
||||||
|
|
||||||
|
Vector stores require the embedding dimension to be specified when creating collections/indexes. However:
|
||||||
|
- Different embedding models produce different dimensions (e.g., 384, 768, 1536)
|
||||||
|
- The dimension is not known until the first embedding is generated
|
||||||
|
- A single TrustGraph collection may receive embeddings from multiple models
|
||||||
|
- Hardcoding a dimension (e.g., 384) causes failures with other embedding sizes
|
||||||
|
|
||||||
|
## Design Principles
|
||||||
|
|
||||||
|
1. **Lazy Creation**: Collections are created on-demand during first write, not during collection management operations
|
||||||
|
2. **Dimension-Based Naming**: Collection names include the embedding dimension as a suffix
|
||||||
|
3. **Graceful Degradation**: Queries against non-existent collections return empty results, not errors
|
||||||
|
4. **Multi-Dimension Support**: A single logical collection can have multiple physical collections (one per dimension)
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
### Collection Naming Convention
|
||||||
|
|
||||||
|
Vector store collections use dimension suffixes to support multiple embedding sizes:
|
||||||
|
|
||||||
|
**Document Embeddings:**
|
||||||
|
- Qdrant: `d_{user}_{collection}_{dimension}`
|
||||||
|
- Pinecone: `d-{user}-{collection}-{dimension}`
|
||||||
|
- Milvus: `doc_{user}_{collection}_{dimension}`
|
||||||
|
|
||||||
|
**Graph Embeddings:**
|
||||||
|
- Qdrant: `t_{user}_{collection}_{dimension}`
|
||||||
|
- Pinecone: `t-{user}-{collection}-{dimension}`
|
||||||
|
- Milvus: `entity_{user}_{collection}_{dimension}`
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
- `d_alice_papers_384` - Alice's papers collection with 384-dimensional embeddings
|
||||||
|
- `d_alice_papers_768` - Same logical collection with 768-dimensional embeddings
|
||||||
|
- `t_bob_knowledge_1536` - Bob's knowledge graph with 1536-dimensional embeddings
|
||||||
|
|
||||||
|
### Lifecycle Phases
|
||||||
|
|
||||||
|
#### 1. Collection Creation Request
|
||||||
|
|
||||||
|
**Request Flow:**
|
||||||
|
```
|
||||||
|
User/System → Librarian → Storage Management Topic → Vector Stores
|
||||||
|
```
|
||||||
|
|
||||||
|
**Behavior:**
|
||||||
|
- The librarian broadcasts `create-collection` requests to all storage backends
|
||||||
|
- Vector store processors acknowledge the request but **do not create physical collections**
|
||||||
|
- Response is returned immediately with success
|
||||||
|
- Actual collection creation is deferred until first write
|
||||||
|
|
||||||
|
**Rationale:**
|
||||||
|
- Dimension is unknown at creation time
|
||||||
|
- Avoids creating collections with wrong dimensions
|
||||||
|
- Simplifies collection management logic
|
||||||
|
|
||||||
|
#### 2. Write Operations (Lazy Creation)
|
||||||
|
|
||||||
|
**Write Flow:**
|
||||||
|
```
|
||||||
|
Data → Storage Processor → Check Collection → Create if Needed → Insert
|
||||||
|
```
|
||||||
|
|
||||||
|
**Behavior:**
|
||||||
|
1. Extract embedding dimension from the vector: `dim = len(vector)`
|
||||||
|
2. Construct collection name with dimension suffix
|
||||||
|
3. Check if collection exists with that specific dimension
|
||||||
|
4. If not exists:
|
||||||
|
- Create collection with correct dimension
|
||||||
|
- Log: `"Lazily creating collection {name} with dimension {dim}"`
|
||||||
|
5. Insert the embedding into the dimension-specific collection
|
||||||
|
|
||||||
|
**Example Scenario:**
|
||||||
|
```
|
||||||
|
1. User creates collection "papers"
|
||||||
|
→ No physical collections created yet
|
||||||
|
|
||||||
|
2. First document with 384-dim embedding arrives
|
||||||
|
→ Creates d_user_papers_384
|
||||||
|
→ Inserts data
|
||||||
|
|
||||||
|
3. Second document with 768-dim embedding arrives
|
||||||
|
→ Creates d_user_papers_768
|
||||||
|
→ Inserts data
|
||||||
|
|
||||||
|
Result: Two physical collections for one logical collection
|
||||||
|
```
|
||||||
|
|
||||||
|
#### 3. Query Operations
|
||||||
|
|
||||||
|
**Query Flow:**
|
||||||
|
```
|
||||||
|
Query Vector → Determine Dimension → Check Collection → Search or Return Empty
|
||||||
|
```
|
||||||
|
|
||||||
|
**Behavior:**
|
||||||
|
1. Extract dimension from query vector: `dim = len(vector)`
|
||||||
|
2. Construct collection name with dimension suffix
|
||||||
|
3. Check if collection exists
|
||||||
|
4. If exists:
|
||||||
|
- Perform similarity search
|
||||||
|
- Return results
|
||||||
|
5. If not exists:
|
||||||
|
- Log: `"Collection {name} does not exist, returning empty results"`
|
||||||
|
- Return empty list (no error raised)
|
||||||
|
|
||||||
|
**Multiple Dimensions in Same Query:**
|
||||||
|
- If query contains vectors of different dimensions
|
||||||
|
- Each dimension queries its corresponding collection
|
||||||
|
- Results are aggregated
|
||||||
|
- Missing collections are skipped (not treated as errors)
|
||||||
|
|
||||||
|
**Rationale:**
|
||||||
|
- Querying an empty collection is a valid use case
|
||||||
|
- Returning empty results is semantically correct
|
||||||
|
- Avoids errors during system startup or before data ingestion
|
||||||
|
|
||||||
|
#### 4. Collection Deletion
|
||||||
|
|
||||||
|
**Delete Flow:**
|
||||||
|
```
|
||||||
|
Delete Request → List All Collections → Filter by Prefix → Delete All Matches
|
||||||
|
```
|
||||||
|
|
||||||
|
**Behavior:**
|
||||||
|
1. Construct prefix pattern: `d_{user}_{collection}_` (note trailing underscore)
|
||||||
|
2. List all collections in the vector store
|
||||||
|
3. Filter collections matching the prefix
|
||||||
|
4. Delete all matching collections
|
||||||
|
5. Log each deletion: `"Deleted collection {name}"`
|
||||||
|
6. Summary log: `"Deleted {count} collection(s) for {user}/{collection}"`
|
||||||
|
|
||||||
|
**Example:**
|
||||||
|
```
|
||||||
|
Collections in store:
|
||||||
|
- d_alice_papers_384
|
||||||
|
- d_alice_papers_768
|
||||||
|
- d_alice_reports_384
|
||||||
|
- d_bob_papers_384
|
||||||
|
|
||||||
|
Delete "papers" for alice:
|
||||||
|
→ Deletes: d_alice_papers_384, d_alice_papers_768
|
||||||
|
→ Keeps: d_alice_reports_384, d_bob_papers_384
|
||||||
|
```
|
||||||
|
|
||||||
|
**Rationale:**
|
||||||
|
- Ensures complete cleanup of all dimension variants
|
||||||
|
- Pattern matching prevents accidental deletion of unrelated collections
|
||||||
|
- Atomic operation from user perspective (all dimensions deleted together)
|
||||||
|
|
||||||
|
## Behavioral Characteristics
|
||||||
|
|
||||||
|
### Normal Operations
|
||||||
|
|
||||||
|
**Collection Creation:**
|
||||||
|
- ✓ Returns success immediately
|
||||||
|
- ✓ No physical storage allocated
|
||||||
|
- ✓ Fast operation (no backend I/O)
|
||||||
|
|
||||||
|
**First Write:**
|
||||||
|
- ✓ Creates collection with correct dimension
|
||||||
|
- ✓ Slightly slower due to collection creation overhead
|
||||||
|
- ✓ Subsequent writes to same dimension are fast
|
||||||
|
|
||||||
|
**Queries Before Any Writes:**
|
||||||
|
- ✓ Returns empty results
|
||||||
|
- ✓ No errors or exceptions
|
||||||
|
- ✓ System remains stable
|
||||||
|
|
||||||
|
**Mixed Dimension Writes:**
|
||||||
|
- ✓ Automatically creates separate collections per dimension
|
||||||
|
- ✓ Each dimension isolated in its own collection
|
||||||
|
- ✓ No dimension conflicts or schema errors
|
||||||
|
|
||||||
|
**Collection Deletion:**
|
||||||
|
- ✓ Removes all dimension variants
|
||||||
|
- ✓ Complete cleanup
|
||||||
|
- ✓ No orphaned collections
|
||||||
|
|
||||||
|
### Edge Cases
|
||||||
|
|
||||||
|
**Multiple Embedding Models:**
|
||||||
|
```
|
||||||
|
Scenario: User switches from model A (384-dim) to model B (768-dim)
|
||||||
|
Behavior:
|
||||||
|
- Both dimensions coexist in separate collections
|
||||||
|
- Old data (384-dim) remains queryable with 384-dim vectors
|
||||||
|
- New data (768-dim) queryable with 768-dim vectors
|
||||||
|
- Cross-dimension queries return results only for matching dimension
|
||||||
|
```
|
||||||
|
|
||||||
|
**Concurrent First Writes:**
|
||||||
|
```
|
||||||
|
Scenario: Multiple processes write to same collection simultaneously
|
||||||
|
Behavior:
|
||||||
|
- Each process checks for existence before creating
|
||||||
|
- Most vector stores handle concurrent creation gracefully
|
||||||
|
- If race condition occurs, second create is typically idempotent
|
||||||
|
- Final state: Collection exists and both writes succeed
|
||||||
|
```
|
||||||
|
|
||||||
|
**Dimension Migration:**
|
||||||
|
```
|
||||||
|
Scenario: User wants to migrate from 384-dim to 768-dim embeddings
|
||||||
|
Behavior:
|
||||||
|
- No automatic migration
|
||||||
|
- Old collection (384-dim) persists
|
||||||
|
- New collection (768-dim) created on first new write
|
||||||
|
- Both dimensions remain accessible
|
||||||
|
- Manual deletion of old dimension collections possible
|
||||||
|
```
|
||||||
|
|
||||||
|
**Empty Collection Queries:**
|
||||||
|
```
|
||||||
|
Scenario: Query a collection that has never received data
|
||||||
|
Behavior:
|
||||||
|
- Collection doesn't exist (never created)
|
||||||
|
- Query returns empty list
|
||||||
|
- No error state
|
||||||
|
- System logs: "Collection does not exist, returning empty results"
|
||||||
|
```
|
||||||
|
|
||||||
|
## Implementation Notes
|
||||||
|
|
||||||
|
### Storage Backend Specifics
|
||||||
|
|
||||||
|
**Qdrant:**
|
||||||
|
- Uses `collection_exists()` for existence checks
|
||||||
|
- Uses `get_collections()` for listing during deletion
|
||||||
|
- Collection creation requires `VectorParams(size=dim, distance=Distance.COSINE)`
|
||||||
|
|
||||||
|
**Pinecone:**
|
||||||
|
- Uses `has_index()` for existence checks
|
||||||
|
- Uses `list_indexes()` for listing during deletion
|
||||||
|
- Index creation requires waiting for "ready" status
|
||||||
|
- Serverless spec configured with cloud/region
|
||||||
|
|
||||||
|
**Milvus:**
|
||||||
|
- Direct classes (`DocVectors`, `EntityVectors`) manage lifecycle
|
||||||
|
- Internal cache `self.collections[(dim, user, collection)]` for performance
|
||||||
|
- Collection names sanitized (alphanumeric + underscore only)
|
||||||
|
- Supports schema with auto-incrementing IDs
|
||||||
|
|
||||||
|
### Performance Considerations
|
||||||
|
|
||||||
|
**First Write Latency:**
|
||||||
|
- Additional overhead due to collection creation
|
||||||
|
- Qdrant: ~100-500ms
|
||||||
|
- Pinecone: ~10-30 seconds (serverless provisioning)
|
||||||
|
- Milvus: ~500-2000ms (includes indexing)
|
||||||
|
|
||||||
|
**Query Performance:**
|
||||||
|
- Existence check adds minimal overhead (~1-10ms)
|
||||||
|
- No performance impact once collection exists
|
||||||
|
- Each dimension collection is independently optimized
|
||||||
|
|
||||||
|
**Storage Overhead:**
|
||||||
|
- Minimal metadata per collection
|
||||||
|
- Main overhead is per-dimension storage
|
||||||
|
- Trade-off: Storage space vs. dimension flexibility
|
||||||
|
|
||||||
|
## Future Considerations
|
||||||
|
|
||||||
|
**Automatic Dimension Consolidation:**
|
||||||
|
- Could add background process to identify and merge unused dimension variants
|
||||||
|
- Would require re-embedding or dimension reduction
|
||||||
|
|
||||||
|
**Dimension Discovery:**
|
||||||
|
- Could expose API to list all dimensions in use for a collection
|
||||||
|
- Useful for administration and monitoring
|
||||||
|
|
||||||
|
**Default Dimension Preference:**
|
||||||
|
- Could track "primary" dimension per collection
|
||||||
|
- Use for queries when dimension context is unavailable
|
||||||
|
|
||||||
|
**Storage Quotas:**
|
||||||
|
- May need per-collection dimension limits
|
||||||
|
- Prevent proliferation of dimension variants
|
||||||
|
|
||||||
|
## Migration Notes
|
||||||
|
|
||||||
|
**From Pre-Dimension-Suffix System:**
|
||||||
|
- Old collections: `d_{user}_{collection}` (no dimension suffix)
|
||||||
|
- New collections: `d_{user}_{collection}_{dim}` (with dimension suffix)
|
||||||
|
- No automatic migration - old collections remain accessible
|
||||||
|
- Consider manual migration script if needed
|
||||||
|
- Can run both naming schemes simultaneously
|
||||||
|
|
||||||
|
## References
|
||||||
|
|
||||||
|
- Collection Management: `docs/tech-specs/collection-management.md`
|
||||||
|
- Storage Schema: `trustgraph-base/trustgraph/schema/services/storage.py`
|
||||||
|
- Librarian Service: `trustgraph-flow/trustgraph/librarian/service.py`
|
||||||
|
|
@ -34,10 +34,12 @@ class TestMilvusUserCollectionIntegration:
|
||||||
expected_collection_name = make_safe_collection_name(
|
expected_collection_name = make_safe_collection_name(
|
||||||
user, collection, "doc"
|
user, collection, "doc"
|
||||||
)
|
)
|
||||||
|
# Add dimension suffix to expected name
|
||||||
|
expected_collection_name_with_dim = f"{expected_collection_name}_{len(vector)}"
|
||||||
|
|
||||||
# Verify collection was created with correct name
|
# Verify collection was created with correct name (including dimension)
|
||||||
assert (len(vector), user, collection) in doc_vectors.collections
|
assert (len(vector), user, collection) in doc_vectors.collections
|
||||||
assert doc_vectors.collections[(len(vector), user, collection)] == expected_collection_name
|
assert doc_vectors.collections[(len(vector), user, collection)] == expected_collection_name_with_dim
|
||||||
|
|
||||||
@patch('trustgraph.direct.milvus_graph_embeddings.MilvusClient')
|
@patch('trustgraph.direct.milvus_graph_embeddings.MilvusClient')
|
||||||
def test_entity_vectors_collection_creation_with_user_collection(self, mock_milvus_client):
|
def test_entity_vectors_collection_creation_with_user_collection(self, mock_milvus_client):
|
||||||
|
|
@ -60,10 +62,12 @@ class TestMilvusUserCollectionIntegration:
|
||||||
expected_collection_name = make_safe_collection_name(
|
expected_collection_name = make_safe_collection_name(
|
||||||
user, collection, "entity"
|
user, collection, "entity"
|
||||||
)
|
)
|
||||||
|
# Add dimension suffix to expected name
|
||||||
|
expected_collection_name_with_dim = f"{expected_collection_name}_{len(vector)}"
|
||||||
|
|
||||||
# Verify collection was created with correct name
|
# Verify collection was created with correct name (including dimension)
|
||||||
assert (len(vector), user, collection) in entity_vectors.collections
|
assert (len(vector), user, collection) in entity_vectors.collections
|
||||||
assert entity_vectors.collections[(len(vector), user, collection)] == expected_collection_name
|
assert entity_vectors.collections[(len(vector), user, collection)] == expected_collection_name_with_dim
|
||||||
|
|
||||||
@patch('trustgraph.direct.milvus_doc_embeddings.MilvusClient')
|
@patch('trustgraph.direct.milvus_doc_embeddings.MilvusClient')
|
||||||
def test_doc_vectors_search_uses_correct_collection(self, mock_milvus_client):
|
def test_doc_vectors_search_uses_correct_collection(self, mock_milvus_client):
|
||||||
|
|
@ -88,11 +92,12 @@ class TestMilvusUserCollectionIntegration:
|
||||||
# Now search
|
# Now search
|
||||||
result = doc_vectors.search(vector, user, collection, limit=5)
|
result = doc_vectors.search(vector, user, collection, limit=5)
|
||||||
|
|
||||||
# Verify search was called with correct collection name
|
# Verify search was called with correct collection name (including dimension)
|
||||||
expected_collection_name = make_safe_collection_name(user, collection, "doc")
|
expected_collection_name = make_safe_collection_name(user, collection, "doc")
|
||||||
|
expected_collection_name_with_dim = f"{expected_collection_name}_{len(vector)}"
|
||||||
mock_client.search.assert_called_once()
|
mock_client.search.assert_called_once()
|
||||||
search_call = mock_client.search.call_args
|
search_call = mock_client.search.call_args
|
||||||
assert search_call[1]["collection_name"] == expected_collection_name
|
assert search_call[1]["collection_name"] == expected_collection_name_with_dim
|
||||||
|
|
||||||
@patch('trustgraph.direct.milvus_graph_embeddings.MilvusClient')
|
@patch('trustgraph.direct.milvus_graph_embeddings.MilvusClient')
|
||||||
def test_entity_vectors_search_uses_correct_collection(self, mock_milvus_client):
|
def test_entity_vectors_search_uses_correct_collection(self, mock_milvus_client):
|
||||||
|
|
@ -117,11 +122,12 @@ class TestMilvusUserCollectionIntegration:
|
||||||
# Now search
|
# Now search
|
||||||
result = entity_vectors.search(vector, user, collection, limit=5)
|
result = entity_vectors.search(vector, user, collection, limit=5)
|
||||||
|
|
||||||
# Verify search was called with correct collection name
|
# Verify search was called with correct collection name (including dimension)
|
||||||
expected_collection_name = make_safe_collection_name(user, collection, "entity")
|
expected_collection_name = make_safe_collection_name(user, collection, "entity")
|
||||||
|
expected_collection_name_with_dim = f"{expected_collection_name}_{len(vector)}"
|
||||||
mock_client.search.assert_called_once()
|
mock_client.search.assert_called_once()
|
||||||
search_call = mock_client.search.call_args
|
search_call = mock_client.search.call_args
|
||||||
assert search_call[1]["collection_name"] == expected_collection_name
|
assert search_call[1]["collection_name"] == expected_collection_name_with_dim
|
||||||
|
|
||||||
@patch('trustgraph.direct.milvus_doc_embeddings.MilvusClient')
|
@patch('trustgraph.direct.milvus_doc_embeddings.MilvusClient')
|
||||||
def test_doc_vectors_collection_isolation(self, mock_milvus_client):
|
def test_doc_vectors_collection_isolation(self, mock_milvus_client):
|
||||||
|
|
@ -141,10 +147,11 @@ class TestMilvusUserCollectionIntegration:
|
||||||
assert len(doc_vectors.collections) == 3
|
assert len(doc_vectors.collections) == 3
|
||||||
|
|
||||||
collection_names = set(doc_vectors.collections.values())
|
collection_names = set(doc_vectors.collections.values())
|
||||||
|
# All vectors are 3-dimensional, so all names should have _3 suffix
|
||||||
expected_names = {
|
expected_names = {
|
||||||
"doc_user1_collection1",
|
"doc_user1_collection1_3",
|
||||||
"doc_user2_collection2",
|
"doc_user2_collection2_3",
|
||||||
"doc_user1_collection2"
|
"doc_user1_collection2_3"
|
||||||
}
|
}
|
||||||
assert collection_names == expected_names
|
assert collection_names == expected_names
|
||||||
|
|
||||||
|
|
@ -166,10 +173,11 @@ class TestMilvusUserCollectionIntegration:
|
||||||
assert len(entity_vectors.collections) == 3
|
assert len(entity_vectors.collections) == 3
|
||||||
|
|
||||||
collection_names = set(entity_vectors.collections.values())
|
collection_names = set(entity_vectors.collections.values())
|
||||||
|
# All vectors are 3-dimensional, so all names should have _3 suffix
|
||||||
expected_names = {
|
expected_names = {
|
||||||
"entity_user1_collection1",
|
"entity_user1_collection1_3",
|
||||||
"entity_user2_collection2",
|
"entity_user2_collection2_3",
|
||||||
"entity_user1_collection2"
|
"entity_user1_collection2_3"
|
||||||
}
|
}
|
||||||
assert collection_names == expected_names
|
assert collection_names == expected_names
|
||||||
|
|
||||||
|
|
@ -193,14 +201,14 @@ class TestMilvusUserCollectionIntegration:
|
||||||
assert len(doc_vectors.collections) == 3
|
assert len(doc_vectors.collections) == 3
|
||||||
|
|
||||||
collection_names = set(doc_vectors.collections.values())
|
collection_names = set(doc_vectors.collections.values())
|
||||||
|
# Different dimensions now create different collections with dimension suffixes
|
||||||
expected_names = {
|
expected_names = {
|
||||||
"doc_test_user_test_collection", # Same name for all dimensions
|
"doc_test_user_test_collection_2", # 2D vector
|
||||||
"doc_test_user_test_collection", # now stored per dimension in key
|
"doc_test_user_test_collection_3", # 3D vector
|
||||||
"doc_test_user_test_collection" # but collection name is the same
|
"doc_test_user_test_collection_4" # 4D vector
|
||||||
}
|
}
|
||||||
# Note: Now all dimensions use the same collection name, they are differentiated by the key
|
# Each dimension gets its own collection
|
||||||
assert len(collection_names) == 1 # Only one unique collection name
|
assert len(collection_names) == 3 # Three unique collection names
|
||||||
assert "doc_test_user_test_collection" in collection_names
|
|
||||||
assert collection_names == expected_names
|
assert collection_names == expected_names
|
||||||
|
|
||||||
@patch('trustgraph.direct.milvus_doc_embeddings.MilvusClient')
|
@patch('trustgraph.direct.milvus_doc_embeddings.MilvusClient')
|
||||||
|
|
@ -223,7 +231,8 @@ class TestMilvusUserCollectionIntegration:
|
||||||
# Verify only one collection was created
|
# Verify only one collection was created
|
||||||
assert len(doc_vectors.collections) == 1
|
assert len(doc_vectors.collections) == 1
|
||||||
|
|
||||||
expected_collection_name = "doc_test_user_test_collection"
|
# Collection name now includes dimension suffix
|
||||||
|
expected_collection_name = "doc_test_user_test_collection_3"
|
||||||
assert doc_vectors.collections[(3, user, collection)] == expected_collection_name
|
assert doc_vectors.collections[(3, user, collection)] == expected_collection_name
|
||||||
|
|
||||||
@patch('trustgraph.direct.milvus_doc_embeddings.MilvusClient')
|
@patch('trustgraph.direct.milvus_doc_embeddings.MilvusClient')
|
||||||
|
|
@ -235,11 +244,12 @@ class TestMilvusUserCollectionIntegration:
|
||||||
doc_vectors = DocVectors(uri="http://test:19530", prefix="doc")
|
doc_vectors = DocVectors(uri="http://test:19530", prefix="doc")
|
||||||
|
|
||||||
# Test various special character combinations
|
# Test various special character combinations
|
||||||
|
# All expected names now include dimension suffix _3
|
||||||
test_cases = [
|
test_cases = [
|
||||||
("user@domain.com", "test-collection.v1", "doc_user_domain_com_test_collection_v1"),
|
("user@domain.com", "test-collection.v1", "doc_user_domain_com_test_collection_v1_3"),
|
||||||
("user_123", "collection_456", "doc_user_123_collection_456"),
|
("user_123", "collection_456", "doc_user_123_collection_456_3"),
|
||||||
("user with spaces", "collection with spaces", "doc_user_with_spaces_collection_with_spaces"),
|
("user with spaces", "collection with spaces", "doc_user_with_spaces_collection_with_spaces_3"),
|
||||||
("user@@@test", "collection---test", "doc_user_test_collection_test"),
|
("user@@@test", "collection---test", "doc_user_test_collection_test_3"),
|
||||||
]
|
]
|
||||||
|
|
||||||
vector = [0.1, 0.2, 0.3]
|
vector = [0.1, 0.2, 0.3]
|
||||||
|
|
|
||||||
|
|
@ -119,8 +119,8 @@ class TestPineconeDocEmbeddingsQueryProcessor:
|
||||||
|
|
||||||
chunks = await processor.query_document_embeddings(message)
|
chunks = await processor.query_document_embeddings(message)
|
||||||
|
|
||||||
# Verify index was accessed correctly
|
# Verify index was accessed correctly (with dimension suffix)
|
||||||
expected_index_name = "d-test_user-test_collection"
|
expected_index_name = "d-test_user-test_collection-3" # 3 dimensions
|
||||||
processor.pinecone.Index.assert_called_once_with(expected_index_name)
|
processor.pinecone.Index.assert_called_once_with(expected_index_name)
|
||||||
|
|
||||||
# Verify query parameters
|
# Verify query parameters
|
||||||
|
|
@ -264,10 +264,12 @@ class TestPineconeDocEmbeddingsQueryProcessor:
|
||||||
|
|
||||||
chunks = await processor.query_document_embeddings(message)
|
chunks = await processor.query_document_embeddings(message)
|
||||||
|
|
||||||
# Verify same index used for both vectors
|
# Verify different indexes used for different dimensions
|
||||||
expected_index_name = "d-test_user-test_collection"
|
|
||||||
assert processor.pinecone.Index.call_count == 2
|
assert processor.pinecone.Index.call_count == 2
|
||||||
processor.pinecone.Index.assert_called_with(expected_index_name)
|
index_calls = processor.pinecone.Index.call_args_list
|
||||||
|
index_names = [call[0][0] for call in index_calls]
|
||||||
|
assert "d-test_user-test_collection-2" in index_names # 2D vector
|
||||||
|
assert "d-test_user-test_collection-4" in index_names # 4D vector
|
||||||
|
|
||||||
# Verify both queries were made
|
# Verify both queries were made
|
||||||
assert mock_index.query.call_count == 2
|
assert mock_index.query.call_count == 2
|
||||||
|
|
|
||||||
|
|
@ -103,8 +103,8 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase):
|
||||||
result = await processor.query_document_embeddings(mock_message)
|
result = await processor.query_document_embeddings(mock_message)
|
||||||
|
|
||||||
# Assert
|
# Assert
|
||||||
# Verify query was called with correct parameters
|
# Verify query was called with correct parameters (with dimension suffix)
|
||||||
expected_collection = 'd_test_user_test_collection'
|
expected_collection = 'd_test_user_test_collection_3' # 3 dimensions
|
||||||
mock_qdrant_instance.query_points.assert_called_once_with(
|
mock_qdrant_instance.query_points.assert_called_once_with(
|
||||||
collection_name=expected_collection,
|
collection_name=expected_collection,
|
||||||
query=[0.1, 0.2, 0.3],
|
query=[0.1, 0.2, 0.3],
|
||||||
|
|
@ -165,8 +165,8 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase):
|
||||||
# Verify query was called twice
|
# Verify query was called twice
|
||||||
assert mock_qdrant_instance.query_points.call_count == 2
|
assert mock_qdrant_instance.query_points.call_count == 2
|
||||||
|
|
||||||
# Verify both collections were queried
|
# Verify both collections were queried (both 2-dimensional vectors)
|
||||||
expected_collection = 'd_multi_user_multi_collection'
|
expected_collection = 'd_multi_user_multi_collection_2' # 2 dimensions
|
||||||
calls = mock_qdrant_instance.query_points.call_args_list
|
calls = mock_qdrant_instance.query_points.call_args_list
|
||||||
assert calls[0][1]['collection_name'] == expected_collection
|
assert calls[0][1]['collection_name'] == expected_collection
|
||||||
assert calls[1][1]['collection_name'] == expected_collection
|
assert calls[1][1]['collection_name'] == expected_collection
|
||||||
|
|
@ -303,11 +303,11 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase):
|
||||||
calls = mock_qdrant_instance.query_points.call_args_list
|
calls = mock_qdrant_instance.query_points.call_args_list
|
||||||
|
|
||||||
# First call should use 2D collection
|
# First call should use 2D collection
|
||||||
assert calls[0][1]['collection_name'] == 'd_dim_user_dim_collection'
|
assert calls[0][1]['collection_name'] == 'd_dim_user_dim_collection_2' # 2 dimensions
|
||||||
assert calls[0][1]['query'] == [0.1, 0.2]
|
assert calls[0][1]['query'] == [0.1, 0.2]
|
||||||
|
|
||||||
# Second call should use 3D collection
|
# Second call should use 3D collection
|
||||||
assert calls[1][1]['collection_name'] == 'd_dim_user_dim_collection'
|
assert calls[1][1]['collection_name'] == 'd_dim_user_dim_collection_3' # 3 dimensions
|
||||||
assert calls[1][1]['query'] == [0.3, 0.4, 0.5]
|
assert calls[1][1]['query'] == [0.3, 0.4, 0.5]
|
||||||
|
|
||||||
# Verify results
|
# Verify results
|
||||||
|
|
|
||||||
|
|
@ -147,8 +147,8 @@ class TestPineconeGraphEmbeddingsQueryProcessor:
|
||||||
|
|
||||||
entities = await processor.query_graph_embeddings(message)
|
entities = await processor.query_graph_embeddings(message)
|
||||||
|
|
||||||
# Verify index was accessed correctly
|
# Verify index was accessed correctly (with dimension suffix)
|
||||||
expected_index_name = "t-test_user-test_collection"
|
expected_index_name = "t-test_user-test_collection-3" # 3 dimensions
|
||||||
processor.pinecone.Index.assert_called_once_with(expected_index_name)
|
processor.pinecone.Index.assert_called_once_with(expected_index_name)
|
||||||
|
|
||||||
# Verify query parameters
|
# Verify query parameters
|
||||||
|
|
@ -290,10 +290,12 @@ class TestPineconeGraphEmbeddingsQueryProcessor:
|
||||||
|
|
||||||
entities = await processor.query_graph_embeddings(message)
|
entities = await processor.query_graph_embeddings(message)
|
||||||
|
|
||||||
# Verify same index used for both vectors
|
# Verify different indexes used for different dimensions
|
||||||
expected_index_name = "t-test_user-test_collection"
|
|
||||||
assert processor.pinecone.Index.call_count == 2
|
assert processor.pinecone.Index.call_count == 2
|
||||||
processor.pinecone.Index.assert_called_with(expected_index_name)
|
index_calls = processor.pinecone.Index.call_args_list
|
||||||
|
index_names = [call[0][0] for call in index_calls]
|
||||||
|
assert "t-test_user-test_collection-2" in index_names # 2D vector
|
||||||
|
assert "t-test_user-test_collection-4" in index_names # 4D vector
|
||||||
|
|
||||||
# Verify both queries were made
|
# Verify both queries were made
|
||||||
assert mock_index.query.call_count == 2
|
assert mock_index.query.call_count == 2
|
||||||
|
|
|
||||||
|
|
@ -175,8 +175,8 @@ class TestQdrantGraphEmbeddingsQuery(IsolatedAsyncioTestCase):
|
||||||
result = await processor.query_graph_embeddings(mock_message)
|
result = await processor.query_graph_embeddings(mock_message)
|
||||||
|
|
||||||
# Assert
|
# Assert
|
||||||
# Verify query was called with correct parameters
|
# Verify query was called with correct parameters (with dimension suffix)
|
||||||
expected_collection = 't_test_user_test_collection'
|
expected_collection = 't_test_user_test_collection_3' # 3 dimensions
|
||||||
mock_qdrant_instance.query_points.assert_called_once_with(
|
mock_qdrant_instance.query_points.assert_called_once_with(
|
||||||
collection_name=expected_collection,
|
collection_name=expected_collection,
|
||||||
query=[0.1, 0.2, 0.3],
|
query=[0.1, 0.2, 0.3],
|
||||||
|
|
@ -235,8 +235,8 @@ class TestQdrantGraphEmbeddingsQuery(IsolatedAsyncioTestCase):
|
||||||
# Verify query was called twice
|
# Verify query was called twice
|
||||||
assert mock_qdrant_instance.query_points.call_count == 2
|
assert mock_qdrant_instance.query_points.call_count == 2
|
||||||
|
|
||||||
# Verify both collections were queried
|
# Verify both collections were queried (both 2-dimensional vectors)
|
||||||
expected_collection = 't_multi_user_multi_collection'
|
expected_collection = 't_multi_user_multi_collection_2' # 2 dimensions
|
||||||
calls = mock_qdrant_instance.query_points.call_args_list
|
calls = mock_qdrant_instance.query_points.call_args_list
|
||||||
assert calls[0][1]['collection_name'] == expected_collection
|
assert calls[0][1]['collection_name'] == expected_collection
|
||||||
assert calls[1][1]['collection_name'] == expected_collection
|
assert calls[1][1]['collection_name'] == expected_collection
|
||||||
|
|
@ -374,11 +374,11 @@ class TestQdrantGraphEmbeddingsQuery(IsolatedAsyncioTestCase):
|
||||||
calls = mock_qdrant_instance.query_points.call_args_list
|
calls = mock_qdrant_instance.query_points.call_args_list
|
||||||
|
|
||||||
# First call should use 2D collection
|
# First call should use 2D collection
|
||||||
assert calls[0][1]['collection_name'] == 't_dim_user_dim_collection'
|
assert calls[0][1]['collection_name'] == 't_dim_user_dim_collection_2' # 2 dimensions
|
||||||
assert calls[0][1]['query'] == [0.1, 0.2]
|
assert calls[0][1]['query'] == [0.1, 0.2]
|
||||||
|
|
||||||
# Second call should use 3D collection
|
# Second call should use 3D collection
|
||||||
assert calls[1][1]['collection_name'] == 't_dim_user_dim_collection'
|
assert calls[1][1]['collection_name'] == 't_dim_user_dim_collection_3' # 3 dimensions
|
||||||
assert calls[1][1]['query'] == [0.3, 0.4, 0.5]
|
assert calls[1][1]['query'] == [0.3, 0.4, 0.5]
|
||||||
|
|
||||||
# Verify results
|
# Verify results
|
||||||
|
|
|
||||||
|
|
@ -134,8 +134,8 @@ class TestPineconeDocEmbeddingsStorageProcessor:
|
||||||
with patch('uuid.uuid4', side_effect=['id1', 'id2']):
|
with patch('uuid.uuid4', side_effect=['id1', 'id2']):
|
||||||
await processor.store_document_embeddings(message)
|
await processor.store_document_embeddings(message)
|
||||||
|
|
||||||
# Verify index name and operations
|
# Verify index name and operations (with dimension suffix)
|
||||||
expected_index_name = "d-test_user-test_collection"
|
expected_index_name = "d-test_user-test_collection-3" # 3 dimensions
|
||||||
processor.pinecone.Index.assert_called_with(expected_index_name)
|
processor.pinecone.Index.assert_called_with(expected_index_name)
|
||||||
|
|
||||||
# Verify upsert was called for each vector
|
# Verify upsert was called for each vector
|
||||||
|
|
@ -179,7 +179,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_document_embeddings_index_validation(self, processor):
|
async def test_store_document_embeddings_index_validation(self, processor):
|
||||||
"""Test that writing to non-existent index raises ValueError"""
|
"""Test that writing to non-existent index creates it lazily"""
|
||||||
message = MagicMock()
|
message = MagicMock()
|
||||||
message.metadata = MagicMock()
|
message.metadata = MagicMock()
|
||||||
message.metadata.user = 'test_user'
|
message.metadata.user = 'test_user'
|
||||||
|
|
@ -191,12 +191,24 @@ class TestPineconeDocEmbeddingsStorageProcessor:
|
||||||
)
|
)
|
||||||
message.chunks = [chunk]
|
message.chunks = [chunk]
|
||||||
|
|
||||||
# Mock index doesn't exist
|
# Mock index doesn't exist initially
|
||||||
processor.pinecone.has_index.return_value = False
|
processor.pinecone.has_index.return_value = False
|
||||||
|
mock_index = MagicMock()
|
||||||
|
processor.pinecone.Index.return_value = mock_index
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
with patch('uuid.uuid4', return_value='test-id'):
|
||||||
await processor.store_document_embeddings(message)
|
await processor.store_document_embeddings(message)
|
||||||
|
|
||||||
|
# Verify index was created with correct dimension
|
||||||
|
expected_index_name = "d-test_user-test_collection-3" # 3 dimensions
|
||||||
|
processor.pinecone.create_index.assert_called_once()
|
||||||
|
create_call = processor.pinecone.create_index.call_args
|
||||||
|
assert create_call[1]['name'] == expected_index_name
|
||||||
|
assert create_call[1]['dimension'] == 3
|
||||||
|
|
||||||
|
# Verify upsert was still called
|
||||||
|
mock_index.upsert.assert_called_once()
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_document_embeddings_empty_chunk(self, processor):
|
async def test_store_document_embeddings_empty_chunk(self, processor):
|
||||||
"""Test storing document embeddings with empty chunk (should be skipped)"""
|
"""Test storing document embeddings with empty chunk (should be skipped)"""
|
||||||
|
|
@ -345,7 +357,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_document_embeddings_validation_before_creation(self, processor):
|
async def test_store_document_embeddings_validation_before_creation(self, processor):
|
||||||
"""Test that validation error occurs before creation attempts"""
|
"""Test that lazy creation happens when index doesn't exist"""
|
||||||
message = MagicMock()
|
message = MagicMock()
|
||||||
message.metadata = MagicMock()
|
message.metadata = MagicMock()
|
||||||
message.metadata.user = 'test_user'
|
message.metadata.user = 'test_user'
|
||||||
|
|
@ -359,13 +371,18 @@ class TestPineconeDocEmbeddingsStorageProcessor:
|
||||||
|
|
||||||
# Mock index doesn't exist
|
# Mock index doesn't exist
|
||||||
processor.pinecone.has_index.return_value = False
|
processor.pinecone.has_index.return_value = False
|
||||||
|
mock_index = MagicMock()
|
||||||
|
processor.pinecone.Index.return_value = mock_index
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
with patch('uuid.uuid4', return_value='test-id'):
|
||||||
await processor.store_document_embeddings(message)
|
await processor.store_document_embeddings(message)
|
||||||
|
|
||||||
|
# Verify index was created
|
||||||
|
processor.pinecone.create_index.assert_called_once()
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_document_embeddings_validates_before_timeout(self, processor):
|
async def test_store_document_embeddings_validates_before_timeout(self, processor):
|
||||||
"""Test that validation error occurs before timeout checks"""
|
"""Test that lazy creation works correctly"""
|
||||||
message = MagicMock()
|
message = MagicMock()
|
||||||
message.metadata = MagicMock()
|
message.metadata = MagicMock()
|
||||||
message.metadata.user = 'test_user'
|
message.metadata.user = 'test_user'
|
||||||
|
|
@ -379,10 +396,16 @@ class TestPineconeDocEmbeddingsStorageProcessor:
|
||||||
|
|
||||||
# Mock index doesn't exist
|
# Mock index doesn't exist
|
||||||
processor.pinecone.has_index.return_value = False
|
processor.pinecone.has_index.return_value = False
|
||||||
|
mock_index = MagicMock()
|
||||||
|
processor.pinecone.Index.return_value = mock_index
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
with patch('uuid.uuid4', return_value='test-id'):
|
||||||
await processor.store_document_embeddings(message)
|
await processor.store_document_embeddings(message)
|
||||||
|
|
||||||
|
# Verify index was created and used
|
||||||
|
processor.pinecone.create_index.assert_called_once()
|
||||||
|
mock_index.upsert.assert_called_once()
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_document_embeddings_unicode_content(self, processor):
|
async def test_store_document_embeddings_unicode_content(self, processor):
|
||||||
"""Test storing document embeddings with Unicode content"""
|
"""Test storing document embeddings with Unicode content"""
|
||||||
|
|
|
||||||
|
|
@ -103,8 +103,8 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
await processor.store_document_embeddings(mock_message)
|
await processor.store_document_embeddings(mock_message)
|
||||||
|
|
||||||
# Assert
|
# Assert
|
||||||
# Verify collection existence was checked
|
# Verify collection existence was checked (with dimension suffix)
|
||||||
expected_collection = 'd_test_user_test_collection'
|
expected_collection = 'd_test_user_test_collection_3' # 3 dimensions in vector [0.1, 0.2, 0.3]
|
||||||
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
|
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
|
||||||
|
|
||||||
# Verify upsert was called
|
# Verify upsert was called
|
||||||
|
|
@ -112,7 +112,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
# Verify upsert parameters
|
# Verify upsert parameters
|
||||||
upsert_call_args = mock_qdrant_instance.upsert.call_args
|
upsert_call_args = mock_qdrant_instance.upsert.call_args
|
||||||
assert upsert_call_args[1]['collection_name'] == expected_collection
|
assert upsert_call_args[1]['collection_name'] == 'd_test_user_test_collection_3'
|
||||||
assert len(upsert_call_args[1]['points']) == 1
|
assert len(upsert_call_args[1]['points']) == 1
|
||||||
|
|
||||||
point = upsert_call_args[1]['points'][0]
|
point = upsert_call_args[1]['points'][0]
|
||||||
|
|
@ -272,18 +272,21 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
# Assert
|
# Assert
|
||||||
# Should not call upsert for empty chunks
|
# Should not call upsert for empty chunks
|
||||||
mock_qdrant_instance.upsert.assert_not_called()
|
mock_qdrant_instance.upsert.assert_not_called()
|
||||||
# But collection_exists should be called for validation
|
# collection_exists should NOT be called since we return early for empty chunks
|
||||||
mock_qdrant_instance.collection_exists.assert_called_once()
|
mock_qdrant_instance.collection_exists.assert_not_called()
|
||||||
|
|
||||||
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
||||||
|
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
|
||||||
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
|
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
|
||||||
async def test_collection_creation_when_not_exists(self, mock_base_init, mock_qdrant_client):
|
async def test_collection_creation_when_not_exists(self, mock_base_init, mock_uuid, mock_qdrant_client):
|
||||||
"""Test that writing to non-existent collection raises ValueError"""
|
"""Test that writing to non-existent collection creates it lazily"""
|
||||||
# Arrange
|
# Arrange
|
||||||
mock_base_init.return_value = None
|
mock_base_init.return_value = None
|
||||||
mock_qdrant_instance = MagicMock()
|
mock_qdrant_instance = MagicMock()
|
||||||
mock_qdrant_instance.collection_exists.return_value = False # Collection doesn't exist
|
mock_qdrant_instance.collection_exists.return_value = False # Collection doesn't exist
|
||||||
mock_qdrant_client.return_value = mock_qdrant_instance
|
mock_qdrant_client.return_value = mock_qdrant_instance
|
||||||
|
mock_uuid.uuid4.return_value = MagicMock()
|
||||||
|
mock_uuid.uuid4.return_value.__str__ = MagicMock(return_value='test-uuid')
|
||||||
|
|
||||||
config = {
|
config = {
|
||||||
'store_uri': 'http://localhost:6333',
|
'store_uri': 'http://localhost:6333',
|
||||||
|
|
@ -305,19 +308,36 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
mock_message.chunks = [mock_chunk]
|
mock_message.chunks = [mock_chunk]
|
||||||
|
|
||||||
# Act & Assert
|
# Act
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
await processor.store_document_embeddings(mock_message)
|
||||||
await processor.store_document_embeddings(mock_message)
|
|
||||||
|
# Assert - collection should be lazily created
|
||||||
|
expected_collection = 'd_new_user_new_collection_5' # 5 dimensions
|
||||||
|
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
|
||||||
|
mock_qdrant_instance.create_collection.assert_called_once()
|
||||||
|
|
||||||
|
# Verify create_collection was called with correct parameters
|
||||||
|
create_call = mock_qdrant_instance.create_collection.call_args
|
||||||
|
assert create_call[1]['collection_name'] == expected_collection
|
||||||
|
assert create_call[1]['vectors_config'].size == 5
|
||||||
|
|
||||||
|
# Verify upsert was still called
|
||||||
|
mock_qdrant_instance.upsert.assert_called_once()
|
||||||
|
|
||||||
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
||||||
|
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
|
||||||
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
|
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
|
||||||
async def test_collection_creation_exception(self, mock_base_init, mock_qdrant_client):
|
async def test_collection_creation_exception(self, mock_base_init, mock_uuid, mock_qdrant_client):
|
||||||
"""Test that validation error occurs before connection errors"""
|
"""Test that collection creation errors are propagated"""
|
||||||
# Arrange
|
# Arrange
|
||||||
mock_base_init.return_value = None
|
mock_base_init.return_value = None
|
||||||
mock_qdrant_instance = MagicMock()
|
mock_qdrant_instance = MagicMock()
|
||||||
mock_qdrant_instance.collection_exists.return_value = False # Collection doesn't exist
|
mock_qdrant_instance.collection_exists.return_value = False # Collection doesn't exist
|
||||||
|
# Simulate creation failure
|
||||||
|
mock_qdrant_instance.create_collection.side_effect = Exception("Connection error")
|
||||||
mock_qdrant_client.return_value = mock_qdrant_instance
|
mock_qdrant_client.return_value = mock_qdrant_instance
|
||||||
|
mock_uuid.uuid4.return_value = MagicMock()
|
||||||
|
mock_uuid.uuid4.return_value.__str__ = MagicMock(return_value='test-uuid')
|
||||||
|
|
||||||
config = {
|
config = {
|
||||||
'store_uri': 'http://localhost:6333',
|
'store_uri': 'http://localhost:6333',
|
||||||
|
|
@ -339,8 +359,8 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
mock_message.chunks = [mock_chunk]
|
mock_message.chunks = [mock_chunk]
|
||||||
|
|
||||||
# Act & Assert
|
# Act & Assert - should propagate the creation error
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
with pytest.raises(Exception, match="Connection error"):
|
||||||
await processor.store_document_embeddings(mock_message)
|
await processor.store_document_embeddings(mock_message)
|
||||||
|
|
||||||
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
||||||
|
|
@ -398,7 +418,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
await processor.store_document_embeddings(mock_message2)
|
await processor.store_document_embeddings(mock_message2)
|
||||||
|
|
||||||
# Assert
|
# Assert
|
||||||
expected_collection = 'd_cache_user_cache_collection'
|
expected_collection = 'd_cache_user_cache_collection_3' # 3 dimensions
|
||||||
|
|
||||||
# Verify collection existence is checked on each write
|
# Verify collection existence is checked on each write
|
||||||
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
|
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
|
||||||
|
|
@ -407,14 +427,17 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
mock_qdrant_instance.upsert.assert_called_once()
|
mock_qdrant_instance.upsert.assert_called_once()
|
||||||
|
|
||||||
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
||||||
|
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
|
||||||
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
|
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
|
||||||
async def test_different_dimensions_different_collections(self, mock_base_init, mock_qdrant_client):
|
async def test_different_dimensions_different_collections(self, mock_base_init, mock_uuid, mock_qdrant_client):
|
||||||
"""Test that different vector dimensions create different collections"""
|
"""Test that different vector dimensions create different collections"""
|
||||||
# Arrange
|
# Arrange
|
||||||
mock_base_init.return_value = None
|
mock_base_init.return_value = None
|
||||||
mock_qdrant_instance = MagicMock()
|
mock_qdrant_instance = MagicMock()
|
||||||
mock_qdrant_instance.collection_exists.return_value = True
|
mock_qdrant_instance.collection_exists.return_value = True
|
||||||
mock_qdrant_client.return_value = mock_qdrant_instance
|
mock_qdrant_client.return_value = mock_qdrant_instance
|
||||||
|
mock_uuid.uuid4.return_value = MagicMock()
|
||||||
|
mock_uuid.uuid4.return_value.__str__ = MagicMock(return_value='test-uuid')
|
||||||
|
|
||||||
config = {
|
config = {
|
||||||
'store_uri': 'http://localhost:6333',
|
'store_uri': 'http://localhost:6333',
|
||||||
|
|
@ -443,16 +466,20 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
await processor.store_document_embeddings(mock_message)
|
await processor.store_document_embeddings(mock_message)
|
||||||
|
|
||||||
# Assert
|
# Assert
|
||||||
# Should check existence of the same collection (dimensions no longer create separate collections)
|
# Should check existence of DIFFERENT collections for each dimension
|
||||||
expected_collection = 'd_dim_user_dim_collection'
|
assert mock_qdrant_instance.collection_exists.call_count == 2
|
||||||
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
|
|
||||||
|
|
||||||
# Should upsert to the same collection for both vectors
|
# Verify the two different collection names were checked
|
||||||
|
collection_exists_calls = [call[0][0] for call in mock_qdrant_instance.collection_exists.call_args_list]
|
||||||
|
assert 'd_dim_user_dim_collection_2' in collection_exists_calls # 2-dim vector
|
||||||
|
assert 'd_dim_user_dim_collection_3' in collection_exists_calls # 3-dim vector
|
||||||
|
|
||||||
|
# Should upsert to different collections for each vector
|
||||||
assert mock_qdrant_instance.upsert.call_count == 2
|
assert mock_qdrant_instance.upsert.call_count == 2
|
||||||
|
|
||||||
upsert_calls = mock_qdrant_instance.upsert.call_args_list
|
upsert_calls = mock_qdrant_instance.upsert.call_args_list
|
||||||
assert upsert_calls[0][1]['collection_name'] == expected_collection
|
assert upsert_calls[0][1]['collection_name'] == 'd_dim_user_dim_collection_2'
|
||||||
assert upsert_calls[1][1]['collection_name'] == expected_collection
|
assert upsert_calls[1][1]['collection_name'] == 'd_dim_user_dim_collection_3'
|
||||||
|
|
||||||
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
|
||||||
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
|
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
|
||||||
|
|
|
||||||
|
|
@ -134,8 +134,8 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
|
||||||
with patch('uuid.uuid4', side_effect=['id1', 'id2']):
|
with patch('uuid.uuid4', side_effect=['id1', 'id2']):
|
||||||
await processor.store_graph_embeddings(message)
|
await processor.store_graph_embeddings(message)
|
||||||
|
|
||||||
# Verify index name and operations
|
# Verify index name and operations (with dimension suffix)
|
||||||
expected_index_name = "t-test_user-test_collection"
|
expected_index_name = "t-test_user-test_collection-3" # 3 dimensions
|
||||||
processor.pinecone.Index.assert_called_with(expected_index_name)
|
processor.pinecone.Index.assert_called_with(expected_index_name)
|
||||||
|
|
||||||
# Verify upsert was called for each vector
|
# Verify upsert was called for each vector
|
||||||
|
|
@ -179,7 +179,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_graph_embeddings_index_validation(self, processor):
|
async def test_store_graph_embeddings_index_validation(self, processor):
|
||||||
"""Test that writing to non-existent index raises ValueError"""
|
"""Test that writing to non-existent index creates it lazily"""
|
||||||
message = MagicMock()
|
message = MagicMock()
|
||||||
message.metadata = MagicMock()
|
message.metadata = MagicMock()
|
||||||
message.metadata.user = 'test_user'
|
message.metadata.user = 'test_user'
|
||||||
|
|
@ -193,10 +193,22 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
|
||||||
|
|
||||||
# Mock index doesn't exist
|
# Mock index doesn't exist
|
||||||
processor.pinecone.has_index.return_value = False
|
processor.pinecone.has_index.return_value = False
|
||||||
|
mock_index = MagicMock()
|
||||||
|
processor.pinecone.Index.return_value = mock_index
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
with patch('uuid.uuid4', return_value='test-id'):
|
||||||
await processor.store_graph_embeddings(message)
|
await processor.store_graph_embeddings(message)
|
||||||
|
|
||||||
|
# Verify index was created with correct dimension
|
||||||
|
expected_index_name = "t-test_user-test_collection-3" # 3 dimensions
|
||||||
|
processor.pinecone.create_index.assert_called_once()
|
||||||
|
create_call = processor.pinecone.create_index.call_args
|
||||||
|
assert create_call[1]['name'] == expected_index_name
|
||||||
|
assert create_call[1]['dimension'] == 3
|
||||||
|
|
||||||
|
# Verify upsert was still called
|
||||||
|
mock_index.upsert.assert_called_once()
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_graph_embeddings_empty_entity_value(self, processor):
|
async def test_store_graph_embeddings_empty_entity_value(self, processor):
|
||||||
"""Test storing graph embeddings with empty entity value (should be skipped)"""
|
"""Test storing graph embeddings with empty entity value (should be skipped)"""
|
||||||
|
|
@ -267,11 +279,16 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
|
||||||
with patch('uuid.uuid4', side_effect=['id1', 'id2', 'id3']):
|
with patch('uuid.uuid4', side_effect=['id1', 'id2', 'id3']):
|
||||||
await processor.store_graph_embeddings(message)
|
await processor.store_graph_embeddings(message)
|
||||||
|
|
||||||
# Verify same index was used for all dimensions
|
# Verify different indexes were used for different dimensions
|
||||||
expected_index_name = 't-test_user-test_collection'
|
index_calls = processor.pinecone.Index.call_args_list
|
||||||
processor.pinecone.Index.assert_called_with(expected_index_name)
|
assert len(index_calls) == 3
|
||||||
|
# Extract index names from calls
|
||||||
|
index_names = [call[0][0] for call in index_calls]
|
||||||
|
assert 't-test_user-test_collection-2' in index_names # 2D vector
|
||||||
|
assert 't-test_user-test_collection-4' in index_names # 4D vector
|
||||||
|
assert 't-test_user-test_collection-3' in index_names # 3D vector
|
||||||
|
|
||||||
# Verify all vectors were upserted to the same index
|
# Verify all vectors were upserted (to their respective indexes)
|
||||||
assert mock_index.upsert.call_count == 3
|
assert mock_index.upsert.call_count == 3
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
@ -316,7 +333,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_graph_embeddings_validation_before_creation(self, processor):
|
async def test_store_graph_embeddings_validation_before_creation(self, processor):
|
||||||
"""Test that validation error occurs before any creation attempts"""
|
"""Test that lazy creation happens when index doesn't exist"""
|
||||||
message = MagicMock()
|
message = MagicMock()
|
||||||
message.metadata = MagicMock()
|
message.metadata = MagicMock()
|
||||||
message.metadata.user = 'test_user'
|
message.metadata.user = 'test_user'
|
||||||
|
|
@ -330,13 +347,18 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
|
||||||
|
|
||||||
# Mock index doesn't exist
|
# Mock index doesn't exist
|
||||||
processor.pinecone.has_index.return_value = False
|
processor.pinecone.has_index.return_value = False
|
||||||
|
mock_index = MagicMock()
|
||||||
|
processor.pinecone.Index.return_value = mock_index
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
with patch('uuid.uuid4', return_value='test-id'):
|
||||||
await processor.store_graph_embeddings(message)
|
await processor.store_graph_embeddings(message)
|
||||||
|
|
||||||
|
# Verify index was created
|
||||||
|
processor.pinecone.create_index.assert_called_once()
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_graph_embeddings_validates_before_timeout(self, processor):
|
async def test_store_graph_embeddings_validates_before_timeout(self, processor):
|
||||||
"""Test that validation error occurs before timeout checks"""
|
"""Test that lazy creation works correctly"""
|
||||||
message = MagicMock()
|
message = MagicMock()
|
||||||
message.metadata = MagicMock()
|
message.metadata = MagicMock()
|
||||||
message.metadata.user = 'test_user'
|
message.metadata.user = 'test_user'
|
||||||
|
|
@ -350,10 +372,16 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
|
||||||
|
|
||||||
# Mock index doesn't exist
|
# Mock index doesn't exist
|
||||||
processor.pinecone.has_index.return_value = False
|
processor.pinecone.has_index.return_value = False
|
||||||
|
mock_index = MagicMock()
|
||||||
|
processor.pinecone.Index.return_value = mock_index
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
with patch('uuid.uuid4', return_value='test-id'):
|
||||||
await processor.store_graph_embeddings(message)
|
await processor.store_graph_embeddings(message)
|
||||||
|
|
||||||
|
# Verify index was created and used
|
||||||
|
processor.pinecone.create_index.assert_called_once()
|
||||||
|
mock_index.upsert.assert_called_once()
|
||||||
|
|
||||||
def test_add_args_method(self):
|
def test_add_args_method(self):
|
||||||
"""Test that add_args properly configures argument parser"""
|
"""Test that add_args properly configures argument parser"""
|
||||||
from argparse import ArgumentParser
|
from argparse import ArgumentParser
|
||||||
|
|
|
||||||
|
|
@ -44,29 +44,6 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
assert hasattr(processor, 'qdrant')
|
assert hasattr(processor, 'qdrant')
|
||||||
assert processor.qdrant == mock_qdrant_instance
|
assert processor.qdrant == mock_qdrant_instance
|
||||||
|
|
||||||
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
|
|
||||||
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
|
|
||||||
async def test_get_collection_validates_existence(self, mock_base_init, mock_qdrant_client):
|
|
||||||
"""Test get_collection validates that collection exists"""
|
|
||||||
# Arrange
|
|
||||||
mock_base_init.return_value = None
|
|
||||||
mock_qdrant_instance = MagicMock()
|
|
||||||
mock_qdrant_instance.collection_exists.return_value = False
|
|
||||||
mock_qdrant_client.return_value = mock_qdrant_instance
|
|
||||||
|
|
||||||
config = {
|
|
||||||
'store_uri': 'http://localhost:6333',
|
|
||||||
'api_key': 'test-api-key',
|
|
||||||
'taskgroup': AsyncMock(),
|
|
||||||
'id': 'test-qdrant-processor'
|
|
||||||
}
|
|
||||||
|
|
||||||
processor = Processor(**config)
|
|
||||||
|
|
||||||
# Act & Assert
|
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
|
||||||
processor.get_collection(user='test_user', collection='test_collection')
|
|
||||||
|
|
||||||
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
|
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
|
||||||
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')
|
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')
|
||||||
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
|
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
|
||||||
|
|
@ -103,8 +80,8 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
await processor.store_graph_embeddings(mock_message)
|
await processor.store_graph_embeddings(mock_message)
|
||||||
|
|
||||||
# Assert
|
# Assert
|
||||||
# Verify collection existence was checked
|
# Verify collection existence was checked (with dimension suffix)
|
||||||
expected_collection = 't_test_user_test_collection'
|
expected_collection = 't_test_user_test_collection_3' # 3 dimensions in vector [0.1, 0.2, 0.3]
|
||||||
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
|
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
|
||||||
|
|
||||||
# Verify upsert was called
|
# Verify upsert was called
|
||||||
|
|
@ -112,105 +89,13 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
# Verify upsert parameters
|
# Verify upsert parameters
|
||||||
upsert_call_args = mock_qdrant_instance.upsert.call_args
|
upsert_call_args = mock_qdrant_instance.upsert.call_args
|
||||||
assert upsert_call_args[1]['collection_name'] == expected_collection
|
assert upsert_call_args[1]['collection_name'] == 't_test_user_test_collection_3'
|
||||||
assert len(upsert_call_args[1]['points']) == 1
|
assert len(upsert_call_args[1]['points']) == 1
|
||||||
|
|
||||||
point = upsert_call_args[1]['points'][0]
|
point = upsert_call_args[1]['points'][0]
|
||||||
assert point.vector == [0.1, 0.2, 0.3]
|
assert point.vector == [0.1, 0.2, 0.3]
|
||||||
assert point.payload['entity'] == 'test_entity'
|
assert point.payload['entity'] == 'test_entity'
|
||||||
|
|
||||||
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
|
|
||||||
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
|
|
||||||
async def test_get_collection_uses_existing_collection(self, mock_base_init, mock_qdrant_client):
|
|
||||||
"""Test get_collection uses existing collection without creating new one"""
|
|
||||||
# Arrange
|
|
||||||
mock_base_init.return_value = None
|
|
||||||
mock_qdrant_instance = MagicMock()
|
|
||||||
mock_qdrant_instance.collection_exists.return_value = True # Collection exists
|
|
||||||
mock_qdrant_client.return_value = mock_qdrant_instance
|
|
||||||
|
|
||||||
config = {
|
|
||||||
'store_uri': 'http://localhost:6333',
|
|
||||||
'api_key': 'test-api-key',
|
|
||||||
'taskgroup': AsyncMock(),
|
|
||||||
'id': 'test-qdrant-processor'
|
|
||||||
}
|
|
||||||
|
|
||||||
processor = Processor(**config)
|
|
||||||
|
|
||||||
# Act
|
|
||||||
collection_name = processor.get_collection(user='existing_user', collection='existing_collection')
|
|
||||||
|
|
||||||
# Assert
|
|
||||||
expected_name = 't_existing_user_existing_collection'
|
|
||||||
assert collection_name == expected_name
|
|
||||||
|
|
||||||
# Verify collection existence check was performed
|
|
||||||
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_name)
|
|
||||||
# Verify create_collection was NOT called
|
|
||||||
mock_qdrant_instance.create_collection.assert_not_called()
|
|
||||||
|
|
||||||
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
|
|
||||||
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
|
|
||||||
async def test_get_collection_validates_on_each_call(self, mock_base_init, mock_qdrant_client):
|
|
||||||
"""Test get_collection validates collection existence on each call"""
|
|
||||||
# Arrange
|
|
||||||
mock_base_init.return_value = None
|
|
||||||
mock_qdrant_instance = MagicMock()
|
|
||||||
mock_qdrant_instance.collection_exists.return_value = True
|
|
||||||
mock_qdrant_client.return_value = mock_qdrant_instance
|
|
||||||
|
|
||||||
config = {
|
|
||||||
'store_uri': 'http://localhost:6333',
|
|
||||||
'api_key': 'test-api-key',
|
|
||||||
'taskgroup': AsyncMock(),
|
|
||||||
'id': 'test-qdrant-processor'
|
|
||||||
}
|
|
||||||
|
|
||||||
processor = Processor(**config)
|
|
||||||
|
|
||||||
# First call
|
|
||||||
collection_name1 = processor.get_collection(user='cache_user', collection='cache_collection')
|
|
||||||
|
|
||||||
# Reset mock to track second call
|
|
||||||
mock_qdrant_instance.reset_mock()
|
|
||||||
mock_qdrant_instance.collection_exists.return_value = True
|
|
||||||
|
|
||||||
# Act - Second call with same parameters
|
|
||||||
collection_name2 = processor.get_collection(user='cache_user', collection='cache_collection')
|
|
||||||
|
|
||||||
# Assert
|
|
||||||
expected_name = 't_cache_user_cache_collection'
|
|
||||||
assert collection_name1 == expected_name
|
|
||||||
assert collection_name2 == expected_name
|
|
||||||
|
|
||||||
# Verify collection existence check happens on each call
|
|
||||||
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_name)
|
|
||||||
mock_qdrant_instance.create_collection.assert_not_called()
|
|
||||||
|
|
||||||
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
|
|
||||||
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
|
|
||||||
async def test_get_collection_creation_exception(self, mock_base_init, mock_qdrant_client):
|
|
||||||
"""Test get_collection raises ValueError when collection doesn't exist"""
|
|
||||||
# Arrange
|
|
||||||
mock_base_init.return_value = None
|
|
||||||
mock_qdrant_instance = MagicMock()
|
|
||||||
mock_qdrant_instance.collection_exists.return_value = False
|
|
||||||
mock_qdrant_client.return_value = mock_qdrant_instance
|
|
||||||
|
|
||||||
config = {
|
|
||||||
'store_uri': 'http://localhost:6333',
|
|
||||||
'api_key': 'test-api-key',
|
|
||||||
'taskgroup': AsyncMock(),
|
|
||||||
'id': 'test-qdrant-processor'
|
|
||||||
}
|
|
||||||
|
|
||||||
processor = Processor(**config)
|
|
||||||
|
|
||||||
# Act & Assert
|
|
||||||
with pytest.raises(ValueError, match="Collection .* does not exist"):
|
|
||||||
processor.get_collection(user='error_user', collection='error_collection')
|
|
||||||
|
|
||||||
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
|
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
|
||||||
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')
|
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')
|
||||||
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
|
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
|
||||||
|
|
|
||||||
|
|
@ -50,24 +50,26 @@ class DocVectors:
|
||||||
logger.debug(f"Reload at {self.next_reload}")
|
logger.debug(f"Reload at {self.next_reload}")
|
||||||
|
|
||||||
def collection_exists(self, user, collection):
|
def collection_exists(self, user, collection):
|
||||||
"""Check if collection exists (dimension-independent check)"""
|
"""
|
||||||
collection_name = make_safe_collection_name(user, collection, self.prefix)
|
Check if any collection exists for this user/collection combination.
|
||||||
return self.client.has_collection(collection_name)
|
Since collections are dimension-specific, this checks if ANY dimension variant exists.
|
||||||
|
"""
|
||||||
|
base_name = make_safe_collection_name(user, collection, self.prefix)
|
||||||
|
prefix = f"{base_name}_"
|
||||||
|
all_collections = self.client.list_collections()
|
||||||
|
return any(coll.startswith(prefix) for coll in all_collections)
|
||||||
|
|
||||||
def create_collection(self, user, collection, dimension=384):
|
def create_collection(self, user, collection, dimension=384):
|
||||||
"""Create collection with default dimension"""
|
"""
|
||||||
collection_name = make_safe_collection_name(user, collection, self.prefix)
|
No-op for explicit collection creation.
|
||||||
|
Collections are created lazily on first insert with actual dimension.
|
||||||
if self.client.has_collection(collection_name):
|
"""
|
||||||
logger.info(f"Collection {collection_name} already exists")
|
logger.info(f"Collection creation requested for {user}/{collection} - will be created lazily on first insert")
|
||||||
return
|
|
||||||
|
|
||||||
self.init_collection(dimension, user, collection)
|
|
||||||
logger.info(f"Created Milvus collection {collection_name} with dimension {dimension}")
|
|
||||||
|
|
||||||
def init_collection(self, dimension, user, collection):
|
def init_collection(self, dimension, user, collection):
|
||||||
|
|
||||||
collection_name = make_safe_collection_name(user, collection, self.prefix)
|
base_name = make_safe_collection_name(user, collection, self.prefix)
|
||||||
|
collection_name = f"{base_name}_{dimension}"
|
||||||
|
|
||||||
pkey_field = FieldSchema(
|
pkey_field = FieldSchema(
|
||||||
name="id",
|
name="id",
|
||||||
|
|
@ -115,6 +117,7 @@ class DocVectors:
|
||||||
)
|
)
|
||||||
|
|
||||||
self.collections[(dimension, user, collection)] = collection_name
|
self.collections[(dimension, user, collection)] = collection_name
|
||||||
|
logger.info(f"Created Milvus collection {collection_name} with dimension {dimension}")
|
||||||
|
|
||||||
def insert(self, embeds, doc, user, collection):
|
def insert(self, embeds, doc, user, collection):
|
||||||
|
|
||||||
|
|
@ -139,8 +142,15 @@ class DocVectors:
|
||||||
|
|
||||||
dim = len(embeds)
|
dim = len(embeds)
|
||||||
|
|
||||||
|
# Check if collection exists - return empty if not
|
||||||
if (dim, user, collection) not in self.collections:
|
if (dim, user, collection) not in self.collections:
|
||||||
self.init_collection(dim, user, collection)
|
base_name = make_safe_collection_name(user, collection, self.prefix)
|
||||||
|
collection_name = f"{base_name}_{dim}"
|
||||||
|
if not self.client.has_collection(collection_name):
|
||||||
|
logger.info(f"Collection {collection_name} does not exist, returning empty results")
|
||||||
|
return []
|
||||||
|
# Collection exists but not in cache, add it
|
||||||
|
self.collections[(dim, user, collection)] = collection_name
|
||||||
|
|
||||||
coll = self.collections[(dim, user, collection)]
|
coll = self.collections[(dim, user, collection)]
|
||||||
|
|
||||||
|
|
@ -172,19 +182,27 @@ class DocVectors:
|
||||||
return res
|
return res
|
||||||
|
|
||||||
def delete_collection(self, user, collection):
|
def delete_collection(self, user, collection):
|
||||||
"""Delete a collection for the given user and collection"""
|
"""
|
||||||
collection_name = make_safe_collection_name(user, collection, self.prefix)
|
Delete all dimension variants of the collection for the given user/collection.
|
||||||
|
Since collections are created with dimension suffixes, we need to find and delete all.
|
||||||
|
"""
|
||||||
|
base_name = make_safe_collection_name(user, collection, self.prefix)
|
||||||
|
prefix = f"{base_name}_"
|
||||||
|
|
||||||
# Check if collection exists
|
# Get all collections and filter for matches
|
||||||
if self.client.has_collection(collection_name):
|
all_collections = self.client.list_collections()
|
||||||
# Drop the collection
|
matching_collections = [coll for coll in all_collections if coll.startswith(prefix)]
|
||||||
self.client.drop_collection(collection_name)
|
|
||||||
logger.info(f"Deleted Milvus collection: {collection_name}")
|
|
||||||
|
|
||||||
# Remove from our local cache
|
if not matching_collections:
|
||||||
keys_to_remove = [key for key in self.collections.keys() if key[1] == user and key[2] == collection]
|
logger.info(f"No collections found matching prefix {prefix}")
|
||||||
for key in keys_to_remove:
|
|
||||||
del self.collections[key]
|
|
||||||
else:
|
else:
|
||||||
logger.info(f"Collection {collection_name} does not exist, nothing to delete")
|
for collection_name in matching_collections:
|
||||||
|
self.client.drop_collection(collection_name)
|
||||||
|
logger.info(f"Deleted Milvus collection: {collection_name}")
|
||||||
|
logger.info(f"Deleted {len(matching_collections)} collection(s) for {user}/{collection}")
|
||||||
|
|
||||||
|
# Remove from our local cache
|
||||||
|
keys_to_remove = [key for key in self.collections.keys() if key[1] == user and key[2] == collection]
|
||||||
|
for key in keys_to_remove:
|
||||||
|
del self.collections[key]
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -50,24 +50,26 @@ class EntityVectors:
|
||||||
logger.debug(f"Reload at {self.next_reload}")
|
logger.debug(f"Reload at {self.next_reload}")
|
||||||
|
|
||||||
def collection_exists(self, user, collection):
|
def collection_exists(self, user, collection):
|
||||||
"""Check if collection exists (dimension-independent check)"""
|
"""
|
||||||
collection_name = make_safe_collection_name(user, collection, self.prefix)
|
Check if any collection exists for this user/collection combination.
|
||||||
return self.client.has_collection(collection_name)
|
Since collections are dimension-specific, this checks if ANY dimension variant exists.
|
||||||
|
"""
|
||||||
|
base_name = make_safe_collection_name(user, collection, self.prefix)
|
||||||
|
prefix = f"{base_name}_"
|
||||||
|
all_collections = self.client.list_collections()
|
||||||
|
return any(coll.startswith(prefix) for coll in all_collections)
|
||||||
|
|
||||||
def create_collection(self, user, collection, dimension=384):
|
def create_collection(self, user, collection, dimension=384):
|
||||||
"""Create collection with default dimension"""
|
"""
|
||||||
collection_name = make_safe_collection_name(user, collection, self.prefix)
|
No-op for explicit collection creation.
|
||||||
|
Collections are created lazily on first insert with actual dimension.
|
||||||
if self.client.has_collection(collection_name):
|
"""
|
||||||
logger.info(f"Collection {collection_name} already exists")
|
logger.info(f"Collection creation requested for {user}/{collection} - will be created lazily on first insert")
|
||||||
return
|
|
||||||
|
|
||||||
self.init_collection(dimension, user, collection)
|
|
||||||
logger.info(f"Created Milvus collection {collection_name} with dimension {dimension}")
|
|
||||||
|
|
||||||
def init_collection(self, dimension, user, collection):
|
def init_collection(self, dimension, user, collection):
|
||||||
|
|
||||||
collection_name = make_safe_collection_name(user, collection, self.prefix)
|
base_name = make_safe_collection_name(user, collection, self.prefix)
|
||||||
|
collection_name = f"{base_name}_{dimension}"
|
||||||
|
|
||||||
pkey_field = FieldSchema(
|
pkey_field = FieldSchema(
|
||||||
name="id",
|
name="id",
|
||||||
|
|
@ -115,6 +117,7 @@ class EntityVectors:
|
||||||
)
|
)
|
||||||
|
|
||||||
self.collections[(dimension, user, collection)] = collection_name
|
self.collections[(dimension, user, collection)] = collection_name
|
||||||
|
logger.info(f"Created Milvus collection {collection_name} with dimension {dimension}")
|
||||||
|
|
||||||
def insert(self, embeds, entity, user, collection):
|
def insert(self, embeds, entity, user, collection):
|
||||||
|
|
||||||
|
|
@ -139,8 +142,15 @@ class EntityVectors:
|
||||||
|
|
||||||
dim = len(embeds)
|
dim = len(embeds)
|
||||||
|
|
||||||
|
# Check if collection exists - return empty if not
|
||||||
if (dim, user, collection) not in self.collections:
|
if (dim, user, collection) not in self.collections:
|
||||||
self.init_collection(dim, user, collection)
|
base_name = make_safe_collection_name(user, collection, self.prefix)
|
||||||
|
collection_name = f"{base_name}_{dim}"
|
||||||
|
if not self.client.has_collection(collection_name):
|
||||||
|
logger.info(f"Collection {collection_name} does not exist, returning empty results")
|
||||||
|
return []
|
||||||
|
# Collection exists but not in cache, add it
|
||||||
|
self.collections[(dim, user, collection)] = collection_name
|
||||||
|
|
||||||
coll = self.collections[(dim, user, collection)]
|
coll = self.collections[(dim, user, collection)]
|
||||||
|
|
||||||
|
|
@ -172,19 +182,27 @@ class EntityVectors:
|
||||||
return res
|
return res
|
||||||
|
|
||||||
def delete_collection(self, user, collection):
|
def delete_collection(self, user, collection):
|
||||||
"""Delete a collection for the given user and collection"""
|
"""
|
||||||
collection_name = make_safe_collection_name(user, collection, self.prefix)
|
Delete all dimension variants of the collection for the given user/collection.
|
||||||
|
Since collections are created with dimension suffixes, we need to find and delete all.
|
||||||
|
"""
|
||||||
|
base_name = make_safe_collection_name(user, collection, self.prefix)
|
||||||
|
prefix = f"{base_name}_"
|
||||||
|
|
||||||
# Check if collection exists
|
# Get all collections and filter for matches
|
||||||
if self.client.has_collection(collection_name):
|
all_collections = self.client.list_collections()
|
||||||
# Drop the collection
|
matching_collections = [coll for coll in all_collections if coll.startswith(prefix)]
|
||||||
self.client.drop_collection(collection_name)
|
|
||||||
logger.info(f"Deleted Milvus collection: {collection_name}")
|
|
||||||
|
|
||||||
# Remove from our local cache
|
if not matching_collections:
|
||||||
keys_to_remove = [key for key in self.collections.keys() if key[1] == user and key[2] == collection]
|
logger.info(f"No collections found matching prefix {prefix}")
|
||||||
for key in keys_to_remove:
|
|
||||||
del self.collections[key]
|
|
||||||
else:
|
else:
|
||||||
logger.info(f"Collection {collection_name} does not exist, nothing to delete")
|
for collection_name in matching_collections:
|
||||||
|
self.client.drop_collection(collection_name)
|
||||||
|
logger.info(f"Deleted Milvus collection: {collection_name}")
|
||||||
|
logger.info(f"Deleted {len(matching_collections)} collection(s) for {user}/{collection}")
|
||||||
|
|
||||||
|
# Remove from our local cache
|
||||||
|
keys_to_remove = [key for key in self.collections.keys() if key[1] == user and key[2] == collection]
|
||||||
|
for key in keys_to_remove:
|
||||||
|
del self.collections[key]
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -47,39 +47,6 @@ class Processor(DocumentEmbeddingsQueryService):
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
self.last_index_name = None
|
|
||||||
|
|
||||||
def ensure_index_exists(self, index_name, dim):
|
|
||||||
"""Ensure index exists, create if it doesn't"""
|
|
||||||
if index_name != self.last_index_name:
|
|
||||||
if not self.pinecone.has_index(index_name):
|
|
||||||
try:
|
|
||||||
self.pinecone.create_index(
|
|
||||||
name=index_name,
|
|
||||||
dimension=dim,
|
|
||||||
metric="cosine",
|
|
||||||
spec=ServerlessSpec(
|
|
||||||
cloud="aws",
|
|
||||||
region="us-east-1",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
logger.info(f"Created index: {index_name}")
|
|
||||||
|
|
||||||
# Wait for index to be ready
|
|
||||||
import time
|
|
||||||
for i in range(0, 1000):
|
|
||||||
if self.pinecone.describe_index(index_name).status["ready"]:
|
|
||||||
break
|
|
||||||
time.sleep(1)
|
|
||||||
|
|
||||||
if not self.pinecone.describe_index(index_name).status["ready"]:
|
|
||||||
raise RuntimeError("Gave up waiting for index creation")
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Pinecone index creation failed: {e}")
|
|
||||||
raise e
|
|
||||||
self.last_index_name = index_name
|
|
||||||
|
|
||||||
async def query_document_embeddings(self, msg):
|
async def query_document_embeddings(self, msg):
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
|
@ -94,11 +61,13 @@ class Processor(DocumentEmbeddingsQueryService):
|
||||||
|
|
||||||
dim = len(vec)
|
dim = len(vec)
|
||||||
|
|
||||||
index_name = (
|
# Use dimension suffix in index name
|
||||||
"d-" + msg.user + "-" + msg.collection
|
index_name = f"d-{msg.user}-{msg.collection}-{dim}"
|
||||||
)
|
|
||||||
|
|
||||||
self.ensure_index_exists(index_name, dim)
|
# Check if index exists - skip if not
|
||||||
|
if not self.pinecone.has_index(index_name):
|
||||||
|
logger.info(f"Index {index_name} does not exist, skipping this vector")
|
||||||
|
continue
|
||||||
|
|
||||||
index = self.pinecone.Index(index_name)
|
index = self.pinecone.Index(index_name)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -38,28 +38,6 @@ class Processor(DocumentEmbeddingsQueryService):
|
||||||
)
|
)
|
||||||
|
|
||||||
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
|
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
|
||||||
self.last_collection = None
|
|
||||||
|
|
||||||
def ensure_collection_exists(self, collection, dim):
|
|
||||||
"""Ensure collection exists, create if it doesn't"""
|
|
||||||
if collection != self.last_collection:
|
|
||||||
if not self.qdrant.collection_exists(collection):
|
|
||||||
try:
|
|
||||||
self.qdrant.create_collection(
|
|
||||||
collection_name=collection,
|
|
||||||
vectors_config=VectorParams(
|
|
||||||
size=dim, distance=Distance.COSINE
|
|
||||||
),
|
|
||||||
)
|
|
||||||
logger.info(f"Created collection: {collection}")
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Qdrant collection creation failed: {e}")
|
|
||||||
raise e
|
|
||||||
self.last_collection = collection
|
|
||||||
|
|
||||||
def collection_exists(self, collection):
|
|
||||||
"""Check if collection exists (no implicit creation)"""
|
|
||||||
return self.qdrant.collection_exists(collection)
|
|
||||||
|
|
||||||
def collection_exists(self, collection):
|
def collection_exists(self, collection):
|
||||||
"""Check if collection exists (no implicit creation)"""
|
"""Check if collection exists (no implicit creation)"""
|
||||||
|
|
@ -71,16 +49,17 @@ class Processor(DocumentEmbeddingsQueryService):
|
||||||
|
|
||||||
chunks = []
|
chunks = []
|
||||||
|
|
||||||
collection = (
|
|
||||||
"d_" + msg.user + "_" + msg.collection
|
|
||||||
)
|
|
||||||
|
|
||||||
# Check if collection exists - return empty if not
|
|
||||||
if not self.collection_exists(collection):
|
|
||||||
logger.info(f"Collection {collection} does not exist, returning empty results")
|
|
||||||
return []
|
|
||||||
|
|
||||||
for vec in msg.vectors:
|
for vec in msg.vectors:
|
||||||
|
|
||||||
|
# Use dimension suffix in collection name
|
||||||
|
dim = len(vec)
|
||||||
|
collection = f"d_{msg.user}_{msg.collection}_{dim}"
|
||||||
|
|
||||||
|
# Check if collection exists - return empty if not
|
||||||
|
if not self.collection_exists(collection):
|
||||||
|
logger.info(f"Collection {collection} does not exist, returning empty results")
|
||||||
|
continue
|
||||||
|
|
||||||
search_result = self.qdrant.query_points(
|
search_result = self.qdrant.query_points(
|
||||||
collection_name=collection,
|
collection_name=collection,
|
||||||
query=vec,
|
query=vec,
|
||||||
|
|
|
||||||
|
|
@ -49,39 +49,6 @@ class Processor(GraphEmbeddingsQueryService):
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
self.last_index_name = None
|
|
||||||
|
|
||||||
def ensure_index_exists(self, index_name, dim):
|
|
||||||
"""Ensure index exists, create if it doesn't"""
|
|
||||||
if index_name != self.last_index_name:
|
|
||||||
if not self.pinecone.has_index(index_name):
|
|
||||||
try:
|
|
||||||
self.pinecone.create_index(
|
|
||||||
name=index_name,
|
|
||||||
dimension=dim,
|
|
||||||
metric="cosine",
|
|
||||||
spec=ServerlessSpec(
|
|
||||||
cloud="aws",
|
|
||||||
region="us-east-1",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
logger.info(f"Created index: {index_name}")
|
|
||||||
|
|
||||||
# Wait for index to be ready
|
|
||||||
import time
|
|
||||||
for i in range(0, 1000):
|
|
||||||
if self.pinecone.describe_index(index_name).status["ready"]:
|
|
||||||
break
|
|
||||||
time.sleep(1)
|
|
||||||
|
|
||||||
if not self.pinecone.describe_index(index_name).status["ready"]:
|
|
||||||
raise RuntimeError("Gave up waiting for index creation")
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Pinecone index creation failed: {e}")
|
|
||||||
raise e
|
|
||||||
self.last_index_name = index_name
|
|
||||||
|
|
||||||
def create_value(self, ent):
|
def create_value(self, ent):
|
||||||
if ent.startswith("http://") or ent.startswith("https://"):
|
if ent.startswith("http://") or ent.startswith("https://"):
|
||||||
return Value(value=ent, is_uri=True)
|
return Value(value=ent, is_uri=True)
|
||||||
|
|
@ -103,11 +70,13 @@ class Processor(GraphEmbeddingsQueryService):
|
||||||
|
|
||||||
dim = len(vec)
|
dim = len(vec)
|
||||||
|
|
||||||
index_name = (
|
# Use dimension suffix in index name
|
||||||
"t-" + msg.user + "-" + msg.collection
|
index_name = f"t-{msg.user}-{msg.collection}-{dim}"
|
||||||
)
|
|
||||||
|
|
||||||
self.ensure_index_exists(index_name, dim)
|
# Check if index exists - skip if not
|
||||||
|
if not self.pinecone.has_index(index_name):
|
||||||
|
logger.info(f"Index {index_name} does not exist, skipping this vector")
|
||||||
|
continue
|
||||||
|
|
||||||
index = self.pinecone.Index(index_name)
|
index = self.pinecone.Index(index_name)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -38,28 +38,6 @@ class Processor(GraphEmbeddingsQueryService):
|
||||||
)
|
)
|
||||||
|
|
||||||
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
|
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
|
||||||
self.last_collection = None
|
|
||||||
|
|
||||||
def ensure_collection_exists(self, collection, dim):
|
|
||||||
"""Ensure collection exists, create if it doesn't"""
|
|
||||||
if collection != self.last_collection:
|
|
||||||
if not self.qdrant.collection_exists(collection):
|
|
||||||
try:
|
|
||||||
self.qdrant.create_collection(
|
|
||||||
collection_name=collection,
|
|
||||||
vectors_config=VectorParams(
|
|
||||||
size=dim, distance=Distance.COSINE
|
|
||||||
),
|
|
||||||
)
|
|
||||||
logger.info(f"Created collection: {collection}")
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Qdrant collection creation failed: {e}")
|
|
||||||
raise e
|
|
||||||
self.last_collection = collection
|
|
||||||
|
|
||||||
def collection_exists(self, collection):
|
|
||||||
"""Check if collection exists (no implicit creation)"""
|
|
||||||
return self.qdrant.collection_exists(collection)
|
|
||||||
|
|
||||||
def collection_exists(self, collection):
|
def collection_exists(self, collection):
|
||||||
"""Check if collection exists (no implicit creation)"""
|
"""Check if collection exists (no implicit creation)"""
|
||||||
|
|
@ -78,17 +56,17 @@ class Processor(GraphEmbeddingsQueryService):
|
||||||
entity_set = set()
|
entity_set = set()
|
||||||
entities = []
|
entities = []
|
||||||
|
|
||||||
collection = (
|
|
||||||
"t_" + msg.user + "_" + msg.collection
|
|
||||||
)
|
|
||||||
|
|
||||||
# Check if collection exists - return empty if not
|
|
||||||
if not self.collection_exists(collection):
|
|
||||||
logger.info(f"Collection {collection} does not exist, returning empty results")
|
|
||||||
return []
|
|
||||||
|
|
||||||
for vec in msg.vectors:
|
for vec in msg.vectors:
|
||||||
|
|
||||||
|
# Use dimension suffix in collection name
|
||||||
|
dim = len(vec)
|
||||||
|
collection = f"t_{msg.user}_{msg.collection}_{dim}"
|
||||||
|
|
||||||
|
# Check if collection exists - return empty if not
|
||||||
|
if not self.collection_exists(collection):
|
||||||
|
logger.info(f"Collection {collection} does not exist, skipping this vector")
|
||||||
|
continue
|
||||||
|
|
||||||
# Heuristic hack, get (2*limit), so that we have more chance
|
# Heuristic hack, get (2*limit), so that we have more chance
|
||||||
# of getting (limit) entities
|
# of getting (limit) entities
|
||||||
search_result = self.qdrant.query_points(
|
search_result = self.qdrant.query_points(
|
||||||
|
|
|
||||||
|
|
@ -132,20 +132,20 @@ class Processor(DocumentEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_create_collection(self, request):
|
async def handle_create_collection(self, request):
|
||||||
"""Create a Milvus collection for document embeddings"""
|
"""
|
||||||
|
No-op for collection creation - collections are created lazily on first write
|
||||||
|
with the correct dimension determined from the actual embeddings.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
if self.vecstore.collection_exists(request.user, request.collection):
|
logger.info(f"Collection create request for {request.user}/{request.collection} - will be created lazily on first write")
|
||||||
logger.info(f"Collection {request.user}/{request.collection} already exists")
|
self.vecstore.create_collection(request.user, request.collection)
|
||||||
else:
|
|
||||||
self.vecstore.create_collection(request.user, request.collection)
|
|
||||||
logger.info(f"Created collection {request.user}/{request.collection}")
|
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(error=None)
|
response = StorageManagementResponse(error=None)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to create collection: {e}", exc_info=True)
|
logger.error(f"Failed to handle create collection request: {e}", exc_info=True)
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=Error(
|
error=Error(
|
||||||
type="creation_error",
|
type="creation_error",
|
||||||
|
|
|
||||||
|
|
@ -123,19 +123,6 @@ class Processor(DocumentEmbeddingsStoreService):
|
||||||
|
|
||||||
async def store_document_embeddings(self, message):
|
async def store_document_embeddings(self, message):
|
||||||
|
|
||||||
index_name = (
|
|
||||||
"d-" + message.metadata.user + "-" + message.metadata.collection
|
|
||||||
)
|
|
||||||
|
|
||||||
# Validate collection exists before accepting writes
|
|
||||||
if not self.pinecone.has_index(index_name):
|
|
||||||
error_msg = (
|
|
||||||
f"Collection {message.metadata.collection} does not exist. "
|
|
||||||
f"Create it first with tg-set-collection."
|
|
||||||
)
|
|
||||||
logger.error(error_msg)
|
|
||||||
raise ValueError(error_msg)
|
|
||||||
|
|
||||||
for emb in message.chunks:
|
for emb in message.chunks:
|
||||||
|
|
||||||
if emb.chunk is None or emb.chunk == b"": continue
|
if emb.chunk is None or emb.chunk == b"": continue
|
||||||
|
|
@ -145,6 +132,17 @@ class Processor(DocumentEmbeddingsStoreService):
|
||||||
|
|
||||||
for vec in emb.vectors:
|
for vec in emb.vectors:
|
||||||
|
|
||||||
|
# Create index name with dimension suffix for lazy creation
|
||||||
|
dim = len(vec)
|
||||||
|
index_name = (
|
||||||
|
f"d-{message.metadata.user}-{message.metadata.collection}-{dim}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Lazily create index if it doesn't exist
|
||||||
|
if not self.pinecone.has_index(index_name):
|
||||||
|
logger.info(f"Lazily creating Pinecone index {index_name} with dimension {dim}")
|
||||||
|
self.create_index(index_name, dim)
|
||||||
|
|
||||||
index = self.pinecone.Index(index_name)
|
index = self.pinecone.Index(index_name)
|
||||||
|
|
||||||
# Generate unique ID for each vector
|
# Generate unique ID for each vector
|
||||||
|
|
@ -220,23 +218,19 @@ class Processor(DocumentEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_create_collection(self, request):
|
async def handle_create_collection(self, request):
|
||||||
"""Create a Pinecone index for document embeddings"""
|
"""
|
||||||
|
No-op for collection creation - indexes are created lazily on first write
|
||||||
|
with the correct dimension determined from the actual embeddings.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
index_name = f"d-{request.user}-{request.collection}"
|
logger.info(f"Collection create request for {request.user}/{request.collection} - will be created lazily on first write")
|
||||||
|
|
||||||
if self.pinecone.has_index(index_name):
|
|
||||||
logger.info(f"Pinecone index {index_name} already exists")
|
|
||||||
else:
|
|
||||||
# Create with default dimension - will need to be recreated if dimension doesn't match
|
|
||||||
self.create_index(index_name, dim=384)
|
|
||||||
logger.info(f"Created Pinecone index: {index_name}")
|
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(error=None)
|
response = StorageManagementResponse(error=None)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to create collection: {e}", exc_info=True)
|
logger.error(f"Failed to handle create collection request: {e}", exc_info=True)
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=Error(
|
error=Error(
|
||||||
type="creation_error",
|
type="creation_error",
|
||||||
|
|
@ -246,22 +240,34 @@ class Processor(DocumentEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_delete_collection(self, request):
|
async def handle_delete_collection(self, request):
|
||||||
"""Delete the collection for document embeddings"""
|
"""
|
||||||
|
Delete all dimension variants of the index for document embeddings.
|
||||||
|
Since indexes are created with dimension suffixes (e.g., d-user-coll-384),
|
||||||
|
we need to find and delete all matching indexes.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
index_name = f"d-{request.user}-{request.collection}"
|
prefix = f"d-{request.user}-{request.collection}-"
|
||||||
|
|
||||||
if self.pinecone.has_index(index_name):
|
# Get all indexes and filter for matches
|
||||||
self.pinecone.delete_index(index_name)
|
all_indexes = self.pinecone.list_indexes()
|
||||||
logger.info(f"Deleted Pinecone index: {index_name}")
|
matching_indexes = [
|
||||||
|
idx.name for idx in all_indexes
|
||||||
|
if idx.name.startswith(prefix)
|
||||||
|
]
|
||||||
|
|
||||||
|
if not matching_indexes:
|
||||||
|
logger.info(f"No indexes found matching prefix {prefix}")
|
||||||
else:
|
else:
|
||||||
logger.info(f"Index {index_name} does not exist, nothing to delete")
|
for index_name in matching_indexes:
|
||||||
|
self.pinecone.delete_index(index_name)
|
||||||
|
logger.info(f"Deleted Pinecone index: {index_name}")
|
||||||
|
logger.info(f"Deleted {len(matching_indexes)} index(es) for {request.user}/{request.collection}")
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=None # No error means success
|
error=None # No error means success
|
||||||
)
|
)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
logger.info(f"Successfully deleted collection {request.user}/{request.collection}")
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to delete collection: {e}")
|
logger.error(f"Failed to delete collection: {e}")
|
||||||
|
|
|
||||||
|
|
@ -79,20 +79,6 @@ class Processor(DocumentEmbeddingsStoreService):
|
||||||
|
|
||||||
async def store_document_embeddings(self, message):
|
async def store_document_embeddings(self, message):
|
||||||
|
|
||||||
# Validate collection exists before accepting writes
|
|
||||||
collection = (
|
|
||||||
"d_" + message.metadata.user + "_" +
|
|
||||||
message.metadata.collection
|
|
||||||
)
|
|
||||||
|
|
||||||
if not self.qdrant.collection_exists(collection):
|
|
||||||
error_msg = (
|
|
||||||
f"Collection {message.metadata.collection} does not exist. "
|
|
||||||
f"Create it first with tg-set-collection."
|
|
||||||
)
|
|
||||||
logger.error(error_msg)
|
|
||||||
raise ValueError(error_msg)
|
|
||||||
|
|
||||||
for emb in message.chunks:
|
for emb in message.chunks:
|
||||||
|
|
||||||
chunk = emb.chunk.decode("utf-8")
|
chunk = emb.chunk.decode("utf-8")
|
||||||
|
|
@ -100,6 +86,23 @@ class Processor(DocumentEmbeddingsStoreService):
|
||||||
|
|
||||||
for vec in emb.vectors:
|
for vec in emb.vectors:
|
||||||
|
|
||||||
|
# Create collection name with dimension suffix for lazy creation
|
||||||
|
dim = len(vec)
|
||||||
|
collection = (
|
||||||
|
f"d_{message.metadata.user}_{message.metadata.collection}_{dim}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Lazily create collection if it doesn't exist
|
||||||
|
if not self.qdrant.collection_exists(collection):
|
||||||
|
logger.info(f"Lazily creating Qdrant collection {collection} with dimension {dim}")
|
||||||
|
self.qdrant.create_collection(
|
||||||
|
collection_name=collection,
|
||||||
|
vectors_config=VectorParams(
|
||||||
|
size=dim,
|
||||||
|
distance=Distance.COSINE
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
self.qdrant.upsert(
|
self.qdrant.upsert(
|
||||||
collection_name=collection,
|
collection_name=collection,
|
||||||
points=[
|
points=[
|
||||||
|
|
@ -160,30 +163,19 @@ class Processor(DocumentEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_create_collection(self, request):
|
async def handle_create_collection(self, request):
|
||||||
"""Create a Qdrant collection for document embeddings"""
|
"""
|
||||||
|
No-op for collection creation - collections are created lazily on first write
|
||||||
|
with the correct dimension determined from the actual embeddings.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
collection_name = f"d_{request.user}_{request.collection}"
|
logger.info(f"Collection create request for {request.user}/{request.collection} - will be created lazily on first write")
|
||||||
|
|
||||||
if self.qdrant.collection_exists(collection_name):
|
|
||||||
logger.info(f"Qdrant collection {collection_name} already exists")
|
|
||||||
else:
|
|
||||||
# Create collection with default dimension (will be recreated with correct dim on first write if needed)
|
|
||||||
# Using a placeholder dimension - actual dimension determined by first embedding
|
|
||||||
self.qdrant.create_collection(
|
|
||||||
collection_name=collection_name,
|
|
||||||
vectors_config=VectorParams(
|
|
||||||
size=384, # Default dimension, common for many models
|
|
||||||
distance=Distance.COSINE
|
|
||||||
)
|
|
||||||
)
|
|
||||||
logger.info(f"Created Qdrant collection: {collection_name}")
|
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(error=None)
|
response = StorageManagementResponse(error=None)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to create collection: {e}", exc_info=True)
|
logger.error(f"Failed to handle create collection request: {e}", exc_info=True)
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=Error(
|
error=Error(
|
||||||
type="creation_error",
|
type="creation_error",
|
||||||
|
|
@ -193,22 +185,34 @@ class Processor(DocumentEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_delete_collection(self, request):
|
async def handle_delete_collection(self, request):
|
||||||
"""Delete the collection for document embeddings"""
|
"""
|
||||||
|
Delete all dimension variants of the collection for document embeddings.
|
||||||
|
Since collections are created with dimension suffixes (e.g., d_user_coll_384),
|
||||||
|
we need to find and delete all matching collections.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
collection_name = f"d_{request.user}_{request.collection}"
|
prefix = f"d_{request.user}_{request.collection}_"
|
||||||
|
|
||||||
if self.qdrant.collection_exists(collection_name):
|
# Get all collections and filter for matches
|
||||||
self.qdrant.delete_collection(collection_name)
|
all_collections = self.qdrant.get_collections().collections
|
||||||
logger.info(f"Deleted Qdrant collection: {collection_name}")
|
matching_collections = [
|
||||||
|
coll.name for coll in all_collections
|
||||||
|
if coll.name.startswith(prefix)
|
||||||
|
]
|
||||||
|
|
||||||
|
if not matching_collections:
|
||||||
|
logger.info(f"No collections found matching prefix {prefix}")
|
||||||
else:
|
else:
|
||||||
logger.info(f"Collection {collection_name} does not exist, nothing to delete")
|
for collection_name in matching_collections:
|
||||||
|
self.qdrant.delete_collection(collection_name)
|
||||||
|
logger.info(f"Deleted Qdrant collection: {collection_name}")
|
||||||
|
logger.info(f"Deleted {len(matching_collections)} collection(s) for {request.user}/{request.collection}")
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=None # No error means success
|
error=None # No error means success
|
||||||
)
|
)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
logger.info(f"Successfully deleted collection {request.user}/{request.collection}")
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to delete collection: {e}")
|
logger.error(f"Failed to delete collection: {e}")
|
||||||
|
|
|
||||||
|
|
@ -128,20 +128,20 @@ class Processor(GraphEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_create_collection(self, request):
|
async def handle_create_collection(self, request):
|
||||||
"""Create a Milvus collection for graph embeddings"""
|
"""
|
||||||
|
No-op for collection creation - collections are created lazily on first write
|
||||||
|
with the correct dimension determined from the actual embeddings.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
if self.vecstore.collection_exists(request.user, request.collection):
|
logger.info(f"Collection create request for {request.user}/{request.collection} - will be created lazily on first write")
|
||||||
logger.info(f"Collection {request.user}/{request.collection} already exists")
|
self.vecstore.create_collection(request.user, request.collection)
|
||||||
else:
|
|
||||||
self.vecstore.create_collection(request.user, request.collection)
|
|
||||||
logger.info(f"Created collection {request.user}/{request.collection}")
|
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(error=None)
|
response = StorageManagementResponse(error=None)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to create collection: {e}", exc_info=True)
|
logger.error(f"Failed to handle create collection request: {e}", exc_info=True)
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=Error(
|
error=Error(
|
||||||
type="creation_error",
|
type="creation_error",
|
||||||
|
|
|
||||||
|
|
@ -123,19 +123,6 @@ class Processor(GraphEmbeddingsStoreService):
|
||||||
|
|
||||||
async def store_graph_embeddings(self, message):
|
async def store_graph_embeddings(self, message):
|
||||||
|
|
||||||
index_name = (
|
|
||||||
"t-" + message.metadata.user + "-" + message.metadata.collection
|
|
||||||
)
|
|
||||||
|
|
||||||
# Validate collection exists before accepting writes
|
|
||||||
if not self.pinecone.has_index(index_name):
|
|
||||||
error_msg = (
|
|
||||||
f"Collection {message.metadata.collection} does not exist. "
|
|
||||||
f"Create it first with tg-set-collection."
|
|
||||||
)
|
|
||||||
logger.error(error_msg)
|
|
||||||
raise ValueError(error_msg)
|
|
||||||
|
|
||||||
for entity in message.entities:
|
for entity in message.entities:
|
||||||
|
|
||||||
if entity.entity.value == "" or entity.entity.value is None:
|
if entity.entity.value == "" or entity.entity.value is None:
|
||||||
|
|
@ -143,6 +130,17 @@ class Processor(GraphEmbeddingsStoreService):
|
||||||
|
|
||||||
for vec in entity.vectors:
|
for vec in entity.vectors:
|
||||||
|
|
||||||
|
# Create index name with dimension suffix for lazy creation
|
||||||
|
dim = len(vec)
|
||||||
|
index_name = (
|
||||||
|
f"t-{message.metadata.user}-{message.metadata.collection}-{dim}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Lazily create index if it doesn't exist
|
||||||
|
if not self.pinecone.has_index(index_name):
|
||||||
|
logger.info(f"Lazily creating Pinecone index {index_name} with dimension {dim}")
|
||||||
|
self.create_index(index_name, dim)
|
||||||
|
|
||||||
index = self.pinecone.Index(index_name)
|
index = self.pinecone.Index(index_name)
|
||||||
|
|
||||||
# Generate unique ID for each vector
|
# Generate unique ID for each vector
|
||||||
|
|
@ -218,23 +216,19 @@ class Processor(GraphEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_create_collection(self, request):
|
async def handle_create_collection(self, request):
|
||||||
"""Create a Pinecone index for graph embeddings"""
|
"""
|
||||||
|
No-op for collection creation - indexes are created lazily on first write
|
||||||
|
with the correct dimension determined from the actual embeddings.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
index_name = f"t-{request.user}-{request.collection}"
|
logger.info(f"Collection create request for {request.user}/{request.collection} - will be created lazily on first write")
|
||||||
|
|
||||||
if self.pinecone.has_index(index_name):
|
|
||||||
logger.info(f"Pinecone index {index_name} already exists")
|
|
||||||
else:
|
|
||||||
# Create with default dimension - will need to be recreated if dimension doesn't match
|
|
||||||
self.create_index(index_name, dim=384)
|
|
||||||
logger.info(f"Created Pinecone index: {index_name}")
|
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(error=None)
|
response = StorageManagementResponse(error=None)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to create collection: {e}", exc_info=True)
|
logger.error(f"Failed to handle create collection request: {e}", exc_info=True)
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=Error(
|
error=Error(
|
||||||
type="creation_error",
|
type="creation_error",
|
||||||
|
|
@ -244,22 +238,34 @@ class Processor(GraphEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_delete_collection(self, request):
|
async def handle_delete_collection(self, request):
|
||||||
"""Delete the collection for graph embeddings"""
|
"""
|
||||||
|
Delete all dimension variants of the index for graph embeddings.
|
||||||
|
Since indexes are created with dimension suffixes (e.g., t-user-coll-384),
|
||||||
|
we need to find and delete all matching indexes.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
index_name = f"t-{request.user}-{request.collection}"
|
prefix = f"t-{request.user}-{request.collection}-"
|
||||||
|
|
||||||
if self.pinecone.has_index(index_name):
|
# Get all indexes and filter for matches
|
||||||
self.pinecone.delete_index(index_name)
|
all_indexes = self.pinecone.list_indexes()
|
||||||
logger.info(f"Deleted Pinecone index: {index_name}")
|
matching_indexes = [
|
||||||
|
idx.name for idx in all_indexes
|
||||||
|
if idx.name.startswith(prefix)
|
||||||
|
]
|
||||||
|
|
||||||
|
if not matching_indexes:
|
||||||
|
logger.info(f"No indexes found matching prefix {prefix}")
|
||||||
else:
|
else:
|
||||||
logger.info(f"Index {index_name} does not exist, nothing to delete")
|
for index_name in matching_indexes:
|
||||||
|
self.pinecone.delete_index(index_name)
|
||||||
|
logger.info(f"Deleted Pinecone index: {index_name}")
|
||||||
|
logger.info(f"Deleted {len(matching_indexes)} index(es) for {request.user}/{request.collection}")
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=None # No error means success
|
error=None # No error means success
|
||||||
)
|
)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
logger.info(f"Successfully deleted collection {request.user}/{request.collection}")
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to delete collection: {e}")
|
logger.error(f"Failed to delete collection: {e}")
|
||||||
|
|
|
||||||
|
|
@ -69,22 +69,6 @@ class Processor(GraphEmbeddingsStoreService):
|
||||||
metrics=storage_response_metrics,
|
metrics=storage_response_metrics,
|
||||||
)
|
)
|
||||||
|
|
||||||
def get_collection(self, user, collection):
|
|
||||||
"""Get collection name and validate it exists"""
|
|
||||||
cname = (
|
|
||||||
"t_" + user + "_" + collection
|
|
||||||
)
|
|
||||||
|
|
||||||
if not self.qdrant.collection_exists(cname):
|
|
||||||
error_msg = (
|
|
||||||
f"Collection {collection} does not exist. "
|
|
||||||
f"Create it first with tg-set-collection."
|
|
||||||
)
|
|
||||||
logger.error(error_msg)
|
|
||||||
raise ValueError(error_msg)
|
|
||||||
|
|
||||||
return cname
|
|
||||||
|
|
||||||
async def start(self):
|
async def start(self):
|
||||||
"""Start the processor and its storage management consumer"""
|
"""Start the processor and its storage management consumer"""
|
||||||
await super().start()
|
await super().start()
|
||||||
|
|
@ -101,10 +85,23 @@ class Processor(GraphEmbeddingsStoreService):
|
||||||
|
|
||||||
for vec in entity.vectors:
|
for vec in entity.vectors:
|
||||||
|
|
||||||
collection = self.get_collection(
|
# Create collection name with dimension suffix for lazy creation
|
||||||
message.metadata.user, message.metadata.collection
|
dim = len(vec)
|
||||||
|
collection = (
|
||||||
|
f"t_{message.metadata.user}_{message.metadata.collection}_{dim}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Lazily create collection if it doesn't exist
|
||||||
|
if not self.qdrant.collection_exists(collection):
|
||||||
|
logger.info(f"Lazily creating Qdrant collection {collection} with dimension {dim}")
|
||||||
|
self.qdrant.create_collection(
|
||||||
|
collection_name=collection,
|
||||||
|
vectors_config=VectorParams(
|
||||||
|
size=dim,
|
||||||
|
distance=Distance.COSINE
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
self.qdrant.upsert(
|
self.qdrant.upsert(
|
||||||
collection_name=collection,
|
collection_name=collection,
|
||||||
points=[
|
points=[
|
||||||
|
|
@ -165,30 +162,19 @@ class Processor(GraphEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_create_collection(self, request):
|
async def handle_create_collection(self, request):
|
||||||
"""Create a Qdrant collection for graph embeddings"""
|
"""
|
||||||
|
No-op for collection creation - collections are created lazily on first write
|
||||||
|
with the correct dimension determined from the actual embeddings.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
collection_name = f"t_{request.user}_{request.collection}"
|
logger.info(f"Collection create request for {request.user}/{request.collection} - will be created lazily on first write")
|
||||||
|
|
||||||
if self.qdrant.collection_exists(collection_name):
|
|
||||||
logger.info(f"Qdrant collection {collection_name} already exists")
|
|
||||||
else:
|
|
||||||
# Create collection with default dimension (will be recreated with correct dim on first write if needed)
|
|
||||||
# Using a placeholder dimension - actual dimension determined by first embedding
|
|
||||||
self.qdrant.create_collection(
|
|
||||||
collection_name=collection_name,
|
|
||||||
vectors_config=VectorParams(
|
|
||||||
size=384, # Default dimension, common for many models
|
|
||||||
distance=Distance.COSINE
|
|
||||||
)
|
|
||||||
)
|
|
||||||
logger.info(f"Created Qdrant collection: {collection_name}")
|
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(error=None)
|
response = StorageManagementResponse(error=None)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to create collection: {e}", exc_info=True)
|
logger.error(f"Failed to handle create collection request: {e}", exc_info=True)
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=Error(
|
error=Error(
|
||||||
type="creation_error",
|
type="creation_error",
|
||||||
|
|
@ -198,22 +184,34 @@ class Processor(GraphEmbeddingsStoreService):
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
|
|
||||||
async def handle_delete_collection(self, request):
|
async def handle_delete_collection(self, request):
|
||||||
"""Delete the collection for graph embeddings"""
|
"""
|
||||||
|
Delete all dimension variants of the collection for graph embeddings.
|
||||||
|
Since collections are created with dimension suffixes (e.g., t_user_coll_384),
|
||||||
|
we need to find and delete all matching collections.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
collection_name = f"t_{request.user}_{request.collection}"
|
prefix = f"t_{request.user}_{request.collection}_"
|
||||||
|
|
||||||
if self.qdrant.collection_exists(collection_name):
|
# Get all collections and filter for matches
|
||||||
self.qdrant.delete_collection(collection_name)
|
all_collections = self.qdrant.get_collections().collections
|
||||||
logger.info(f"Deleted Qdrant collection: {collection_name}")
|
matching_collections = [
|
||||||
|
coll.name for coll in all_collections
|
||||||
|
if coll.name.startswith(prefix)
|
||||||
|
]
|
||||||
|
|
||||||
|
if not matching_collections:
|
||||||
|
logger.info(f"No collections found matching prefix {prefix}")
|
||||||
else:
|
else:
|
||||||
logger.info(f"Collection {collection_name} does not exist, nothing to delete")
|
for collection_name in matching_collections:
|
||||||
|
self.qdrant.delete_collection(collection_name)
|
||||||
|
logger.info(f"Deleted Qdrant collection: {collection_name}")
|
||||||
|
logger.info(f"Deleted {len(matching_collections)} collection(s) for {request.user}/{request.collection}")
|
||||||
|
|
||||||
# Send success response
|
# Send success response
|
||||||
response = StorageManagementResponse(
|
response = StorageManagementResponse(
|
||||||
error=None # No error means success
|
error=None # No error means success
|
||||||
)
|
)
|
||||||
await self.storage_response_producer.send(response)
|
await self.storage_response_producer.send(response)
|
||||||
logger.info(f"Successfully deleted collection {request.user}/{request.collection}")
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to delete collection: {e}")
|
logger.error(f"Failed to delete collection: {e}")
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue