Skip to content

Commit 076b3bf

Browse files
alekszievralekszievr
committed
Add type to DataPoint metadata
1 parent 3935660 commit 076b3bf

File tree

19 files changed

+154
-58
lines changed

19 files changed

+154
-58
lines changed

‎cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,26 @@
11
import asyncio
22
# from datetime import datetime
33
import json
4-
from uuid import UUID
54
from textwrap import dedent
5+
from uuid import UUID
6+
67
from falkordb import FalkorDB
78

89
from cognee.exceptions import InvalidValueError
9-
from cognee.infrastructure.engine import DataPoint
10-
from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
10+
from cognee.infrastructure.databases.graph.graph_db_interface import \
11+
GraphDBInterface
1112
from cognee.infrastructure.databases.vector.embeddings import EmbeddingEngine
12-
from cognee.infrastructure.databases.vector.vector_db_interface import VectorDBInterface
13+
from cognee.infrastructure.databases.vector.vector_db_interface import \
14+
VectorDBInterface
15+
from cognee.infrastructure.engine import DataPoint
16+
1317

1418
class IndexSchema(DataPoint):
1519
text: str
1620

1721
_metadata: dict = {
18-
"index_fields": ["text"]
22+
"index_fields": ["text"],
23+
"type": "IndexSchema"
1924
}
2025

2126
class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
@@ -333,5 +338,9 @@ async def delete_graph(self):
333338
except Exception as e:
334339
print(f"Error deleting graph: {e}")
335340

341+
async def prune(self):
342+
await self.delete_graph()
343+
print(f"Error deleting graph: {e}")
344+
336345
async def prune(self):
337346
await self.delete_graph()

‎cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,29 @@
1-
from typing import List, Optional, get_type_hints, Generic, TypeVar
21
import asyncio
2+
from typing import Generic, List, Optional, TypeVar, get_type_hints
33
from uuid import UUID
4+
45
import lancedb
6+
from lancedb.pydantic import LanceModel, Vector
57
from pydantic import BaseModel
6-
from lancedb.pydantic import Vector, LanceModel
78

89
from cognee.exceptions import InvalidValueError
910
from cognee.infrastructure.engine import DataPoint
1011
from cognee.infrastructure.files.storage import LocalStorage
1112
from cognee.modules.storage.utils import copy_model, get_own_properties
13+
14+
from ..embeddings.EmbeddingEngine import EmbeddingEngine
1215
from ..models.ScoredResult import ScoredResult
13-
from ..vector_db_interface import VectorDBInterface
1416
from ..utils import normalize_distances
15-
from ..embeddings.EmbeddingEngine import EmbeddingEngine
17+
from ..vector_db_interface import VectorDBInterface
18+
1619

1720
class IndexSchema(DataPoint):
1821
id: str
1922
text: str
2023

2124
_metadata: dict = {
22-
"index_fields": ["text"]
25+
"index_fields": ["text"],
26+
"type": "IndexSchema"
2327
}
2428

2529
class LanceDBAdapter(VectorDBInterface):

‎cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@
44
import logging
55
from typing import List, Optional
66
from uuid import UUID
7+
78
from cognee.infrastructure.engine import DataPoint
8-
from ..vector_db_interface import VectorDBInterface
9-
from ..models.ScoredResult import ScoredResult
9+
1010
from ..embeddings.EmbeddingEngine import EmbeddingEngine
11+
from ..models.ScoredResult import ScoredResult
12+
from ..vector_db_interface import VectorDBInterface
1113

1214
logger = logging.getLogger("MilvusAdapter")
1315

@@ -16,7 +18,8 @@ class IndexSchema(DataPoint):
1618
text: str
1719

1820
_metadata: dict = {
19-
"index_fields": ["text"]
21+
"index_fields": ["text"],
22+
"type": "IndexSchema"
2023
}
2124

2225

‎cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,30 @@
11
import asyncio
2-
from uuid import UUID
32
from typing import List, Optional, get_type_hints
3+
from uuid import UUID
4+
5+
from sqlalchemy import JSON, Column, Table, delete, select
6+
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
47
from sqlalchemy.orm import Mapped, mapped_column
5-
from sqlalchemy import JSON, Column, Table, select, delete
6-
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
78

89
from cognee.exceptions import InvalidValueError
910
from cognee.infrastructure.databases.exceptions import EntityNotFoundError
1011
from cognee.infrastructure.engine import DataPoint
1112

12-
from .serialize_data import serialize_data
13+
from ...relational.ModelBase import Base
14+
from ...relational.sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter
15+
from ..embeddings.EmbeddingEngine import EmbeddingEngine
1316
from ..models.ScoredResult import ScoredResult
14-
from ..vector_db_interface import VectorDBInterface
1517
from ..utils import normalize_distances
16-
from ..embeddings.EmbeddingEngine import EmbeddingEngine
17-
from ...relational.sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter
18-
from ...relational.ModelBase import Base
18+
from ..vector_db_interface import VectorDBInterface
19+
from .serialize_data import serialize_data
20+
1921

2022
class IndexSchema(DataPoint):
2123
text: str
2224

2325
_metadata: dict = {
24-
"index_fields": ["text"]
26+
"index_fields": ["text"],
27+
"type": "IndexSchema"
2528
}
2629

2730
def singleton(class_):

‎cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,25 @@
11
import logging
2+
from typing import Dict, List, Optional
23
from uuid import UUID
3-
from typing import List, Dict, Optional
4+
45
from qdrant_client import AsyncQdrantClient, models
56

67
from cognee.exceptions import InvalidValueError
7-
from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult
8+
from cognee.infrastructure.databases.vector.models.ScoredResult import \
9+
ScoredResult
810
from cognee.infrastructure.engine import DataPoint
9-
from ..vector_db_interface import VectorDBInterface
11+
1012
from ..embeddings.EmbeddingEngine import EmbeddingEngine
13+
from ..vector_db_interface import VectorDBInterface
1114

1215
logger = logging.getLogger("QDrantAdapter")
1316

1417
class IndexSchema(DataPoint):
1518
text: str
1619

1720
_metadata: dict = {
18-
"index_fields": ["text"]
21+
"index_fields": ["text"],
22+
"type": "IndexSchema"
1923
}
2024

2125
# class CollectionConfig(BaseModel, extra = "forbid"):

‎cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,19 @@
55

66
from cognee.exceptions import InvalidValueError
77
from cognee.infrastructure.engine import DataPoint
8-
from ..vector_db_interface import VectorDBInterface
9-
from ..models.ScoredResult import ScoredResult
8+
109
from ..embeddings.EmbeddingEngine import EmbeddingEngine
10+
from ..models.ScoredResult import ScoredResult
11+
from ..vector_db_interface import VectorDBInterface
1112

1213
logger = logging.getLogger("WeaviateAdapter")
1314

1415
class IndexSchema(DataPoint):
1516
text: str
1617

1718
_metadata: dict = {
18-
"index_fields": ["text"]
19+
"index_fields": ["text"],
20+
"type": "IndexSchema"
1921
}
2022

2123
class WeaviateAdapter(VectorDBInterface):
@@ -243,3 +245,4 @@ async def delete_data_points(self, collection_name: str, data_point_ids: list[st
243245

244246
async def prune(self):
245247
self.client.collections.delete_all()
248+
self.client.collections.delete_all()

‎cognee/infrastructure/engine/models/DataPoint.py‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
1-
from typing_extensions import TypedDict
2-
from uuid import UUID, uuid4
3-
from typing import Optional
41
from datetime import datetime, timezone
2+
from typing import Optional
3+
from uuid import UUID, uuid4
4+
55
from pydantic import BaseModel, Field
6+
from typing_extensions import TypedDict
7+
68

79
class MetaData(TypedDict):
810
index_fields: list[str]
@@ -13,7 +15,8 @@ class DataPoint(BaseModel):
1315
updated_at: Optional[datetime] = datetime.now(timezone.utc)
1416
topological_rank: Optional[int] = 0
1517
_metadata: Optional[MetaData] = {
16-
"index_fields": []
18+
"index_fields": [],
19+
"type": "DataPoint"
1720
}
1821

1922
# class Config:
@@ -40,3 +43,5 @@ def get_embeddable_properties(self, data_point):
4043
@classmethod
4144
def get_embeddable_property_names(self, data_point):
4245
return data_point._metadata["index_fields"] or []
46+
def get_embeddable_property_names(self, data_point):
47+
return data_point._metadata["index_fields"] or []

‎cognee/modules/chunking/models/DocumentChunk.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
from typing import List, Optional
2+
23
from cognee.infrastructure.engine import DataPoint
34
from cognee.modules.data.processing.document_types import Document
45
from cognee.modules.engine.models import Entity
56

7+
68
class DocumentChunk(DataPoint):
79
__tablename__ = "document_chunk"
810
text: str
@@ -14,4 +16,5 @@ class DocumentChunk(DataPoint):
1416

1517
_metadata: Optional[dict] = {
1618
"index_fields": ["text"],
19+
"type": "DocumentChunk"
1720
}

‎cognee/modules/data/processing/document_types/Document.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,18 @@
1-
from cognee.infrastructure.engine import DataPoint
21
from uuid import UUID
32

3+
from cognee.infrastructure.engine import DataPoint
4+
5+
46
class Document(DataPoint):
5-
type: str
67
name: str
78
raw_data_location: str
89
metadata_id: UUID
910
mime_type: str
11+
_metadata: dict = {
12+
"type": "Document"
13+
}
1014

1115
def read(self, chunk_size: int) -> str:
1216
pass
17+
def read(self, chunk_size: int) -> str:
18+
pass

‎cognee/modules/engine/models/Entity.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,5 @@ class Entity(DataPoint):
1010

1111
_metadata: dict = {
1212
"index_fields": ["name"],
13+
"type": "Entity"
1314
}

0 commit comments

Comments
 (0)