apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-23 20:21:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Remove implicit collection creation

Browse source
This commit is contained in:
Cyber MacGeddon 2025-09-30 12:36:23 +01:00
parent 9ec52d7c65
commit df80995d2e
6 changed files with 149 additions and 55 deletions
Download patch file Download diff file Expand all files Collapse all files

34
trustgraph-flow/trustgraph/direct/cassandra_kg.py
View file

@ -29,6 +29,7 @@ class KnowledgeGraph:
self.po_table = "triples_p"
self.object_table = "triples_o"
self.collection_table = "triples_collection" # For SPO queries and deletion
self.collection_metadata_table = "collection_metadata" # For tracking which collections exist
if username and password:
ssl_context = SSLContext(PROTOCOL_TLSv1_2)
@ -116,7 +117,17 @@ class KnowledgeGraph:
);
""");
logger.info("Optimized multi-table schema initialized (4 tables)")
# Table 5: Collection metadata tracking
# Tracks which collections exist without polluting triple data
self.session.execute(f"""
CREATE TABLE IF NOT EXISTS {self.collection_metadata_table} (
collection text,
created_at timestamp,
PRIMARY KEY (collection)
);
""");
logger.info("Optimized multi-table schema initialized (5 tables)")
def prepare_statements(self):
"""Prepare statements for optimal performance"""
@ -248,10 +259,10 @@ class KnowledgeGraph:
)
def collection_exists(self, collection):
"""Check if collection exists by querying triples_collection table"""
"""Check if collection exists by querying collection_metadata table"""
try:
result = self.session.execute(
f"SELECT collection FROM {self.collection_table} WHERE collection = %s LIMIT 1",
f"SELECT collection FROM {self.collection_metadata_table} WHERE collection = %s LIMIT 1",
(collection,)
)
return bool(list(result))
@ -260,15 +271,14 @@ class KnowledgeGraph:
return False
def create_collection(self, collection):
"""Create collection by inserting marker row in triples_collection table"""
"""Create collection by inserting metadata row"""
try:
# Insert a system marker triple to establish collection exists
# This won't interfere with data as it uses a reserved system namespace
import datetime
self.session.execute(
f"INSERT INTO {self.collection_table} (collection, s, p, o) VALUES (%s, %s, %s, %s)",
(collection, "__system__", "__collection_created__", "__marker__")
f"INSERT INTO {self.collection_metadata_table} (collection, created_at) VALUES (%s, %s)",
(collection, datetime.datetime.now())
)
logger.info(f"Created collection marker for {collection}")
logger.info(f"Created collection metadata for {collection}")
except Exception as e:
logger.error(f"Error creating collection: {e}")
raise e
@ -324,6 +334,12 @@ class KnowledgeGraph:
if count % 100 != 0:
self.session.execute(batch)
# Step 3: Delete collection metadata
self.session.execute(
f"DELETE FROM {self.collection_metadata_table} WHERE collection = %s",
(collection,)
)
logger.info(f"Deleted {count} triples from collection {collection}")
def close(self):

34
trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py
View file

@ -39,20 +39,9 @@ class Processor(DocumentEmbeddingsQueryService):
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
def ensure_collection_exists(self, collection, dim):
"""Ensure collection exists, create if it doesn't"""
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
def collection_exists(self, collection):
"""Check if collection exists (no implicit creation)"""
return self.qdrant.collection_exists(collection)
async def query_document_embeddings(self, msg):
@ -60,15 +49,16 @@ class Processor(DocumentEmbeddingsQueryService):
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:
dim = len(vec)
collection = (
"d_" + msg.user + "_" + msg.collection
)
self.ensure_collection_exists(collection, dim)
search_result = self.qdrant.query_points(
collection_name=collection,
query=vec,

33
trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py
View file

@ -39,20 +39,9 @@ class Processor(GraphEmbeddingsQueryService):
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
def ensure_collection_exists(self, collection, dim):
"""Ensure collection exists, create if it doesn't"""
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
def collection_exists(self, collection):
"""Check if collection exists (no implicit creation)"""
return self.qdrant.collection_exists(collection)
def create_value(self, ent):
if ent.startswith("http://") or ent.startswith("https://"):
@ -67,15 +56,17 @@ class Processor(GraphEmbeddingsQueryService):
entity_set = set()
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:
dim = len(vec)
collection = (
"t_" + msg.user + "_" + msg.collection
)
self.ensure_collection_exists(collection, dim)
# Heuristic hack, get (2*limit), so that we have more chance
# of getting (limit) entities
search_result = self.qdrant.query_points(

37
trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py
View file

@ -140,7 +140,9 @@ class Processor(DocumentEmbeddingsStoreService):
logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}")
try:
if request.operation == "delete-collection":
if request.operation == "create-collection":
await self.handle_create_collection(request)
elif request.operation == "delete-collection":
await self.handle_delete_collection(request)
else:
response = StorageManagementResponse(
@ -161,6 +163,39 @@ class Processor(DocumentEmbeddingsStoreService):
)
await self.storage_response_producer.send(response)
async def handle_create_collection(self, request):
"""Create a Qdrant collection for document embeddings"""
try:
collection_name = f"d_{request.user}_{request.collection}"
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
response = StorageManagementResponse(error=None)
await self.storage_response_producer.send(response)
except Exception as e:
logger.error(f"Failed to create collection: {e}", exc_info=True)
response = StorageManagementResponse(
error=Error(
type="creation_error",
message=str(e)
)
)
await self.storage_response_producer.send(response)
async def handle_delete_collection(self, request):
"""Delete the collection for document embeddings"""
try:

37
trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
View file

@ -147,7 +147,9 @@ class Processor(GraphEmbeddingsStoreService):
logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}")
try:
if request.operation == "delete-collection":
if request.operation == "create-collection":
await self.handle_create_collection(request)
elif request.operation == "delete-collection":
await self.handle_delete_collection(request)
else:
response = StorageManagementResponse(
@ -168,6 +170,39 @@ class Processor(GraphEmbeddingsStoreService):
)
await self.storage_response_producer.send(response)
async def handle_create_collection(self, request):
"""Create a Qdrant collection for graph embeddings"""
try:
collection_name = f"t_{request.user}_{request.collection}"
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
response = StorageManagementResponse(error=None)
await self.storage_response_producer.send(response)
except Exception as e:
logger.error(f"Failed to create collection: {e}", exc_info=True)
response = StorageManagementResponse(
error=Error(
type="creation_error",
message=str(e)
)
)
await self.storage_response_producer.send(response)
async def handle_delete_collection(self, request):
"""Delete the collection for graph embeddings"""
try:

29
trustgraph-flow/trustgraph/storage/objects/cassandra/write.py
View file

@ -434,7 +434,16 @@ class Processor(FlowProcessor):
logger.info(f"Received storage management request: {msg.operation} for {msg.user}/{msg.collection}")
try:
if msg.operation == "delete-collection":
if msg.operation == "create-collection":
await self.create_collection(msg.user, msg.collection)
# Send success response
response = StorageManagementResponse(
error=None # No error means success
)
await self.storage_response_producer.send(response)
logger.info(f"Successfully created collection {msg.user}/{msg.collection}")
elif msg.operation == "delete-collection":
await self.delete_collection(msg.user, msg.collection)
# Send success response
@ -467,6 +476,24 @@ class Processor(FlowProcessor):
)
await self.storage_response_producer.send(response)
async def create_collection(self, user: str, collection: str):
"""Create/verify collection exists in Cassandra object store"""
# Connect if not already connected
self.connect_cassandra()
# Sanitize names for safety
safe_keyspace = self.sanitize_name(user)
# Ensure keyspace exists
if safe_keyspace not in self.known_keyspaces:
self.ensure_keyspace(safe_keyspace)
self.known_keyspaces.add(safe_keyspace)
# For Cassandra objects, collection is just a property in rows
# No need to create separate tables per collection
# Just mark that we've seen this collection
logger.info(f"Collection {collection} ready for user {user} (using keyspace {safe_keyspace})")
async def delete_collection(self, user: str, collection: str):
"""Delete all data for a specific collection"""
# Connect if not already connected