import asyncio
import contextlib
import csv
import datetime
import json
import logging
import os
import tempfile
import time
from typing import IO, Any, AsyncGenerator, Optional, Tuple
from uuid import UUID
import asyncpg
import httpx
from asyncpg.exceptions import UniqueViolationError
from fastapi import HTTPException
from core.base.abstractions import (
Community,
Entity,
Graph,
GraphExtractionStatus,
R2RException,
Relationship,
StoreType,
VectorQuantizationType,
)
from core.base.api.models import GraphResponse
from core.base.providers.database import Handler
from core.base.utils import (
_get_vector_column_str,
generate_entity_document_id,
)
from .base import PostgresConnectionManager
from .collections import PostgresCollectionsHandler
logger = logging.getLogger()
class PostgresEntitiesHandler(Handler):
def __init__(self, *args: Any, **kwargs: Any) -> None:
self.project_name: str = kwargs.get("project_name") # type: ignore
self.connection_manager: PostgresConnectionManager = kwargs.get(
"connection_manager"
) # type: ignore
self.dimension: int = kwargs.get("dimension") # type: ignore
self.quantization_type: VectorQuantizationType = kwargs.get(
"quantization_type"
) # type: ignore
self.relationships_handler: PostgresRelationshipsHandler = (
PostgresRelationshipsHandler(*args, **kwargs)
)
def _get_table_name(self, table: str) -> str:
"""Get the fully qualified table name."""
return f'"{self.project_name}"."{table}"'
def _get_entity_table_for_store(self, store_type: StoreType) -> str:
"""Get the appropriate table name for the store type."""
return f"{store_type.value}_entities"
def _get_parent_constraint(self, store_type: StoreType) -> str:
"""Get the appropriate foreign key constraint for the store type."""
if store_type == StoreType.GRAPHS:
return f"""
CONSTRAINT fk_graph
FOREIGN KEY(parent_id)
REFERENCES {self._get_table_name("graphs")}(id)
ON DELETE CASCADE
"""
else:
return f"""
CONSTRAINT fk_document
FOREIGN KEY(parent_id)
REFERENCES {self._get_table_name("documents")}(id)
ON DELETE CASCADE
"""
async def create_tables(self) -> None:
"""Create separate tables for graph and document entities."""
vector_column_str = _get_vector_column_str(
self.dimension, self.quantization_type
)
for store_type in StoreType:
table_name = self._get_entity_table_for_store(store_type)
parent_constraint = self._get_parent_constraint(store_type)
QUERY = f"""
CREATE TABLE IF NOT EXISTS {self._get_table_name(table_name)} (
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
name TEXT NOT NULL,
category TEXT,
description TEXT,
parent_id UUID NOT NULL,
description_embedding {vector_column_str},
chunk_ids UUID[],
metadata JSONB,
created_at TIMESTAMPTZ DEFAULT NOW(),
updated_at TIMESTAMPTZ DEFAULT NOW(),
{parent_constraint}
);
CREATE INDEX IF NOT EXISTS {table_name}_name_idx
ON {self._get_table_name(table_name)} (name);
CREATE INDEX IF NOT EXISTS {table_name}_parent_id_idx
ON {self._get_table_name(table_name)} (parent_id);
CREATE INDEX IF NOT EXISTS {table_name}_category_idx
ON {self._get_table_name(table_name)} (category);
"""
await self.connection_manager.execute_query(QUERY)
async def create(
self,
parent_id: UUID,
store_type: StoreType,
name: str,
category: Optional[str] = None,
description: Optional[str] = None,
description_embedding: Optional[list[float] | str] = None,
chunk_ids: Optional[list[UUID]] = None,
metadata: Optional[dict[str, Any] | str] = None,
) -> Entity:
"""Create a new entity in the specified store."""
table_name = self._get_entity_table_for_store(store_type)
if isinstance(metadata, str):
with contextlib.suppress(json.JSONDecodeError):
metadata = json.loads(metadata)
if isinstance(description_embedding, list):
description_embedding = str(description_embedding)
query = f"""
INSERT INTO {self._get_table_name(table_name)}
(name, category, description, parent_id, description_embedding, chunk_ids, metadata)
VALUES ($1, $2, $3, $4, $5, $6, $7)
RETURNING id, name, category, description, parent_id, chunk_ids, metadata
"""
params = [
name,
category,
description,
parent_id,
description_embedding,
chunk_ids,
json.dumps(metadata) if metadata else None,
]
result = await self.connection_manager.fetchrow_query(
query=query,
params=params,
)
return Entity(
id=result["id"],
name=result["name"],
category=result["category"],
description=result["description"],
parent_id=result["parent_id"],
chunk_ids=result["chunk_ids"],
metadata=result["metadata"],
)
async def get(
self,
parent_id: UUID,
store_type: StoreType,
offset: int,
limit: int,
entity_ids: Optional[list[UUID]] = None,
entity_names: Optional[list[str]] = None,
include_embeddings: bool = False,
):
"""Retrieve entities from the specified store."""
table_name = self._get_entity_table_for_store(store_type)
conditions = ["parent_id = $1"]
params: list[Any] = [parent_id]
param_index = 2
if entity_ids:
conditions.append(f"id = ANY(${param_index})")
params.append(entity_ids)
param_index += 1
if entity_names:
conditions.append(f"name = ANY(${param_index})")
params.append(entity_names)
param_index += 1
select_fields = """
id, name, category, description, parent_id,
chunk_ids, metadata
"""
if include_embeddings:
select_fields += ", description_embedding"
COUNT_QUERY = f"""
SELECT COUNT(*)
FROM {self._get_table_name(table_name)}
WHERE {" AND ".join(conditions)}
"""
count_params = params[: param_index - 1]
count = (
await self.connection_manager.fetch_query(
COUNT_QUERY, count_params
)
)[0]["count"]
QUERY = f"""
SELECT {select_fields}
FROM {self._get_table_name(table_name)}
WHERE {" AND ".join(conditions)}
ORDER BY created_at
OFFSET ${param_index}
"""
params.append(offset)
param_index += 1
if limit != -1:
QUERY += f" LIMIT ${param_index}"
params.append(limit)
rows = await self.connection_manager.fetch_query(QUERY, params)
entities = []
for row in rows:
# Convert the Record to a dictionary
entity_dict = dict(row)
# Process metadata if it exists and is a string
if isinstance(entity_dict["metadata"], str):
with contextlib.suppress(json.JSONDecodeError):
entity_dict["metadata"] = json.loads(
entity_dict["metadata"]
)
entities.append(Entity(**entity_dict))
return entities, count
async def update(
self,
entity_id: UUID,
store_type: StoreType,
name: Optional[str] = None,
description: Optional[str] = None,
description_embedding: Optional[list[float] | str] = None,
category: Optional[str] = None,
metadata: Optional[dict] = None,
) -> Entity:
"""Update an entity in the specified store."""
table_name = self._get_entity_table_for_store(store_type)
update_fields = []
params: list[Any] = []
param_index = 1
if isinstance(metadata, str):
with contextlib.suppress(json.JSONDecodeError):
metadata = json.loads(metadata)
if name is not None:
update_fields.append(f"name = ${param_index}")
params.append(name)
param_index += 1
if description is not None:
update_fields.append(f"description = ${param_index}")
params.append(description)
param_index += 1
if description_embedding is not None:
update_fields.append(f"description_embedding = ${param_index}")
params.append(description_embedding)
param_index += 1
if category is not None:
update_fields.append(f"category = ${param_index}")
params.append(category)
param_index += 1
if metadata is not None:
update_fields.append(f"metadata = ${param_index}")
params.append(json.dumps(metadata))
param_index += 1
if not update_fields:
raise R2RException(status_code=400, message="No fields to update")
update_fields.append("updated_at = NOW()")
params.append(entity_id)
query = f"""
UPDATE {self._get_table_name(table_name)}
SET {", ".join(update_fields)}
WHERE id = ${param_index}\
RETURNING id, name, category, description, parent_id, chunk_ids, metadata
"""
try:
result = await self.connection_manager.fetchrow_query(
query=query,
params=params,
)
return Entity(
id=result["id"],
name=result["name"],
category=result["category"],
description=result["description"],
parent_id=result["parent_id"],
chunk_ids=result["chunk_ids"],
metadata=result["metadata"],
)
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"An error occurred while updating the entity: {e}",
) from e
async def delete(
self,
parent_id: UUID,
entity_ids: Optional[list[UUID]] = None,
store_type: StoreType = StoreType.GRAPHS,
) -> None:
"""Delete entities from the specified store. If entity_ids is not
provided, deletes all entities for the given parent_id.
Args:
parent_id (UUID): Parent ID (collection_id or document_id)
entity_ids (Optional[list[UUID]]): Specific entity IDs to delete. If None, deletes all entities for parent_id
store_type (StoreType): Type of store (graph or document)
Returns:
list[UUID]: List of deleted entity IDs
Raises:
R2RException: If specific entities were requested but not all found
"""
table_name = self._get_entity_table_for_store(store_type)
if entity_ids is None:
# Delete all entities for the parent_id
QUERY = f"""
DELETE FROM {self._get_table_name(table_name)}
WHERE parent_id = $1
RETURNING id
"""
results = await self.connection_manager.fetch_query(
QUERY, [parent_id]
)
else:
# Delete specific entities
QUERY = f"""
DELETE FROM {self._get_table_name(table_name)}
WHERE id = ANY($1) AND parent_id = $2
RETURNING id
"""
results = await self.connection_manager.fetch_query(
QUERY, [entity_ids, parent_id]
)
# Check if all requested entities were deleted
deleted_ids = [row["id"] for row in results]
if entity_ids and len(deleted_ids) != len(entity_ids):
raise R2RException(
f"Some entities not found in {store_type} store or no permission to delete",
404,
)
async def get_duplicate_name_blocks(
self,
parent_id: UUID,
store_type: StoreType,
) -> list[list[Entity]]:
"""Find all groups of entities that share identical names within the
same parent.
Returns a list of entity groups, where each group contains entities
with the same name. For each group, includes the n most dissimilar
descriptions based on cosine similarity.
"""
table_name = self._get_entity_table_for_store(store_type)
# First get the duplicate names and their descriptions with embeddings
query = f"""
WITH duplicates AS (
SELECT name
FROM {self._get_table_name(table_name)}
WHERE parent_id = $1
GROUP BY name
HAVING COUNT(*) > 1
)
SELECT
e.id, e.name, e.category, e.description,
e.parent_id, e.chunk_ids, e.metadata
FROM {self._get_table_name(table_name)} e
WHERE e.parent_id = $1
AND e.name IN (SELECT name FROM duplicates)
ORDER BY e.name;
"""
rows = await self.connection_manager.fetch_query(query, [parent_id])
# Group entities by name
name_groups: dict[str, list[Entity]] = {}
for row in rows:
entity_dict = dict(row)
if isinstance(entity_dict["metadata"], str):
with contextlib.suppress(json.JSONDecodeError):
entity_dict["metadata"] = json.loads(
entity_dict["metadata"]
)
entity = Entity(**entity_dict)
name_groups.setdefault(entity.name, []).append(entity)
return list(name_groups.values())
async def merge_duplicate_name_blocks(
self,
parent_id: UUID,
store_type: StoreType,
) -> list[tuple[list[Entity], Entity]]:
"""Merge entities that share identical names.
Returns list of tuples: (original_entities, merged_entity)
"""
duplicate_blocks = await self.get_duplicate_name_blocks(
parent_id, store_type
)
merged_results: list[tuple[list[Entity], Entity]] = []
for block in duplicate_blocks:
# Create a new merged entity from the block
merged_entity = await self._create_merged_entity(block)
merged_results.append((block, merged_entity))
table_name = self._get_entity_table_for_store(store_type)
async with self.connection_manager.transaction():
# Insert the merged entity
new_id = await self._insert_merged_entity(
merged_entity, table_name
)
merged_entity.id = new_id
# Get the old entity IDs
old_ids = [str(entity.id) for entity in block]
relationship_table = self.relationships_handler._get_relationship_table_for_store(
store_type
)
# Update relationships where old entities appear as subjects
subject_update_query = f"""
UPDATE {self._get_table_name(relationship_table)}
SET subject_id = $1
WHERE subject_id = ANY($2::uuid[])
AND parent_id = $3
"""
await self.connection_manager.execute_query(
subject_update_query, [new_id, old_ids, parent_id]
)
# Update relationships where old entities appear as objects
object_update_query = f"""
UPDATE {self._get_table_name(relationship_table)}
SET object_id = $1
WHERE object_id = ANY($2::uuid[])
AND parent_id = $3
"""
await self.connection_manager.execute_query(
object_update_query, [new_id, old_ids, parent_id]
)
# Delete the original entities
delete_query = f"""
DELETE FROM {self._get_table_name(table_name)}
WHERE id = ANY($1::uuid[])
"""
await self.connection_manager.execute_query(
delete_query, [old_ids]
)
return merged_results
async def _insert_merged_entity(
self, entity: Entity, table_name: str
) -> UUID:
"""Insert merged entity and return its new ID."""
new_id = generate_entity_document_id()
query = f"""
INSERT INTO {self._get_table_name(table_name)}
(id, name, category, description, parent_id, chunk_ids, metadata)
VALUES ($1, $2, $3, $4, $5, $6, $7)
RETURNING id
"""
values = [
new_id,
entity.name,
entity.category,
entity.description,
entity.parent_id,
entity.chunk_ids,
json.dumps(entity.metadata) if entity.metadata else None,
]
result = await self.connection_manager.fetch_query(query, values)
return result[0]["id"]
async def _create_merged_entity(self, entities: list[Entity]) -> Entity:
"""Create a merged entity from a list of duplicate entities.
Uses various strategies to combine fields.
"""
if not entities:
raise ValueError("Cannot merge empty list of entities")
# Take the first non-None category, or None if all are None
category = next(
(e.category for e in entities if e.category is not None), None
)
# Combine descriptions with newlines if they differ
descriptions = {e.description for e in entities if e.description}
description = "\n\n".join(descriptions) if descriptions else None
# Combine chunk_ids, removing duplicates
chunk_ids = list(
{
chunk_id
for entity in entities
for chunk_id in (entity.chunk_ids or [])
}
)
# Merge metadata dictionaries
merged_metadata: dict[str, Any] = {}
for entity in entities:
if entity.metadata:
merged_metadata |= entity.metadata
# Create new merged entity (without actually inserting to DB)
return Entity(
id=UUID(
"00000000-0000-0000-0000-000000000000"
), # Placeholder UUID
name=entities[0].name, # All entities in block have same name
category=category,
description=description,
parent_id=entities[0].parent_id,
chunk_ids=chunk_ids or None,
metadata=merged_metadata or None,
)
async def export_to_csv(
self,
parent_id: UUID,
store_type: StoreType,
columns: Optional[list[str]] = None,
filters: Optional[dict] = None,
include_header: bool = True,
) -> tuple[str, IO]:
"""Creates a CSV file from the PostgreSQL data and returns the path to
the temp file."""
valid_columns = {
"id",
"name",
"category",
"description",
"parent_id",
"chunk_ids",
"metadata",
"created_at",
"updated_at",
}
if not columns:
columns = list(valid_columns)
elif invalid_cols := set(columns) - valid_columns:
raise ValueError(f"Invalid columns: {invalid_cols}")
select_stmt = f"""
SELECT
id::text,
name,
category,
description,
parent_id::text,
chunk_ids::text,
metadata::text,
to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') AS created_at,
to_char(updated_at, 'YYYY-MM-DD HH24:MI:SS') AS updated_at
FROM {self._get_table_name(self._get_entity_table_for_store(store_type))}
"""
conditions = ["parent_id = $1"]
params: list[Any] = [parent_id]
param_index = 2
if filters:
for field, value in filters.items():
if field not in valid_columns:
continue
if isinstance(value, dict):
for op, val in value.items():
if op == "$eq":
conditions.append(f"{field} = ${param_index}")
params.append(val)
param_index += 1
elif op == "$gt":
conditions.append(f"{field} > ${param_index}")
params.append(val)
param_index += 1
elif op == "$lt":
conditions.append(f"{field} < ${param_index}")
params.append(val)
param_index += 1
else:
# Direct equality
conditions.append(f"{field} = ${param_index}")
params.append(value)
param_index += 1
if conditions:
select_stmt = f"{select_stmt} WHERE {' AND '.join(conditions)}"
select_stmt = f"{select_stmt} ORDER BY created_at DESC"
temp_file = None
try:
temp_file = tempfile.NamedTemporaryFile(
mode="w", delete=True, suffix=".csv"
)
writer = csv.writer(temp_file, quoting=csv.QUOTE_ALL)
async with self.connection_manager.pool.get_connection() as conn: # type: ignore
async with conn.transaction():
cursor = await conn.cursor(select_stmt, *params)
if include_header:
writer.writerow(columns)
chunk_size = 1000
while True:
rows = await cursor.fetch(chunk_size)
if not rows:
break
for row in rows:
row_dict = {
"id": row[0],
"name": row[1],
"category": row[2],
"description": row[3],
"parent_id": row[4],
"chunk_ids": row[5],
"metadata": row[6],
"created_at": row[7],
"updated_at": row[8],
}
writer.writerow([row_dict[col] for col in columns])
temp_file.flush()
return temp_file.name, temp_file
except Exception as e:
if temp_file:
temp_file.close()
raise HTTPException(
status_code=500,
detail=f"Failed to export data: {str(e)}",
) from e
class PostgresRelationshipsHandler(Handler):
def __init__(self, *args: Any, **kwargs: Any) -> None:
self.project_name: str = kwargs.get("project_name") # type: ignore
self.connection_manager: PostgresConnectionManager = kwargs.get(
"connection_manager"
) # type: ignore
self.dimension: int = kwargs.get("dimension") # type: ignore
self.quantization_type: VectorQuantizationType = kwargs.get(
"quantization_type"
) # type: ignore
def _get_table_name(self, table: str) -> str:
"""Get the fully qualified table name."""
return f'"{self.project_name}"."{table}"'
def _get_relationship_table_for_store(self, store_type: StoreType) -> str:
"""Get the appropriate table name for the store type."""
return f"{store_type.value}_relationships"
def _get_parent_constraint(self, store_type: StoreType) -> str:
"""Get the appropriate foreign key constraint for the store type."""
if store_type == StoreType.GRAPHS:
return f"""
CONSTRAINT fk_graph
FOREIGN KEY(parent_id)
REFERENCES {self._get_table_name("graphs")}(id)
ON DELETE CASCADE
"""
else:
return f"""
CONSTRAINT fk_document
FOREIGN KEY(parent_id)
REFERENCES {self._get_table_name("documents")}(id)
ON DELETE CASCADE
"""
async def create_tables(self) -> None:
"""Create separate tables for graph and document relationships."""
for store_type in StoreType:
table_name = self._get_relationship_table_for_store(store_type)
parent_constraint = self._get_parent_constraint(store_type)
vector_column_str = _get_vector_column_str(
self.dimension, self.quantization_type
)
QUERY = f"""
CREATE TABLE IF NOT EXISTS {self._get_table_name(table_name)} (
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
subject TEXT NOT NULL,
predicate TEXT NOT NULL,
object TEXT NOT NULL,
description TEXT,
description_embedding {vector_column_str},
subject_id UUID,
object_id UUID,
weight FLOAT DEFAULT 1.0,
chunk_ids UUID[],
parent_id UUID NOT NULL,
metadata JSONB,
created_at TIMESTAMPTZ DEFAULT NOW(),
updated_at TIMESTAMPTZ DEFAULT NOW(),
{parent_constraint}
);
CREATE INDEX IF NOT EXISTS {table_name}_subject_idx
ON {self._get_table_name(table_name)} (subject);
CREATE INDEX IF NOT EXISTS {table_name}_object_idx
ON {self._get_table_name(table_name)} (object);
CREATE INDEX IF NOT EXISTS {table_name}_predicate_idx
ON {self._get_table_name(table_name)} (predicate);
CREATE INDEX IF NOT EXISTS {table_name}_parent_id_idx
ON {self._get_table_name(table_name)} (parent_id);
CREATE INDEX IF NOT EXISTS {table_name}_subject_id_idx
ON {self._get_table_name(table_name)} (subject_id);
CREATE INDEX IF NOT EXISTS {table_name}_object_id_idx
ON {self._get_table_name(table_name)} (object_id);
"""
await self.connection_manager.execute_query(QUERY)
async def create(
self,
subject: str,
subject_id: UUID,
predicate: str,
object: str,
object_id: UUID,
parent_id: UUID,
store_type: StoreType,
description: str | None = None,
weight: float | None = 1.0,
chunk_ids: Optional[list[UUID]] = None,
description_embedding: Optional[list[float] | str] = None,
metadata: Optional[dict[str, Any] | str] = None,
) -> Relationship:
"""Create a new relationship in the specified store."""
table_name = self._get_relationship_table_for_store(store_type)
if isinstance(metadata, str):
with contextlib.suppress(json.JSONDecodeError):
metadata = json.loads(metadata)
if isinstance(description_embedding, list):
description_embedding = str(description_embedding)
query = f"""
INSERT INTO {self._get_table_name(table_name)}
(subject, predicate, object, description, subject_id, object_id,
weight, chunk_ids, parent_id, description_embedding, metadata)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
RETURNING id, subject, predicate, object, description, subject_id, object_id, weight, chunk_ids, parent_id, metadata
"""
params = [
subject,
predicate,
object,
description,
subject_id,
object_id,
weight,
chunk_ids,
parent_id,
description_embedding,
json.dumps(metadata) if metadata else None,
]
result = await self.connection_manager.fetchrow_query(
query=query,
params=params,
)
return Relationship(
id=result["id"],
subject=result["subject"],
predicate=result["predicate"],
object=result["object"],
description=result["description"],
subject_id=result["subject_id"],
object_id=result["object_id"],
weight=result["weight"],
chunk_ids=result["chunk_ids"],
parent_id=result["parent_id"],
metadata=result["metadata"],
)
async def get(
self,
parent_id: UUID,
store_type: StoreType,
offset: int,
limit: int,
relationship_ids: Optional[list[UUID]] = None,
entity_names: Optional[list[str]] = None,
relationship_types: Optional[list[str]] = None,
include_metadata: bool = False,
):
"""Get relationships from the specified store.
Args:
parent_id: UUID of the parent (collection_id or document_id)
store_type: Type of store (graph or document)
offset: Number of records to skip
limit: Maximum number of records to return (-1 for no limit)
relationship_ids: Optional list of specific relationship IDs to retrieve
entity_names: Optional list of entity names to filter by (matches subject or object)
relationship_types: Optional list of relationship types (predicates) to filter by
include_metadata: Whether to include metadata in the response
Returns:
Tuple of (list of relationships, total count)
"""
table_name = self._get_relationship_table_for_store(store_type)
conditions = ["parent_id = $1"]
params: list[Any] = [parent_id]
param_index = 2
if relationship_ids:
conditions.append(f"id = ANY(${param_index})")
params.append(relationship_ids)
param_index += 1
if entity_names:
conditions.append(
f"(subject = ANY(${param_index}) OR object = ANY(${param_index}))"
)
params.append(entity_names)
param_index += 1
if relationship_types:
conditions.append(f"predicate = ANY(${param_index})")
params.append(relationship_types)
param_index += 1
select_fields = """
id, subject, predicate, object, description,
subject_id, object_id, weight, chunk_ids,
parent_id
"""
if include_metadata:
select_fields += ", metadata"
# Count query
COUNT_QUERY = f"""
SELECT COUNT(*)
FROM {self._get_table_name(table_name)}
WHERE {" AND ".join(conditions)}
"""
count_params = params[: param_index - 1]
count = (
await self.connection_manager.fetch_query(
COUNT_QUERY, count_params
)
)[0]["count"]
# Main query
QUERY = f"""
SELECT {select_fields}
FROM {self._get_table_name(table_name)}
WHERE {" AND ".join(conditions)}
ORDER BY created_at
OFFSET ${param_index}
"""
params.append(offset)
param_index += 1
if limit != -1:
QUERY += f" LIMIT ${param_index}"
params.append(limit)
rows = await self.connection_manager.fetch_query(QUERY, params)
relationships = []
for row in rows:
relationship_dict = dict(row)
if include_metadata and isinstance(
relationship_dict["metadata"], str
):
with contextlib.suppress(json.JSONDecodeError):
relationship_dict["metadata"] = json.loads(
relationship_dict["metadata"]
)
elif not include_metadata:
relationship_dict.pop("metadata", None)
relationships.append(Relationship(**relationship_dict))
return relationships, count
async def update(
self,
relationship_id: UUID,
store_type: StoreType,
subject: Optional[str],
subject_id: Optional[UUID],
predicate: Optional[str],
object: Optional[str],
object_id: Optional[UUID],
description: Optional[str],
description_embedding: Optional[list[float] | str],
weight: Optional[float],
metadata: Optional[dict[str, Any] | str],
) -> Relationship:
"""Update multiple relationships in the specified store."""
table_name = self._get_relationship_table_for_store(store_type)
update_fields = []
params: list = []
param_index = 1
if isinstance(metadata, str):
with contextlib.suppress(json.JSONDecodeError):
metadata = json.loads(metadata)
if subject is not None:
update_fields.append(f"subject = ${param_index}")
params.append(subject)
param_index += 1
if subject_id is not None:
update_fields.append(f"subject_id = ${param_index}")
params.append(subject_id)
param_index += 1
if predicate is not None:
update_fields.append(f"predicate = ${param_index}")
params.append(predicate)
param_index += 1
if object is not None:
update_fields.append(f"object = ${param_index}")
params.append(object)
param_index += 1
if object_id is not None:
update_fields.append(f"object_id = ${param_index}")
params.append(object_id)
param_index += 1
if description is not None:
update_fields.append(f"description = ${param_index}")
params.append(description)
param_index += 1
if description_embedding is not None:
update_fields.append(f"description_embedding = ${param_index}")
params.append(description_embedding)
param_index += 1
if weight is not None:
update_fields.append(f"weight = ${param_index}")
params.append(weight)
param_index += 1
if not update_fields:
raise R2RException(status_code=400, message="No fields to update")
update_fields.append("updated_at = NOW()")
params.append(relationship_id)
query = f"""
UPDATE {self._get_table_name(table_name)}
SET {", ".join(update_fields)}
WHERE id = ${param_index}
RETURNING id, subject, predicate, object, description, subject_id, object_id, weight, chunk_ids, parent_id, metadata
"""
try:
result = await self.connection_manager.fetchrow_query(
query=query,
params=params,
)
return Relationship(
id=result["id"],
subject=result["subject"],
predicate=result["predicate"],
object=result["object"],
description=result["description"],
subject_id=result["subject_id"],
object_id=result["object_id"],
weight=result["weight"],
chunk_ids=result["chunk_ids"],
parent_id=result["parent_id"],
metadata=result["metadata"],
)
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"An error occurred while updating the relationship: {e}",
) from e
async def delete(
self,
parent_id: UUID,
relationship_ids: Optional[list[UUID]] = None,
store_type: StoreType = StoreType.GRAPHS,
) -> None:
"""Delete relationships from the specified store. If relationship_ids
is not provided, deletes all relationships for the given parent_id.
Args:
parent_id: UUID of the parent (collection_id or document_id)
relationship_ids: Optional list of specific relationship IDs to delete
store_type: Type of store (graph or document)
Returns:
List of deleted relationship IDs
Raises:
R2RException: If specific relationships were requested but not all found
"""
table_name = self._get_relationship_table_for_store(store_type)
if relationship_ids is None:
QUERY = f"""
DELETE FROM {self._get_table_name(table_name)}
WHERE parent_id = $1
RETURNING id
"""
results = await self.connection_manager.fetch_query(
QUERY, [parent_id]
)
else:
QUERY = f"""
DELETE FROM {self._get_table_name(table_name)}
WHERE id = ANY($1) AND parent_id = $2
RETURNING id
"""
results = await self.connection_manager.fetch_query(
QUERY, [relationship_ids, parent_id]
)
deleted_ids = [row["id"] for row in results]
if relationship_ids and len(deleted_ids) != len(relationship_ids):
raise R2RException(
f"Some relationships not found in {store_type} store or no permission to delete",
404,
)
async def export_to_csv(
self,
parent_id: UUID,
store_type: StoreType,
columns: Optional[list[str]] = None,
filters: Optional[dict] = None,
include_header: bool = True,
) -> tuple[str, IO]:
"""Creates a CSV file from the PostgreSQL data and returns the path to
the temp file."""
valid_columns = {
"id",
"subject",
"predicate",
"object",
"description",
"subject_id",
"object_id",
"weight",
"chunk_ids",
"parent_id",
"metadata",
"created_at",
"updated_at",
}
if not columns:
columns = list(valid_columns)
elif invalid_cols := set(columns) - valid_columns:
raise ValueError(f"Invalid columns: {invalid_cols}")
select_stmt = f"""
SELECT
id::text,
subject,
predicate,
object,
description,
subject_id::text,
object_id::text,
weight,
chunk_ids::text,
parent_id::text,
metadata::text,
to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') AS created_at,
to_char(updated_at, 'YYYY-MM-DD HH24:MI:SS') AS updated_at
FROM {self._get_table_name(self._get_relationship_table_for_store(store_type))}
"""
conditions = ["parent_id = $1"]
params: list[Any] = [parent_id]
param_index = 2
if filters:
for field, value in filters.items():
if field not in valid_columns:
continue
if isinstance(value, dict):
for op, val in value.items():
if op == "$eq":
conditions.append(f"{field} = ${param_index}")
params.append(val)
param_index += 1
elif op == "$gt":
conditions.append(f"{field} > ${param_index}")
params.append(val)
param_index += 1
elif op == "$lt":
conditions.append(f"{field} < ${param_index}")
params.append(val)
param_index += 1
else:
# Direct equality
conditions.append(f"{field} = ${param_index}")
params.append(value)
param_index += 1
if conditions:
select_stmt = f"{select_stmt} WHERE {' AND '.join(conditions)}"
select_stmt = f"{select_stmt} ORDER BY created_at DESC"
temp_file = None
try:
temp_file = tempfile.NamedTemporaryFile(
mode="w", delete=True, suffix=".csv"
)
writer = csv.writer(temp_file, quoting=csv.QUOTE_ALL)
async with self.connection_manager.pool.get_connection() as conn: # type: ignore
async with conn.transaction():
cursor = await conn.cursor(select_stmt, *params)
if include_header:
writer.writerow(columns)
chunk_size = 1000
while True:
rows = await cursor.fetch(chunk_size)
if not rows:
break
for row in rows:
writer.writerow(row)
temp_file.flush()
return temp_file.name, temp_file
except Exception as e:
if temp_file:
temp_file.close()
raise HTTPException(
status_code=500,
detail=f"Failed to export data: {str(e)}",
) from e
class PostgresCommunitiesHandler(Handler):
def __init__(self, *args: Any, **kwargs: Any) -> None:
self.project_name: str = kwargs.get("project_name") # type: ignore
self.connection_manager: PostgresConnectionManager = kwargs.get(
"connection_manager"
) # type: ignore
self.dimension: int = kwargs.get("dimension") # type: ignore
self.quantization_type: VectorQuantizationType = kwargs.get(
"quantization_type"
) # type: ignore
async def create_tables(self) -> None:
vector_column_str = _get_vector_column_str(
self.dimension, self.quantization_type
)
query = f"""
CREATE TABLE IF NOT EXISTS {self._get_table_name("graphs_communities")} (
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
collection_id UUID,
community_id UUID,
level INT,
name TEXT NOT NULL,
summary TEXT NOT NULL,
findings TEXT[],
rating FLOAT,
rating_explanation TEXT,
description_embedding {vector_column_str} NOT NULL,
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
metadata JSONB,
UNIQUE (community_id, level, collection_id)
);"""
await self.connection_manager.execute_query(query)
async def create(
self,
parent_id: UUID,
store_type: StoreType,
name: str,
summary: str,
findings: Optional[list[str]],
rating: Optional[float],
rating_explanation: Optional[str],
description_embedding: Optional[list[float] | str] = None,
) -> Community:
table_name = "graphs_communities"
if isinstance(description_embedding, list):
description_embedding = str(description_embedding)
query = f"""
INSERT INTO {self._get_table_name(table_name)}
(collection_id, name, summary, findings, rating, rating_explanation, description_embedding)
VALUES ($1, $2, $3, $4, $5, $6, $7)
RETURNING id, collection_id, name, summary, findings, rating, rating_explanation, created_at, updated_at
"""
params = [
parent_id,
name,
summary,
findings,
rating,
rating_explanation,
description_embedding,
]
try:
result = await self.connection_manager.fetchrow_query(
query=query,
params=params,
)
return Community(
id=result["id"],
collection_id=result["collection_id"],
name=result["name"],
summary=result["summary"],
findings=result["findings"],
rating=result["rating"],
rating_explanation=result["rating_explanation"],
created_at=result["created_at"],
updated_at=result["updated_at"],
)
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"An error occurred while creating the community: {e}",
) from e
async def update(
self,
community_id: UUID,
store_type: StoreType,
name: Optional[str] = None,
summary: Optional[str] = None,
summary_embedding: Optional[list[float] | str] = None,
findings: Optional[list[str]] = None,
rating: Optional[float] = None,
rating_explanation: Optional[str] = None,
) -> Community:
table_name = "graphs_communities"
update_fields = []
params: list[Any] = []
param_index = 1
if name is not None:
update_fields.append(f"name = ${param_index}")
params.append(name)
param_index += 1
if summary is not None:
update_fields.append(f"summary = ${param_index}")
params.append(summary)
param_index += 1
if summary_embedding is not None:
update_fields.append(f"description_embedding = ${param_index}")
params.append(summary_embedding)
param_index += 1
if findings is not None:
update_fields.append(f"findings = ${param_index}")
params.append(findings)
param_index += 1
if rating is not None:
update_fields.append(f"rating = ${param_index}")
params.append(rating)
param_index += 1
if rating_explanation is not None:
update_fields.append(f"rating_explanation = ${param_index}")
params.append(rating_explanation)
param_index += 1
if not update_fields:
raise R2RException(status_code=400, message="No fields to update")
update_fields.append("updated_at = NOW()")
params.append(community_id)
query = f"""
UPDATE {self._get_table_name(table_name)}
SET {", ".join(update_fields)}
WHERE id = ${param_index}\
RETURNING id, community_id, name, summary, findings, rating, rating_explanation, created_at, updated_at
"""
try:
result = await self.connection_manager.fetchrow_query(
query, params
)
return Community(
id=result["id"],
community_id=result["community_id"],
name=result["name"],
summary=result["summary"],
findings=result["findings"],
rating=result["rating"],
rating_explanation=result["rating_explanation"],
created_at=result["created_at"],
updated_at=result["updated_at"],
)
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"An error occurred while updating the community: {e}",
) from e
async def delete(
self,
parent_id: UUID,
community_id: UUID,
) -> None:
table_name = "graphs_communities"
params = [community_id, parent_id]
# Delete the community
query = f"""
DELETE FROM {self._get_table_name(table_name)}
WHERE id = $1 AND collection_id = $2
"""
try:
await self.connection_manager.execute_query(query, params)
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"An error occurred while deleting the community: {e}",
) from e
async def delete_all_communities(
self,
parent_id: UUID,
) -> None:
table_name = "graphs_communities"
params = [parent_id]
# Delete all communities for the parent_id
query = f"""
DELETE FROM {self._get_table_name(table_name)}
WHERE collection_id = $1
"""
try:
await self.connection_manager.execute_query(query, params)
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"An error occurred while deleting communities: {e}",
) from e
async def get(
self,
parent_id: UUID,
store_type: StoreType,
offset: int,
limit: int,
community_ids: Optional[list[UUID]] = None,
community_names: Optional[list[str]] = None,
include_embeddings: bool = False,
):
"""Retrieve communities from the specified store."""
# Do we ever want to get communities from document store?
table_name = "graphs_communities"
conditions = ["collection_id = $1"]
params: list[Any] = [parent_id]
param_index = 2
if community_ids:
conditions.append(f"id = ANY(${param_index})")
params.append(community_ids)
param_index += 1
if community_names:
conditions.append(f"name = ANY(${param_index})")
params.append(community_names)
param_index += 1
select_fields = """
id, community_id, name, summary, findings, rating,
rating_explanation, level, created_at, updated_at
"""
if include_embeddings:
select_fields += ", description_embedding"
COUNT_QUERY = f"""
SELECT COUNT(*)
FROM {self._get_table_name(table_name)}
WHERE {" AND ".join(conditions)}
"""
count = (
await self.connection_manager.fetch_query(
COUNT_QUERY, params[: param_index - 1]
)
)[0]["count"]
QUERY = f"""
SELECT {select_fields}
FROM {self._get_table_name(table_name)}
WHERE {" AND ".join(conditions)}
ORDER BY created_at
OFFSET ${param_index}
"""
params.append(offset)
param_index += 1
if limit != -1:
QUERY += f" LIMIT ${param_index}"
params.append(limit)
rows = await self.connection_manager.fetch_query(QUERY, params)
communities = []
for row in rows:
community_dict = dict(row)
communities.append(Community(**community_dict))
return communities, count
async def export_to_csv(
self,
parent_id: UUID,
store_type: StoreType,
columns: Optional[list[str]] = None,
filters: Optional[dict] = None,
include_header: bool = True,
) -> tuple[str, IO]:
"""Creates a CSV file from the PostgreSQL data and returns the path to
the temp file."""
valid_columns = {
"id",
"collection_id",
"community_id",
"level",
"name",
"summary",
"findings",
"rating",
"rating_explanation",
"created_at",
"updated_at",
"metadata",
}
if not columns:
columns = list(valid_columns)
elif invalid_cols := set(columns) - valid_columns:
raise ValueError(f"Invalid columns: {invalid_cols}")
table_name = "graphs_communities"
select_stmt = f"""
SELECT
id::text,
collection_id::text,
community_id::text,
level,
name,
summary,
findings::text,
rating,
rating_explanation,
to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') AS created_at,
to_char(updated_at, 'YYYY-MM-DD HH24:MI:SS') AS updated_at,
metadata::text
FROM {self._get_table_name(table_name)}
"""
conditions = ["collection_id = $1"]
params: list[Any] = [parent_id]
param_index = 2
if filters:
for field, value in filters.items():
if field not in valid_columns:
continue
if isinstance(value, dict):
for op, val in value.items():
if op == "$eq":
conditions.append(f"{field} = ${param_index}")
params.append(val)
param_index += 1
elif op == "$gt":
conditions.append(f"{field} > ${param_index}")
params.append(val)
param_index += 1
elif op == "$lt":
conditions.append(f"{field} < ${param_index}")
params.append(val)
param_index += 1
else:
# Direct equality
conditions.append(f"{field} = ${param_index}")
params.append(value)
param_index += 1
if conditions:
select_stmt = f"{select_stmt} WHERE {' AND '.join(conditions)}"
select_stmt = f"{select_stmt} ORDER BY created_at DESC"
temp_file = None
try:
temp_file = tempfile.NamedTemporaryFile(
mode="w", delete=True, suffix=".csv"
)
writer = csv.writer(temp_file, quoting=csv.QUOTE_ALL)
async with self.connection_manager.pool.get_connection() as conn: # type: ignore
async with conn.transaction():
cursor = await conn.cursor(select_stmt, *params)
if include_header:
writer.writerow(columns)
chunk_size = 1000
while True:
rows = await cursor.fetch(chunk_size)
if not rows:
break
for row in rows:
writer.writerow(row)
temp_file.flush()
return temp_file.name, temp_file
except Exception as e:
if temp_file:
temp_file.close()
raise HTTPException(
status_code=500,
detail=f"Failed to export data: {str(e)}",
) from e
class PostgresGraphsHandler(Handler):
"""Handler for Knowledge Graph METHODS in PostgreSQL."""
TABLE_NAME = "graphs"
def __init__(
self,
*args: Any,
**kwargs: Any,
) -> None:
self.project_name: str = kwargs.get("project_name") # type: ignore
self.connection_manager: PostgresConnectionManager = kwargs.get(
"connection_manager"
) # type: ignore
self.dimension: int = kwargs.get("dimension") # type: ignore
self.quantization_type: VectorQuantizationType = kwargs.get(
"quantization_type"
) # type: ignore
self.collections_handler: PostgresCollectionsHandler = kwargs.get(
"collections_handler"
) # type: ignore
self.entities = PostgresEntitiesHandler(*args, **kwargs)
self.relationships = PostgresRelationshipsHandler(*args, **kwargs)
self.communities = PostgresCommunitiesHandler(*args, **kwargs)
self.handlers = [
self.entities,
self.relationships,
self.communities,
]
async def create_tables(self) -> None:
"""Create the graph tables with mandatory collection_id support."""
QUERY = f"""
CREATE TABLE IF NOT EXISTS {self._get_table_name(PostgresGraphsHandler.TABLE_NAME)} (
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
collection_id UUID NOT NULL,
name TEXT NOT NULL,
description TEXT,
status TEXT NOT NULL,
document_ids UUID[],
metadata JSONB,
created_at TIMESTAMPTZ DEFAULT NOW(),
updated_at TIMESTAMPTZ DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS graph_collection_id_idx
ON {self._get_table_name("graphs")} (collection_id);
"""
await self.connection_manager.execute_query(QUERY)
for handler in self.handlers:
await handler.create_tables()
async def create(
self,
collection_id: UUID,
name: Optional[str] = None,
description: Optional[str] = None,
status: str = "pending",
) -> GraphResponse:
"""Create a new graph associated with a collection."""
name = name or f"Graph {collection_id}"
description = description or ""
query = f"""
INSERT INTO {self._get_table_name(PostgresGraphsHandler.TABLE_NAME)}
(id, collection_id, name, description, status)
VALUES ($1, $2, $3, $4, $5)
RETURNING id, collection_id, name, description, status, created_at, updated_at, document_ids
"""
params = [
collection_id,
collection_id,
name,
description,
status,
]
try:
result = await self.connection_manager.fetchrow_query(
query=query,
params=params,
)
return GraphResponse(
id=result["id"],
collection_id=result["collection_id"],
name=result["name"],
description=result["description"],
status=result["status"],
created_at=result["created_at"],
updated_at=result["updated_at"],
document_ids=result["document_ids"] or [],
)
except UniqueViolationError:
raise R2RException(
message="Graph with this ID already exists",
status_code=409,
) from None
async def reset(self, parent_id: UUID) -> None:
"""Completely reset a graph and all associated data."""
await self.entities.delete(
parent_id=parent_id, store_type=StoreType.GRAPHS
)
await self.relationships.delete(
parent_id=parent_id, store_type=StoreType.GRAPHS
)
await self.communities.delete_all_communities(parent_id=parent_id)
# Now, update the graph record to remove any attached document IDs.
# This sets document_ids to an empty UUID array.
query = f"""
UPDATE {self._get_table_name(PostgresGraphsHandler.TABLE_NAME)}
SET document_ids = ARRAY[]::uuid[]
WHERE id = $1;
"""
await self.connection_manager.execute_query(query, [parent_id])
async def list_graphs(
self,
offset: int,
limit: int,
# filter_user_ids: Optional[list[UUID]] = None,
filter_graph_ids: Optional[list[UUID]] = None,
filter_collection_id: Optional[UUID] = None,
) -> dict[str, list[GraphResponse] | int]:
conditions = []
params: list[Any] = []
param_index = 1
if filter_graph_ids:
conditions.append(f"id = ANY(${param_index})")
params.append(filter_graph_ids)
param_index += 1
# if filter_user_ids:
# conditions.append(f"user_id = ANY(${param_index})")
# params.append(filter_user_ids)
# param_index += 1
if filter_collection_id:
conditions.append(f"collection_id = ${param_index}")
params.append(filter_collection_id)
param_index += 1
where_clause = (
f"WHERE {' AND '.join(conditions)}" if conditions else ""
)
query = f"""
WITH RankedGraphs AS (
SELECT
id, collection_id, name, description, status, created_at, updated_at, document_ids,
COUNT(*) OVER() as total_entries,
ROW_NUMBER() OVER (PARTITION BY collection_id ORDER BY created_at DESC) as rn
FROM {self._get_table_name(PostgresGraphsHandler.TABLE_NAME)}
{where_clause}
)
SELECT * FROM RankedGraphs
WHERE rn = 1
ORDER BY created_at DESC
OFFSET ${param_index} LIMIT ${param_index + 1}
"""
params.extend([offset, limit])
try:
results = await self.connection_manager.fetch_query(query, params)
if not results:
return {"results": [], "total_entries": 0}
total_entries = results[0]["total_entries"] if results else 0
graphs = [
GraphResponse(
id=row["id"],
document_ids=row["document_ids"] or [],
name=row["name"],
collection_id=row["collection_id"],
description=row["description"],
status=row["status"],
created_at=row["created_at"],
updated_at=row["updated_at"],
)
for row in results
]
return {"results": graphs, "total_entries": total_entries}
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"An error occurred while fetching graphs: {e}",
) from e
async def get(
self, offset: int, limit: int, graph_id: Optional[UUID] = None
):
if graph_id is None:
params = [offset, limit]
QUERY = f"""
SELECT * FROM {self._get_table_name(PostgresGraphsHandler.TABLE_NAME)}
OFFSET $1 LIMIT $2
"""
ret = await self.connection_manager.fetch_query(QUERY, params)
COUNT_QUERY = f"""
SELECT COUNT(*) FROM {self._get_table_name(PostgresGraphsHandler.TABLE_NAME)}
"""
count = (await self.connection_manager.fetch_query(COUNT_QUERY))[
0
]["count"]
return {
"results": [Graph(**row) for row in ret],
"total_entries": count,
}
else:
QUERY = f"""
SELECT * FROM {self._get_table_name(PostgresGraphsHandler.TABLE_NAME)} WHERE id = $1
"""
params = [graph_id] # type: ignore
return {
"results": [
Graph(
**await self.connection_manager.fetchrow_query(
QUERY, params
)
)
]
}
async def add_documents(self, id: UUID, document_ids: list[UUID]) -> bool:
"""Add documents to the graph by copying their entities and
relationships."""
# Copy entities from document_entity to graphs_entities
ENTITY_COPY_QUERY = f"""
INSERT INTO {self._get_table_name("graphs_entities")} (
name, category, description, parent_id, description_embedding,
chunk_ids, metadata
)
SELECT
name, category, description, $1, description_embedding,
chunk_ids, metadata
FROM {self._get_table_name("documents_entities")}
WHERE parent_id = ANY($2)
"""
await self.connection_manager.execute_query(
ENTITY_COPY_QUERY, [id, document_ids]
)
# Copy relationships from documents_relationships to graphs_relationships
RELATIONSHIP_COPY_QUERY = f"""
INSERT INTO {self._get_table_name("graphs_relationships")} (
subject, predicate, object, description, subject_id, object_id,
weight, chunk_ids, parent_id, metadata, description_embedding
)
SELECT
subject, predicate, object, description, subject_id, object_id,
weight, chunk_ids, $1, metadata, description_embedding
FROM {self._get_table_name("documents_relationships")}
WHERE parent_id = ANY($2)
"""
await self.connection_manager.execute_query(
RELATIONSHIP_COPY_QUERY, [id, document_ids]
)
# Add document_ids to the graph
UPDATE_GRAPH_QUERY = f"""
UPDATE {self._get_table_name(PostgresGraphsHandler.TABLE_NAME)}
SET document_ids = array_cat(
CASE
WHEN document_ids IS NULL THEN ARRAY[]::uuid[]
ELSE document_ids
END,
$2::uuid[]
)
WHERE id = $1
"""
await self.connection_manager.execute_query(
UPDATE_GRAPH_QUERY, [id, document_ids]
)
return True
async def update(
self,
collection_id: UUID,
name: Optional[str] = None,
description: Optional[str] = None,
) -> GraphResponse:
"""Update an existing graph."""
update_fields = []
params: list = []
param_index = 1
if name is not None:
update_fields.append(f"name = ${param_index}")
params.append(name)
param_index += 1
if description is not None:
update_fields.append(f"description = ${param_index}")
params.append(description)
param_index += 1
if not update_fields:
raise R2RException(status_code=400, message="No fields to update")
update_fields.append("updated_at = NOW()")
params.append(collection_id)
query = f"""
UPDATE {self._get_table_name(PostgresGraphsHandler.TABLE_NAME)}
SET {", ".join(update_fields)}
WHERE id = ${param_index}
RETURNING id, name, description, status, created_at, updated_at, collection_id, document_ids
"""
try:
result = await self.connection_manager.fetchrow_query(
query, params
)
if not result:
raise R2RException(status_code=404, message="Graph not found")
return GraphResponse(
id=result["id"],
collection_id=result["collection_id"],
name=result["name"],
description=result["description"],
status=result["status"],
created_at=result["created_at"],
document_ids=result["document_ids"] or [],
updated_at=result["updated_at"],
)
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"An error occurred while updating the graph: {e}",
) from e
async def get_entities(
self,
parent_id: UUID,
offset: int,
limit: int,
entity_ids: Optional[list[UUID]] = None,
entity_names: Optional[list[str]] = None,
include_embeddings: bool = False,
) -> tuple[list[Entity], int]:
"""Get entities for a graph.
Args:
offset: Number of records to skip
limit: Maximum number of records to return (-1 for no limit)
parent_id: UUID of the collection
entity_ids: Optional list of entity IDs to filter by
entity_names: Optional list of entity names to filter by
include_embeddings: Whether to include embeddings in the response
Returns:
Tuple of (list of entities, total count)
"""
conditions = ["parent_id = $1"]
params: list[Any] = [parent_id]
param_index = 2
if entity_ids:
conditions.append(f"id = ANY(${param_index})")
params.append(entity_ids)
param_index += 1
if entity_names:
conditions.append(f"name = ANY(${param_index})")
params.append(entity_names)
param_index += 1
# Count query - uses the same conditions but without offset/limit
COUNT_QUERY = f"""
SELECT COUNT(*)
FROM {self._get_table_name("graphs_entities")}
WHERE {" AND ".join(conditions)}
"""
count = (
await self.connection_manager.fetch_query(COUNT_QUERY, params)
)[0]["count"]
# Define base columns to select
select_fields = """
id, name, category, description, parent_id,
chunk_ids, metadata
"""
if include_embeddings:
select_fields += ", description_embedding"
# Main query for fetching entities with pagination
QUERY = f"""
SELECT {select_fields}
FROM {self._get_table_name("graphs_entities")}
WHERE {" AND ".join(conditions)}
ORDER BY created_at
OFFSET ${param_index}
"""
params.append(offset)
param_index += 1
if limit != -1:
QUERY += f" LIMIT ${param_index}"
params.append(limit)
rows = await self.connection_manager.fetch_query(QUERY, params)
entities = []
for row in rows:
entity_dict = dict(row)
if isinstance(entity_dict["metadata"], str):
with contextlib.suppress(json.JSONDecodeError):
entity_dict["metadata"] = json.loads(
entity_dict["metadata"]
)
entities.append(Entity(**entity_dict))
return entities, count
async def get_relationships(
self,
parent_id: UUID,
offset: int,
limit: int,
relationship_ids: Optional[list[UUID]] = None,
relationship_types: Optional[list[str]] = None,
include_embeddings: bool = False,
) -> tuple[list[Relationship], int]:
"""Get relationships for a graph.
Args:
parent_id: UUID of the graph
offset: Number of records to skip
limit: Maximum number of records to return (-1 for no limit)
relationship_ids: Optional list of relationship IDs to filter by
relationship_types: Optional list of relationship types to filter by
include_metadata: Whether to include metadata in the response
Returns:
Tuple of (list of relationships, total count)
"""
conditions = ["parent_id = $1"]
params: list[Any] = [parent_id]
param_index = 2
if relationship_ids:
conditions.append(f"id = ANY(${param_index})")
params.append(relationship_ids)
param_index += 1
if relationship_types:
conditions.append(f"predicate = ANY(${param_index})")
params.append(relationship_types)
param_index += 1
# Count query - uses the same conditions but without offset/limit
COUNT_QUERY = f"""
SELECT COUNT(*)
FROM {self._get_table_name("graphs_relationships")}
WHERE {" AND ".join(conditions)}
"""
count = (
await self.connection_manager.fetch_query(COUNT_QUERY, params)
)[0]["count"]
# Define base columns to select
select_fields = """
id, subject, predicate, object, weight, chunk_ids, parent_id, metadata
"""
if include_embeddings:
select_fields += ", description_embedding"
# Main query for fetching relationships with pagination
QUERY = f"""
SELECT {select_fields}
FROM {self._get_table_name("graphs_relationships")}
WHERE {" AND ".join(conditions)}
ORDER BY created_at
OFFSET ${param_index}
"""
params.append(offset)
param_index += 1
if limit != -1:
QUERY += f" LIMIT ${param_index}"
params.append(limit)
rows = await self.connection_manager.fetch_query(QUERY, params)
relationships = []
for row in rows:
relationship_dict = dict(row)
if isinstance(relationship_dict["metadata"], str):
with contextlib.suppress(json.JSONDecodeError):
relationship_dict["metadata"] = json.loads(
relationship_dict["metadata"]
)
relationships.append(Relationship(**relationship_dict))
return relationships, count
async def add_entities(
self,
entities: list[Entity],
table_name: str,
conflict_columns: list[str] | None = None,
) -> asyncpg.Record:
"""Upsert entities into the entities_raw table. These are raw entities
extracted from the document.
Args:
entities: list[Entity]: list of entities to upsert
collection_name: str: name of the collection
Returns:
result: asyncpg.Record: result of the upsert operation
"""
if not conflict_columns:
conflict_columns = []
cleaned_entities = []
for entity in entities:
entity_dict = entity.to_dict()
entity_dict["chunk_ids"] = (
entity_dict["chunk_ids"]
if entity_dict.get("chunk_ids")
else []
)
entity_dict["description_embedding"] = (
str(entity_dict["description_embedding"])
if entity_dict.get("description_embedding") # type: ignore
else None
)
cleaned_entities.append(entity_dict)
return await _add_objects(
objects=cleaned_entities,
full_table_name=self._get_table_name(table_name),
connection_manager=self.connection_manager,
conflict_columns=conflict_columns,
)
async def get_all_relationships(
self,
collection_id: UUID | None,
graph_id: UUID | None,
document_ids: Optional[list[UUID]] = None,
) -> list[Relationship]:
QUERY = f"""
SELECT id, subject, predicate, weight, object, parent_id FROM {self._get_table_name("graphs_relationships")} WHERE parent_id = ANY($1)
"""
relationships = await self.connection_manager.fetch_query(
QUERY, [collection_id]
)
return [Relationship(**relationship) for relationship in relationships]
async def has_document(self, graph_id: UUID, document_id: UUID) -> bool:
"""Check if a document exists in the graph's document_ids array.
Args:
graph_id (UUID): ID of the graph to check
document_id (UUID): ID of the document to look for
Returns:
bool: True if document exists in graph, False otherwise
Raises:
R2RException: If graph not found
"""
QUERY = f"""
SELECT EXISTS (
SELECT 1
FROM {self._get_table_name("graphs")}
WHERE id = $1
AND document_ids IS NOT NULL
AND $2 = ANY(document_ids)
) as exists;
"""
result = await self.connection_manager.fetchrow_query(
QUERY, [graph_id, document_id]
)
if result is None:
raise R2RException(f"Graph {graph_id} not found", 404)
return result["exists"]
async def get_communities(
self,
parent_id: UUID,
offset: int,
limit: int,
community_ids: Optional[list[UUID]] = None,
include_embeddings: bool = False,
) -> tuple[list[Community], int]:
"""Get communities for a graph.
Args:
collection_id: UUID of the collection
offset: Number of records to skip
limit: Maximum number of records to return (-1 for no limit)
community_ids: Optional list of community IDs to filter by
include_embeddings: Whether to include embeddings in the response
Returns:
Tuple of (list of communities, total count)
"""
conditions = ["collection_id = $1"]
params: list[Any] = [parent_id]
param_index = 2
if community_ids:
conditions.append(f"id = ANY(${param_index})")
params.append(community_ids)
param_index += 1
select_fields = """
id, collection_id, name, summary, findings, rating, rating_explanation
"""
if include_embeddings:
select_fields += ", description_embedding"
COUNT_QUERY = f"""
SELECT COUNT(*)
FROM {self._get_table_name("graphs_communities")}
WHERE {" AND ".join(conditions)}
"""
count = (
await self.connection_manager.fetch_query(COUNT_QUERY, params)
)[0]["count"]
QUERY = f"""
SELECT {select_fields}
FROM {self._get_table_name("graphs_communities")}
WHERE {" AND ".join(conditions)}
ORDER BY created_at
OFFSET ${param_index}
"""
params.append(offset)
param_index += 1
if limit != -1:
QUERY += f" LIMIT ${param_index}"
params.append(limit)
rows = await self.connection_manager.fetch_query(QUERY, params)
communities = []
for row in rows:
community_dict = dict(row)
communities.append(Community(**community_dict))
return communities, count
async def add_community(self, community: Community) -> None:
# TODO: Fix in the short term.
# we need to do this because postgres insert needs to be a string
community.description_embedding = str(community.description_embedding) # type: ignore[assignment]
non_null_attrs = {
k: v for k, v in community.__dict__.items() if v is not None
}
columns = ", ".join(non_null_attrs.keys())
placeholders = ", ".join(
f"${i + 1}" for i in range(len(non_null_attrs))
)
conflict_columns = ", ".join(
[f"{k} = EXCLUDED.{k}" for k in non_null_attrs]
)
QUERY = f"""
INSERT INTO {self._get_table_name("graphs_communities")} ({columns})
VALUES ({placeholders})
ON CONFLICT (community_id, level, collection_id) DO UPDATE SET
{conflict_columns}
"""
await self.connection_manager.execute_many(
QUERY, [tuple(non_null_attrs.values())]
)
async def delete(self, collection_id: UUID) -> None:
graphs = await self.get(graph_id=collection_id, offset=0, limit=-1)
if len(graphs["results"]) == 0:
raise R2RException(
message=f"Graph not found for collection {collection_id}",
status_code=404,
)
await self.reset(collection_id)
# set status to PENDING for this collection.
QUERY = f"""
UPDATE {self._get_table_name("collections")} SET graph_cluster_status = $1 WHERE id = $2
"""
await self.connection_manager.execute_query(
QUERY, [GraphExtractionStatus.PENDING, collection_id]
)
# Delete the graph
QUERY = f"""
DELETE FROM {self._get_table_name("graphs")} WHERE collection_id = $1
"""
try:
await self.connection_manager.execute_query(QUERY, [collection_id])
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"An error occurred while deleting the graph: {e}",
) from e
async def perform_graph_clustering(
self,
collection_id: UUID,
leiden_params: dict[str, Any],
) -> Tuple[int, Any]:
"""Calls the external clustering service to cluster the graph."""
offset = 0
page_size = 1000
all_relationships = []
while True:
relationships, count = await self.relationships.get(
parent_id=collection_id,
store_type=StoreType.GRAPHS,
offset=offset,
limit=page_size,
)
if not relationships:
break
all_relationships.extend(relationships)
offset += len(relationships)
if offset >= count:
break
logger.info(
f"Clustering over {len(all_relationships)} relationships for {collection_id} with settings: {leiden_params}"
)
if len(all_relationships) == 0:
raise R2RException(
message="No relationships found for clustering",
status_code=400,
)
return await self._cluster_and_add_community_info(
relationships=all_relationships,
leiden_params=leiden_params,
collection_id=collection_id,
)
async def _call_clustering_service(
self, relationships: list[Relationship], leiden_params: dict[str, Any]
) -> list[dict]:
"""Calls the external Graspologic clustering service, sending
relationships and parameters.
Expects a response with 'communities' field.
"""
# Convert relationships to a JSON-friendly format
rel_data = []
for r in relationships:
rel_data.append(
{
"id": str(r.id),
"subject": r.subject,
"object": r.object,
"weight": r.weight if r.weight is not None else 1.0,
}
)
endpoint = os.environ.get("CLUSTERING_SERVICE_URL")
if not endpoint:
raise ValueError("CLUSTERING_SERVICE_URL not set.")
url = f"{endpoint}/cluster"
payload = {"relationships": rel_data, "leiden_params": leiden_params}
async with httpx.AsyncClient() as client:
response = await client.post(url, json=payload, timeout=3600)
response.raise_for_status()
data = response.json()
return data.get("communities", [])
async def _create_graph_and_cluster(
self,
relationships: list[Relationship],
leiden_params: dict[str, Any],
) -> Any:
"""Create a graph and cluster it."""
return await self._call_clustering_service(
relationships, leiden_params
)
async def _cluster_and_add_community_info(
self,
relationships: list[Relationship],
leiden_params: dict[str, Any],
collection_id: UUID,
) -> Tuple[int, Any]:
logger.info(f"Creating graph and clustering for {collection_id}")
await asyncio.sleep(0.1)
start_time = time.time()
hierarchical_communities = await self._create_graph_and_cluster(
relationships=relationships,
leiden_params=leiden_params,
)
logger.info(
f"Computing Leiden communities completed, time {time.time() - start_time:.2f} seconds."
)
if not hierarchical_communities:
num_communities = 0
else:
num_communities = (
max(item["cluster"] for item in hierarchical_communities) + 1
)
logger.info(
f"Generated {num_communities} communities, time {time.time() - start_time:.2f} seconds."
)
return num_communities, hierarchical_communities
async def get_entity_map(
self, offset: int, limit: int, document_id: UUID
) -> dict[str, dict[str, list[dict[str, Any]]]]:
QUERY1 = f"""
WITH entities_list AS (
SELECT DISTINCT name
FROM {self._get_table_name("documents_entities")}
WHERE parent_id = $1
ORDER BY name ASC
LIMIT {limit} OFFSET {offset}
)
SELECT e.name, e.description, e.category,
(SELECT array_agg(DISTINCT x) FROM unnest(e.chunk_ids) x) AS chunk_ids,
e.parent_id
FROM {self._get_table_name("documents_entities")} e
JOIN entities_list el ON e.name = el.name
GROUP BY e.name, e.description, e.category, e.chunk_ids, e.parent_id
ORDER BY e.name;"""
entities_list = await self.connection_manager.fetch_query(
QUERY1, [document_id]
)
entities_list = [Entity(**entity) for entity in entities_list]
QUERY2 = f"""
WITH entities_list AS (
SELECT DISTINCT name
FROM {self._get_table_name("documents_entities")}
WHERE parent_id = $1
ORDER BY name ASC
LIMIT {limit} OFFSET {offset}
)
SELECT DISTINCT t.subject, t.predicate, t.object, t.weight, t.description,
(SELECT array_agg(DISTINCT x) FROM unnest(t.chunk_ids) x) AS chunk_ids, t.parent_id
FROM {self._get_table_name("documents_relationships")} t
JOIN entities_list el ON t.subject = el.name
ORDER BY t.subject, t.predicate, t.object;
"""
relationships_list = await self.connection_manager.fetch_query(
QUERY2, [document_id]
)
relationships_list = [
Relationship(**relationship) for relationship in relationships_list
]
entity_map: dict[str, dict[str, list[Any]]] = {}
for entity in entities_list:
if entity.name not in entity_map:
entity_map[entity.name] = {"entities": [], "relationships": []}
entity_map[entity.name]["entities"].append(entity)
for relationship in relationships_list:
if relationship.subject in entity_map:
entity_map[relationship.subject]["relationships"].append(
relationship
)
if relationship.object in entity_map:
entity_map[relationship.object]["relationships"].append(
relationship
)
return entity_map
async def graph_search(
self, query: str, **kwargs: Any
) -> AsyncGenerator[Any, None]:
"""Perform semantic search with similarity scores while maintaining
exact same structure."""
query_embedding = kwargs.get("query_embedding", None)
if query_embedding is None:
raise ValueError(
"query_embedding must be provided for semantic search"
)
search_type = kwargs.get(
"search_type", "entities"
) # entities | relationships | communities
embedding_type = kwargs.get("embedding_type", "description_embedding")
property_names = kwargs.get("property_names", ["name", "description"])
# Add metadata if not present
if "metadata" not in property_names:
property_names.append("metadata")
filters = kwargs.get("filters", {})
limit = kwargs.get("limit", 10)
use_fulltext_search = kwargs.get("use_fulltext_search", True)
use_hybrid_search = kwargs.get("use_hybrid_search", True)
if use_hybrid_search or use_fulltext_search:
logger.warning(
"Hybrid and fulltext search not supported for graph search, ignoring."
)
table_name = f"graphs_{search_type}"
property_names_str = ", ".join(property_names)
# Build the WHERE clause from filters
params: list[str | int | bytes] = [
json.dumps(query_embedding),
limit,
]
conditions_clause = self._build_filters(filters, params, search_type)
where_clause = (
f"WHERE {conditions_clause}" if conditions_clause else ""
)
# Construct the query
# Note: For vector similarity, we use <=> for distance. The smaller the number, the more similar.
# We'll convert that to similarity_score by doing (1 - distance).
QUERY = f"""
SELECT
{property_names_str},
({embedding_type} <=> $1) as similarity_score
FROM {self._get_table_name(table_name)}
{where_clause}
ORDER BY {embedding_type} <=> $1
LIMIT $2;
"""
results = await self.connection_manager.fetch_query(
QUERY, tuple(params)
)
for result in results:
output = {
prop: result[prop] for prop in property_names if prop in result
}
output["similarity_score"] = (
1 - float(result["similarity_score"])
if result.get("similarity_score")
else "n/a"
)
yield output
def _build_filters(
self, filter_dict: dict, parameters: list[Any], search_type: str
) -> str:
"""Build a WHERE clause from a nested filter dictionary for the graph
search.
- If search_type == "communities", we normally filter by `collection_id`.
- Otherwise (entities/relationships), we normally filter by `parent_id`.
- If user provides `"collection_ids": {...}`, we interpret that as wanting
to filter by multiple collection IDs (i.e. 'parent_id IN (...)' or
'collection_id IN (...)').
"""
# The usual "base" column used by your code
base_id_column = (
"collection_id" if search_type == "communities" else "parent_id"
)
def parse_condition(key: str, value: Any) -> str:
# ----------------------------------------------------------------------
# 1) If it's the normal base_id_column (like "parent_id" or "collection_id")
# ----------------------------------------------------------------------
if key == base_id_column:
if isinstance(value, dict):
op, clause = next(iter(value.items()))
if op == "$eq":
# single equality
parameters.append(str(clause))
return f"{base_id_column} = ${len(parameters)}::uuid"
elif op in ("$in", "$overlap"):
# treat both $in/$overlap as "IN the set" for a single column
array_val = [str(x) for x in clause]
parameters.append(array_val)
return f"{base_id_column} = ANY(${len(parameters)}::uuid[])"
# handle other operators as needed
else:
# direct equality
parameters.append(str(value))
return f"{base_id_column} = ${len(parameters)}::uuid"
# ----------------------------------------------------------------------
# 2) SPECIAL: if user specifically sets "collection_ids" in filters
# We interpret that to mean "Look for rows whose parent_id (or collection_id)
# is in the array of values" – i.e. we do the same logic but we forcibly
# direct it to the same column: parent_id or collection_id.
# ----------------------------------------------------------------------
elif key == "collection_ids":
# If we are searching communities, the relevant field is `collection_id`.
# If searching entities/relationships, the relevant field is `parent_id`.
col_to_use = (
"collection_id"
if search_type == "communities"
else "parent_id"
)
if isinstance(value, dict):
op, clause = next(iter(value.items()))
if op == "$eq":
# single equality => col_to_use = clause
parameters.append(str(clause))
return f"{col_to_use} = ${len(parameters)}::uuid"
elif op in ("$in", "$overlap"):
# "col_to_use = ANY($param::uuid[])"
array_val = [str(x) for x in clause]
parameters.append(array_val)
return (
f"{col_to_use} = ANY(${len(parameters)}::uuid[])"
)
# add more if you want, e.g. $ne, $gt, etc.
else:
# direct equality scenario: "collection_ids": "some-uuid"
parameters.append(str(value))
return f"{col_to_use} = ${len(parameters)}::uuid"
# ----------------------------------------------------------------------
# 3) If key starts with "metadata.", handle metadata-based filters
# ----------------------------------------------------------------------
elif key.startswith("metadata."):
field = key.split("metadata.")[1]
if isinstance(value, dict):
op, clause = next(iter(value.items()))
if op == "$eq":
parameters.append(clause)
return f"(metadata->>'{field}') = ${len(parameters)}"
elif op == "$ne":
parameters.append(clause)
return f"(metadata->>'{field}') != ${len(parameters)}"
elif op == "$gt":
parameters.append(clause)
return f"(metadata->>'{field}')::float > ${len(parameters)}::float"
# etc...
else:
parameters.append(value)
return f"(metadata->>'{field}') = ${len(parameters)}"
# ----------------------------------------------------------------------
# 4) Not recognized => return empty so we skip it
# ----------------------------------------------------------------------
return ""
# --------------------------------------------------------------------------
# 5) parse_filter() is the recursive walker that sees $and/$or or normal fields
# --------------------------------------------------------------------------
def parse_filter(fd: dict) -> str:
filter_conditions = []
for k, v in fd.items():
if k == "$and":
and_parts = [parse_filter(sub) for sub in v if sub]
and_parts = [x for x in and_parts if x.strip()]
if and_parts:
filter_conditions.append(
f"({' AND '.join(and_parts)})"
)
elif k == "$or":
or_parts = [parse_filter(sub) for sub in v if sub]
or_parts = [x for x in or_parts if x.strip()]
if or_parts:
filter_conditions.append(f"({' OR '.join(or_parts)})")
else:
c = parse_condition(k, v)
if c and c.strip():
filter_conditions.append(c)
if not filter_conditions:
return ""
if len(filter_conditions) == 1:
return filter_conditions[0]
return " AND ".join(filter_conditions)
return parse_filter(filter_dict)
async def get_existing_document_entity_chunk_ids(
self, document_id: UUID
) -> list[str]:
QUERY = f"""
SELECT DISTINCT unnest(chunk_ids) AS chunk_id FROM {self._get_table_name("documents_entities")} WHERE parent_id = $1
"""
return [
item["chunk_id"]
for item in await self.connection_manager.fetch_query(
QUERY, [document_id]
)
]
async def get_entity_count(
self,
collection_id: Optional[UUID] = None,
document_id: Optional[UUID] = None,
distinct: bool = False,
entity_table_name: str = "entity",
) -> int:
if collection_id is None and document_id is None:
raise ValueError(
"Either collection_id or document_id must be provided."
)
conditions = ["parent_id = $1"]
params = [str(document_id)]
count_value = "DISTINCT name" if distinct else "*"
QUERY = f"""
SELECT COUNT({count_value}) FROM {self._get_table_name(entity_table_name)}
WHERE {" AND ".join(conditions)}
"""
return (await self.connection_manager.fetch_query(QUERY, params))[0][
"count"
]
async def update_entity_descriptions(self, entities: list[Entity]):
query = f"""
UPDATE {self._get_table_name("graphs_entities")}
SET description = $3, description_embedding = $4
WHERE name = $1 AND graph_id = $2
"""
inputs = [
(
entity.name,
entity.parent_id,
entity.description,
entity.description_embedding,
)
for entity in entities
]
await self.connection_manager.execute_many(query, inputs) # type: ignore
def _json_serialize(obj):
if isinstance(obj, UUID):
return str(obj)
elif isinstance(obj, (datetime.datetime, datetime.date)):
return obj.isoformat()
raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
async def _add_objects(
objects: list[dict],
full_table_name: str,
connection_manager: PostgresConnectionManager,
conflict_columns: list[str] | None = None,
exclude_metadata: list[str] | None = None,
) -> list[UUID]:
"""Bulk insert objects into the specified table using
jsonb_to_recordset."""
if conflict_columns is None:
conflict_columns = []
if exclude_metadata is None:
exclude_metadata = []
# Exclude specified metadata and prepare data
cleaned_objects = []
for obj in objects:
cleaned_obj = {
k: v
for k, v in obj.items()
if k not in exclude_metadata and v is not None
}
cleaned_objects.append(cleaned_obj)
# Serialize the list of objects to JSON
json_data = json.dumps(cleaned_objects, default=_json_serialize)
# Prepare the column definitions for jsonb_to_recordset
columns = cleaned_objects[0].keys()
column_defs = []
for col in columns:
# Map Python types to PostgreSQL types
sample_value = cleaned_objects[0][col]
if "embedding" in col:
pg_type = "vector"
elif "chunk_ids" in col or "document_ids" in col or "graph_ids" in col:
pg_type = "uuid[]"
elif col == "id" or "_id" in col:
pg_type = "uuid"
elif isinstance(sample_value, str):
pg_type = "text"
elif isinstance(sample_value, UUID):
pg_type = "uuid"
elif isinstance(sample_value, (int, float)):
pg_type = "numeric"
elif isinstance(sample_value, list) and all(
isinstance(x, UUID) for x in sample_value
):
pg_type = "uuid[]"
elif isinstance(sample_value, list):
pg_type = "jsonb"
elif isinstance(sample_value, dict):
pg_type = "jsonb"
elif isinstance(sample_value, bool):
pg_type = "boolean"
elif isinstance(sample_value, (datetime.datetime, datetime.date)):
pg_type = "timestamp"
else:
raise TypeError(
f"Unsupported data type for column '{col}': {type(sample_value)}"
)
column_defs.append(f"{col} {pg_type}")
columns_str = ", ".join(columns)
column_defs_str = ", ".join(column_defs)
if conflict_columns:
conflict_columns_str = ", ".join(conflict_columns)
update_columns_str = ", ".join(
f"{col}=EXCLUDED.{col}"
for col in columns
if col not in conflict_columns
)
on_conflict_clause = f"ON CONFLICT ({conflict_columns_str}) DO UPDATE SET {update_columns_str}"
else:
on_conflict_clause = ""
QUERY = f"""
INSERT INTO {full_table_name} ({columns_str})
SELECT {columns_str}
FROM jsonb_to_recordset($1::jsonb)
AS x({column_defs_str})
{on_conflict_clause}
RETURNING id;
"""
# Execute the query
result = await connection_manager.fetch_query(QUERY, [json_data])
# Extract and return the IDs
return [record["id"] for record in result]
