gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: remove edge_text edges from vector store

Browse source
This commit is contained in:
Boris Arzentar 2025-11-05 10:08:06 +01:00
parent 01a4430be1
commit 81c2797cf1
No known key found for this signature in database
GPG key ID: D5CC274C784807B7
3 changed files with 36 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
cognee/modules/graph/methods/upsert_edges.py
View file

@ -29,29 +29,50 @@ async def upsert_edges(
# Set the session-level RLS variable
await set_current_user(session, user_id)
upsert_statement = (
insert(Edge)
.values(
[
edges_to_add = []
for edge in edges:
edges_to_add.append(
{
"id": uuid5(
NAMESPACE_OID,
str(user_id) + str(dataset_id) + str(edge[0]) + str(edge[2]) + str(edge[1]),
),
"slug": generate_edge_id(edge[2]),
"user_id": user_id,
"data_id": data_id,
"dataset_id": dataset_id,
"source_node_id": edge[0],
"destination_node_id": edge[1],
"relationship_name": edge[2],
"label": edge[2],
"attributes": jsonable_encoder(edge[3]),
}
)
if len(edge) == 4 and "edge_text" in edge[3]:
edge_text = edge[3]["edge_text"]
edges_to_add.append(
{
"id": uuid5(
NAMESPACE_OID,
str(user_id) + str(dataset_id) + str(edge[0]) + str(edge[2]) + str(edge[1]),
str(user_id) + str(dataset_id) + str(edge_text),
),
"slug": generate_edge_id(edge[2]),
"slug": generate_edge_id(edge_text),
"user_id": user_id,
"data_id": data_id,
"dataset_id": dataset_id,
"source_node_id": edge[0],
"destination_node_id": edge[1],
"relationship_name": edge[2],
"label": edge[2],
"props": jsonable_encoder(edge[3]),
"relationship_name": edge_text,
"label": edge_text,
"attributes": jsonable_encoder(edge[3]),
}
for edge in edges
]
)
.on_conflict_do_nothing(index_elements=["id"])
)
upsert_statement = (
insert(Edge).values(edges_to_add).on_conflict_do_nothing(index_elements=["id"])
)
await session.execute(upsert_statement)

2
cognee/modules/graph/models/Edge.py
View file

@ -32,7 +32,7 @@ class Edge(Base):
relationship_name: Mapped[str | None] = mapped_column(String(255), nullable=False)
label: Mapped[str | None] = mapped_column(String(255))
props: Mapped[dict | None] = mapped_column(JSON)
attributes: Mapped[dict | None] = mapped_column(JSON)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False

2
cognee/modules/graph/models/Node.py
View file

@ -31,7 +31,7 @@ class Node(Base):
type: Mapped[str] = mapped_column(String(255), nullable=False)
indexed_fields: Mapped[list] = mapped_column(JSON, nullable=False)
# props: Mapped[dict] = mapped_column(JSON)
attributes: Mapped[dict | None] = mapped_column(JSON)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False