From 8f84713b546c8b0f571ef09e0d9eb8554ad660f1 Mon Sep 17 00:00:00 2001 From: Boris Date: Mon, 10 Feb 2025 17:16:13 +0100 Subject: [PATCH] fix: support structured data conversion to data points (#512) ## Description ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin ## Summary by CodeRabbit - New Features - Introduced version tracking and enhanced metadata in core data models for improved data consistency. - Bug Fixes - Improved error handling during graph data loading to prevent disruptions from unexpected identifier formats. - Refactor - Centralized identifier parsing and streamlined model definitions, ensuring smoother and more consistent operations across search, retrieval, and indexing workflows. --- cognee/api/v1/cognify/cognify_v2.py | 4 ---- .../databases/graph/networkx/adapter.py | 12 ++++++++---- .../databases/hybrid/falkordb/FalkorDBAdapter.py | 2 +- .../databases/vector/lancedb/LanceDBAdapter.py | 8 ++++---- .../databases/vector/milvus/MilvusAdapter.py | 4 ++-- .../databases/vector/pgvector/PGVectorAdapter.py | 7 ++++--- .../databases/vector/qdrant/QDrantAdapter.py | 10 +++++----- .../databases/vector/weaviate_db/WeaviateAdapter.py | 6 +++--- cognee/infrastructure/engine/models/DataPoint.py | 1 - cognee/infrastructure/engine/utils/__init__.py | 1 + cognee/infrastructure/engine/utils/parse_id.py | 10 ++++++++++ cognee/modules/chunking/models/DocumentChunk.py | 2 -- cognee/modules/engine/models/Entity.py | 2 -- cognee/modules/engine/models/EntityType.py | 2 -- cognee/modules/graph/cognee_graph/CogneeGraph.py | 2 +- cognee/modules/graph/models/EdgeType.py | 1 - cognee/modules/graph/utils/get_graph_from_model.py | 3 --- cognee/modules/search/methods/search.py | 4 ++-- cognee/shared/CodeGraphEntities.py | 8 -------- cognee/shared/data_models.py | 4 ---- cognee/tasks/storage/add_data_points.py | 4 ++++ cognee/tasks/storage/index_graph_edges.py | 4 ++-- cognee/tasks/summarization/models.py | 3 --- cognee/tasks/temporal_awareness/graphiti_model.py | 2 -- .../temporal_awareness/index_graphiti_objects.py | 8 ++++---- 25 files changed, 51 insertions(+), 63 deletions(-) create mode 100644 cognee/infrastructure/engine/utils/__init__.py create mode 100644 cognee/infrastructure/engine/utils/parse_id.py diff --git a/cognee/api/v1/cognify/cognify_v2.py b/cognee/api/v1/cognify/cognify_v2.py index 9385507ce..dd4ab2229 100644 --- a/cognee/api/v1/cognify/cognify_v2.py +++ b/cognee/api/v1/cognify/cognify_v2.py @@ -25,8 +25,6 @@ from cognee.tasks.documents import ( ) from cognee.tasks.graph import extract_graph_from_data from cognee.tasks.storage import add_data_points -from cognee.modules.data.methods import store_descriptive_metrics -from cognee.tasks.storage.index_graph_edges import index_graph_edges from cognee.tasks.summarization import summarize_text logger = logging.getLogger("cognify.v2") @@ -112,8 +110,6 @@ async def run_cognify_pipeline(dataset: Dataset, user: User, tasks: list[Task]): async for result in pipeline: print(result) - await index_graph_edges() - send_telemetry("cognee.cognify EXECUTION COMPLETED", user.id) await log_pipeline_status( diff --git a/cognee/infrastructure/databases/graph/networkx/adapter.py b/cognee/infrastructure/databases/graph/networkx/adapter.py index 1d0c4c540..708336a5b 100644 --- a/cognee/infrastructure/databases/graph/networkx/adapter.py +++ b/cognee/infrastructure/databases/graph/networkx/adapter.py @@ -5,7 +5,6 @@ import os import json import asyncio import logging -from re import A from typing import Dict, Any, List, Union from uuid import UUID import aiofiles @@ -13,6 +12,7 @@ import aiofiles.os as aiofiles_os import networkx as nx from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface from cognee.infrastructure.engine import DataPoint +from cognee.infrastructure.engine.utils import parse_id from cognee.modules.storage.utils import JSONEncoder import numpy as np @@ -268,7 +268,11 @@ class NetworkXAdapter(GraphDBInterface): for node in graph_data["nodes"]: try: if not isinstance(node["id"], UUID): - node["id"] = UUID(node["id"]) + try: + node["id"] = UUID(node["id"]) + except Exception: + # If conversion fails, keep the original id + pass except Exception as e: logger.error(e) raise e @@ -285,12 +289,12 @@ class NetworkXAdapter(GraphDBInterface): for edge in graph_data["links"]: try: if not isinstance(edge["source"], UUID): - source_id = UUID(edge["source"]) + source_id = parse_id(edge["source"]) else: source_id = edge["source"] if not isinstance(edge["target"], UUID): - target_id = UUID(edge["target"]) + target_id = parse_id(edge["target"]) else: target_id = edge["target"] diff --git a/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py b/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py index 50304cd62..3d9a53548 100644 --- a/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py +++ b/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py @@ -69,7 +69,7 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface): return ",".join([f"{key}:{parse_value(value)}" for key, value in properties.items()]) async def create_data_point_query(self, data_point: DataPoint, vectorized_values: dict): - node_label = type(data_point).__tablename__ + node_label = type(data_point).__name__ property_names = DataPoint.get_embeddable_property_names(data_point) node_properties = await self.stringify_properties( diff --git a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py index d5d2aa865..78ac5f8fa 100644 --- a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py +++ b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py @@ -1,6 +1,5 @@ import asyncio from typing import Generic, List, Optional, TypeVar, get_type_hints -from uuid import UUID import lancedb from lancedb.pydantic import LanceModel, Vector @@ -8,6 +7,7 @@ from pydantic import BaseModel from cognee.exceptions import InvalidValueError from cognee.infrastructure.engine import DataPoint +from cognee.infrastructure.engine.utils import parse_id from cognee.infrastructure.files.storage import LocalStorage from cognee.modules.storage.utils import copy_model, get_own_properties @@ -133,7 +133,7 @@ class LanceDBAdapter(VectorDBInterface): return [ ScoredResult( - id=UUID(result["id"]), + id=parse_id(result["id"]), payload=result["payload"], score=0, ) @@ -162,7 +162,7 @@ class LanceDBAdapter(VectorDBInterface): return [ ScoredResult( - id=UUID(result["id"]), + id=parse_id(result["id"]), payload=result["payload"], score=normalized_values[value_index], ) @@ -195,7 +195,7 @@ class LanceDBAdapter(VectorDBInterface): return [ ScoredResult( - id=UUID(result["id"]), + id=parse_id(result["id"]), payload=result["payload"], score=normalized_values[value_index], ) diff --git a/cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py b/cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py index 246be7e9a..0956a70b5 100644 --- a/cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py +++ b/cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py @@ -3,9 +3,9 @@ from __future__ import annotations import asyncio import logging from typing import List, Optional -from uuid import UUID from cognee.infrastructure.engine import DataPoint +from cognee.infrastructure.engine.utils import parse_id from ..embeddings.EmbeddingEngine import EmbeddingEngine from ..models.ScoredResult import ScoredResult @@ -193,7 +193,7 @@ class MilvusAdapter(VectorDBInterface): return [ ScoredResult( - id=UUID(result["id"]), + id=parse_id(result["id"]), score=result["distance"], payload=result.get("entity", {}), ) diff --git a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py index da8a398a0..3377684b5 100644 --- a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py +++ b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py @@ -9,6 +9,7 @@ from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker from cognee.exceptions import InvalidValueError from cognee.infrastructure.databases.exceptions import EntityNotFoundError from cognee.infrastructure.engine import DataPoint +from cognee.infrastructure.engine.utils import parse_id from ...relational.ModelBase import Base from ...relational.sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter @@ -169,7 +170,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface): results = results.all() return [ - ScoredResult(id=UUID(result.id), payload=result.payload, score=0) + ScoredResult(id=parse_id(result.id), payload=result.payload, score=0) for result in results ] @@ -208,7 +209,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface): # Create and return ScoredResult objects return [ - ScoredResult(id=UUID(str(row.id)), payload=row.payload, score=row.similarity) + ScoredResult(id=parse_id(str(row.id)), payload=row.payload, score=row.similarity) for row in vector_list ] @@ -249,7 +250,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface): for vector in closest_items: vector_list.append( { - "id": UUID(str(vector.id)), + "id": parse_id(str(vector.id)), "payload": vector.payload, "_distance": vector.similarity, } diff --git a/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py b/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py index c3644c98c..bf2cfbcf9 100644 --- a/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py +++ b/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py @@ -1,7 +1,7 @@ import logging from typing import Dict, List, Optional -from uuid import UUID +from cognee.infrastructure.engine.utils import parse_id from qdrant_client import AsyncQdrantClient, models from cognee.exceptions import InvalidValueError @@ -170,10 +170,10 @@ class QDrantAdapter(VectorDBInterface): return [ ScoredResult( - id=UUID(result.id), + id=parse_id(result.id), payload={ **result.payload, - "id": UUID(result.id), + "id": parse_id(result.id), }, score=1 - result.score, ) @@ -209,10 +209,10 @@ class QDrantAdapter(VectorDBInterface): return [ ScoredResult( - id=UUID(result.id), + id=parse_id(result.id), payload={ **result.payload, - "id": UUID(result.id), + "id": parse_id(result.id), }, score=1 - result.score, ) diff --git a/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py b/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py index ac3a85cc7..cf4551dc9 100644 --- a/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py +++ b/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py @@ -1,10 +1,10 @@ import asyncio import logging from typing import List, Optional -from uuid import UUID from cognee.exceptions import InvalidValueError from cognee.infrastructure.engine import DataPoint +from cognee.infrastructure.engine.utils import parse_id from ..embeddings.EmbeddingEngine import EmbeddingEngine from ..models.ScoredResult import ScoredResult @@ -188,7 +188,7 @@ class WeaviateAdapter(VectorDBInterface): return [ ScoredResult( - id=UUID(str(result.uuid)), + id=parse_id(str(result.uuid)), payload=result.properties, score=1 - float(result.metadata.score), ) @@ -221,7 +221,7 @@ class WeaviateAdapter(VectorDBInterface): return [ ScoredResult( - id=UUID(str(result.uuid)), + id=parse_id(str(result.uuid)), payload=result.properties, score=1 - float(result.metadata.score), ) diff --git a/cognee/infrastructure/engine/models/DataPoint.py b/cognee/infrastructure/engine/models/DataPoint.py index 9f123f4c0..9ab8cf5ca 100644 --- a/cognee/infrastructure/engine/models/DataPoint.py +++ b/cognee/infrastructure/engine/models/DataPoint.py @@ -15,7 +15,6 @@ class MetaData(TypedDict): # Updated DataPoint model with versioning and new fields class DataPoint(BaseModel): - __tablename__ = "data_point" id: UUID = Field(default_factory=uuid4) created_at: int = Field( default_factory=lambda: int(datetime.now(timezone.utc).timestamp() * 1000) diff --git a/cognee/infrastructure/engine/utils/__init__.py b/cognee/infrastructure/engine/utils/__init__.py new file mode 100644 index 000000000..bfae36f16 --- /dev/null +++ b/cognee/infrastructure/engine/utils/__init__.py @@ -0,0 +1 @@ +from .parse_id import parse_id diff --git a/cognee/infrastructure/engine/utils/parse_id.py b/cognee/infrastructure/engine/utils/parse_id.py new file mode 100644 index 000000000..de284d287 --- /dev/null +++ b/cognee/infrastructure/engine/utils/parse_id.py @@ -0,0 +1,10 @@ +from uuid import UUID + + +def parse_id(id: any): + if isinstance(id, str): + try: + return UUID(id) + except Exception: + pass + return id diff --git a/cognee/modules/chunking/models/DocumentChunk.py b/cognee/modules/chunking/models/DocumentChunk.py index b7f473d36..1e320959c 100644 --- a/cognee/modules/chunking/models/DocumentChunk.py +++ b/cognee/modules/chunking/models/DocumentChunk.py @@ -6,14 +6,12 @@ from cognee.modules.engine.models import Entity class DocumentChunk(DataPoint): - __tablename__ = "document_chunk" text: str word_count: int token_count: int chunk_index: int cut_type: str is_part_of: Document - pydantic_type: str = "DocumentChunk" contains: List[Entity] = None metadata: dict = {"index_fields": ["text"]} diff --git a/cognee/modules/engine/models/Entity.py b/cognee/modules/engine/models/Entity.py index 4aa771bd4..be228eb06 100644 --- a/cognee/modules/engine/models/Entity.py +++ b/cognee/modules/engine/models/Entity.py @@ -3,10 +3,8 @@ from cognee.modules.engine.models.EntityType import EntityType class Entity(DataPoint): - __tablename__ = "entity" name: str is_a: EntityType description: str - pydantic_type: str = "Entity" metadata: dict = {"index_fields": ["name"]} diff --git a/cognee/modules/engine/models/EntityType.py b/cognee/modules/engine/models/EntityType.py index 1672034b3..df577acbd 100644 --- a/cognee/modules/engine/models/EntityType.py +++ b/cognee/modules/engine/models/EntityType.py @@ -2,9 +2,7 @@ from cognee.infrastructure.engine import DataPoint class EntityType(DataPoint): - __tablename__ = "entity_type" name: str description: str - pydantic_type: str = "EntityType" metadata: dict = {"index_fields": ["name"]} diff --git a/cognee/modules/graph/cognee_graph/CogneeGraph.py b/cognee/modules/graph/cognee_graph/CogneeGraph.py index b53367e8d..2436529e6 100644 --- a/cognee/modules/graph/cognee_graph/CogneeGraph.py +++ b/cognee/modules/graph/cognee_graph/CogneeGraph.py @@ -137,7 +137,7 @@ class CogneeGraph(CogneeAbstractGraph): raise ValueError("Failed to generate query embedding.") edge_distances = await vector_engine.get_distance_from_collection_elements( - "edge_type_relationship_name", query_text=query + "EdgeType_relationship_name", query_text=query ) embedding_map = {result.payload["text"]: result.score for result in edge_distances} diff --git a/cognee/modules/graph/models/EdgeType.py b/cognee/modules/graph/models/EdgeType.py index e42310c3c..9ff48c03b 100644 --- a/cognee/modules/graph/models/EdgeType.py +++ b/cognee/modules/graph/models/EdgeType.py @@ -2,7 +2,6 @@ from cognee.infrastructure.engine import DataPoint class EdgeType(DataPoint): - __tablename__ = "edge_type" relationship_name: str number_of_edges: int diff --git a/cognee/modules/graph/utils/get_graph_from_model.py b/cognee/modules/graph/utils/get_graph_from_model.py index c1b9787fb..5aea2de90 100644 --- a/cognee/modules/graph/utils/get_graph_from_model.py +++ b/cognee/modules/graph/utils/get_graph_from_model.py @@ -61,9 +61,6 @@ async def get_graph_from_model( if include_root and str(data_point.id) not in added_nodes: SimpleDataPointModel = copy_model( type(data_point), - include_fields={ - "__tablename__": (str, data_point.__tablename__), - }, exclude_fields=list(excluded_properties), ) nodes.append(SimpleDataPointModel(**data_point_properties)) diff --git a/cognee/modules/search/methods/search.py b/cognee/modules/search/methods/search.py index ede465d11..23f45278f 100644 --- a/cognee/modules/search/methods/search.py +++ b/cognee/modules/search/methods/search.py @@ -1,8 +1,8 @@ import json -from uuid import UUID from typing import Callable from cognee.exceptions import InvalidValueError +from cognee.infrastructure.engine.utils import parse_id from cognee.modules.retrieval.code_graph_retrieval import code_graph_retrieval from cognee.modules.search.types import SearchType from cognee.modules.storage.utils import JSONEncoder @@ -32,7 +32,7 @@ async def search( for search_result in search_results: document_id = search_result["document_id"] if "document_id" in search_result else None - document_id = UUID(document_id) if isinstance(document_id, str) else document_id + document_id = parse_id(document_id) if document_id is None or document_id in own_document_ids: filtered_search_results.append(search_result) diff --git a/cognee/shared/CodeGraphEntities.py b/cognee/shared/CodeGraphEntities.py index 2d05c56f0..606abc76a 100644 --- a/cognee/shared/CodeGraphEntities.py +++ b/cognee/shared/CodeGraphEntities.py @@ -3,16 +3,12 @@ from cognee.infrastructure.engine import DataPoint class Repository(DataPoint): - __tablename__ = "Repository" path: str - pydantic_type: str = "Repository" metadata: dict = {"index_fields": []} class CodeFile(DataPoint): - __tablename__ = "codefile" extracted_id: str # actually file path - pydantic_type: str = "CodeFile" source_code: Optional[str] = None part_of: Optional[Repository] = None depends_on: Optional[List["CodeFile"]] = None @@ -22,19 +18,15 @@ class CodeFile(DataPoint): class CodePart(DataPoint): - __tablename__ = "codepart" file_path: str # file path # part_of: Optional[CodeFile] = None - pydantic_type: str = "CodePart" source_code: Optional[str] = None metadata: dict = {"index_fields": []} class SourceCodeChunk(DataPoint): - __tablename__ = "sourcecodechunk" code_chunk_of: Optional[CodePart] = None source_code: Optional[str] = None - pydantic_type: str = "SourceCodeChunk" previous_chunk: Optional["SourceCodeChunk"] = None metadata: dict = {"index_fields": ["source_code"]} diff --git a/cognee/shared/data_models.py b/cognee/shared/data_models.py index a36a09010..d23d2841c 100644 --- a/cognee/shared/data_models.py +++ b/cognee/shared/data_models.py @@ -231,7 +231,6 @@ class SummarizedContent(BaseModel): summary: str description: str - pydantic_type: str = "SummarizedContent" class SummarizedFunction(BaseModel): @@ -240,7 +239,6 @@ class SummarizedFunction(BaseModel): inputs: Optional[List[str]] = None outputs: Optional[List[str]] = None decorators: Optional[List[str]] = None - pydantic_type: str = "SummarizedFunction" class SummarizedClass(BaseModel): @@ -248,7 +246,6 @@ class SummarizedClass(BaseModel): description: str methods: Optional[List[SummarizedFunction]] = None decorators: Optional[List[str]] = None - pydantic_type: str = "SummarizedClass" class SummarizedCode(BaseModel): @@ -259,7 +256,6 @@ class SummarizedCode(BaseModel): classes: List[SummarizedClass] = [] functions: List[SummarizedFunction] = [] workflow_description: Optional[str] = None - pydantic_type: str = "SummarizedCode" class GraphDBType(Enum): diff --git a/cognee/tasks/storage/add_data_points.py b/cognee/tasks/storage/add_data_points.py index 540575b5b..c5ab8e702 100644 --- a/cognee/tasks/storage/add_data_points.py +++ b/cognee/tasks/storage/add_data_points.py @@ -3,6 +3,7 @@ from cognee.infrastructure.engine import DataPoint from cognee.infrastructure.databases.graph import get_graph_engine from cognee.modules.graph.utils import deduplicate_nodes_and_edges, get_graph_from_model from .index_data_points import index_data_points +from .index_graph_edges import index_graph_edges async def add_data_points(data_points: list[DataPoint]): @@ -38,4 +39,7 @@ async def add_data_points(data_points: list[DataPoint]): await graph_engine.add_nodes(nodes) await graph_engine.add_edges(edges) + # This step has to happen after adding nodes and edges because we query the graph. + await index_graph_edges() + return data_points diff --git a/cognee/tasks/storage/index_graph_edges.py b/cognee/tasks/storage/index_graph_edges.py index 50d4bb54c..885e4732b 100644 --- a/cognee/tasks/storage/index_graph_edges.py +++ b/cognee/tasks/storage/index_graph_edges.py @@ -51,10 +51,10 @@ async def index_graph_edges(): data_point_type = type(edge) for field_name in edge.metadata["index_fields"]: - index_name = f"{data_point_type.__tablename__}.{field_name}" + index_name = f"{data_point_type.__name__}.{field_name}" if index_name not in created_indexes: - await vector_engine.create_vector_index(data_point_type.__tablename__, field_name) + await vector_engine.create_vector_index(data_point_type.__name__, field_name) created_indexes[index_name] = True if index_name not in index_points: diff --git a/cognee/tasks/summarization/models.py b/cognee/tasks/summarization/models.py index 54932d320..096373c51 100644 --- a/cognee/tasks/summarization/models.py +++ b/cognee/tasks/summarization/models.py @@ -6,7 +6,6 @@ from cognee.shared.CodeGraphEntities import CodeFile, CodePart, SourceCodeChunk class TextSummary(DataPoint): - __tablename__ = "text_summary" text: str made_from: DocumentChunk @@ -14,9 +13,7 @@ class TextSummary(DataPoint): class CodeSummary(DataPoint): - __tablename__ = "code_summary" text: str summarizes: Union[CodeFile, CodePart, SourceCodeChunk] - pydantic_type: str = "CodeSummary" metadata: dict = {"index_fields": ["text"]} diff --git a/cognee/tasks/temporal_awareness/graphiti_model.py b/cognee/tasks/temporal_awareness/graphiti_model.py index 904904091..89aef540b 100644 --- a/cognee/tasks/temporal_awareness/graphiti_model.py +++ b/cognee/tasks/temporal_awareness/graphiti_model.py @@ -3,10 +3,8 @@ from typing import Optional class GraphitiNode(DataPoint): - __tablename__ = "graphitinode" content: Optional[str] = None name: Optional[str] = None summary: Optional[str] = None - pydantic_type: str = "GraphitiNode" metadata: dict = {"index_fields": ["name", "summary", "content"]} diff --git a/cognee/tasks/temporal_awareness/index_graphiti_objects.py b/cognee/tasks/temporal_awareness/index_graphiti_objects.py index 0b0cc10bb..3fb6c9617 100644 --- a/cognee/tasks/temporal_awareness/index_graphiti_objects.py +++ b/cognee/tasks/temporal_awareness/index_graphiti_objects.py @@ -37,10 +37,10 @@ async def index_and_transform_graphiti_nodes_and_edges(): data_point_type = type(graphiti_node) for field_name in graphiti_node.metadata["index_fields"]: - index_name = f"{data_point_type.__tablename__}.{field_name}" + index_name = f"{data_point_type.__name__}.{field_name}" if index_name not in created_indexes: - await vector_engine.create_vector_index(data_point_type.__tablename__, field_name) + await vector_engine.create_vector_index(data_point_type.__name__, field_name) created_indexes[index_name] = True if index_name not in index_points: @@ -66,10 +66,10 @@ async def index_and_transform_graphiti_nodes_and_edges(): data_point_type = type(edge) for field_name in edge.metadata["index_fields"]: - index_name = f"{data_point_type.__tablename__}.{field_name}" + index_name = f"{data_point_type.__name__}.{field_name}" if index_name not in created_indexes: - await vector_engine.create_vector_index(data_point_type.__tablename__, field_name) + await vector_engine.create_vector_index(data_point_type.__name__, field_name) created_indexes[index_name] = True if index_name not in index_points: