V3 graph implmentations #1593

shreyaspimpalgaonkar · 2024-11-15T01:52:04Z

Important

Refactor and enhance knowledge graph functionality by renaming triples to relationships, updating models, handlers, and API endpoints, and adjusting tests accordingly.

Behavior:
- Renamed getTriples to getRelationships in r2rV2ClientIntegrationSuperUser.test.ts and r2rV2ClientIntegrationUser.test.ts.
- Updated KGCreationSettings in models.tsx to use kg_relationships_extraction_prompt and max_knowledge_relationships.
- Modified get_triples to get_relationships in kg.py.
- Changed API endpoints in kg_router.py and graph_router.py to handle relationships instead of triples.
Models:
- Replaced Triple with Relationship in __init__.py and other related files.
- Updated CommunityReport to Community in abstractions/__init__.py.
Handlers:
- Introduced PostgresRelationshipHandler and PostgresEntityHandler in kg.py.
- Added methods for creating, updating, and deleting entities and relationships.
Pipes:
- Renamed KGTriplesExtractionPipe to KGRelationshipsExtractionPipe in pipes/__init__.py.
- Updated logic in community_summary.py and deduplication.py to handle relationships.
Prompts:
- Updated graphrag_relationships_extraction_few_shot.yaml to reflect relationship extraction changes.
Tests:
- Adjusted tests in test_kg_community_summary_pipe.py and test_kg_logic.py to align with new relationship handling.

^{This description was created by}^{for 1cb66cf. It will automatically update as commits are pushed.}

…plmentations

NolanTrem

Looks good, just needs a few tweaks. Would like to take another look tomorrow before we merge

NolanTrem · 2024-11-15T04:44:22Z

js/sdk/src/r2rClient.ts

   */
-  @feature("getTriples")
-  async getTriples(
+  @feature("getRelationships")


Let's keep the SDK methods here to getTriples, as not to break any current implementation. We can also add an @deprecated warning to this such that it warns users to switch to the v3 SDK methods

NolanTrem · 2024-11-15T04:45:03Z

py/cli/commands/v2/kg.py

@@ -229,7 +229,7 @@ async def get_entities(
 @click.option(
    "--collection-id",
    required=True,
-    help="Collection ID to retrieve triples from.",
+    help="Collection ID to retrieve relationships from.",


Similar to the SDK, let's leave the CLI as is. I have a deprecation wrapper that we can add to these methods, and then we will point users to the new command.

NolanTrem · 2024-11-15T04:46:41Z

py/core/base/api/models/__init__.py

    WrappedKGTunePromptResponse,
 )
+
+
+from shared.api.models.kg.responses_v3 import (


I think it's fine to get ride of the validation on the v2 endpoints—it's there for legacy, and should work, but, no need to have two versions of this for legacy code.

NolanTrem · 2024-11-15T04:51:20Z

py/core/main/api/v2/kg_router.py

@@ -308,9 +309,9 @@ async def get_triples(
                description="Number of items to return. Use -1 to return all items.",
            ),
            auth_user=Depends(self.service.providers.auth.auth_wrapper),
-        ) -> WrappedKGTriplesResponse:
+        ) -> WrappedKGRelationshipsResponse:


We should modify these to WrappedRelationshipResponse and WrappedRelationshipsResponse to match the rest of the codebase

NolanTrem · 2024-11-15T04:52:32Z

py/core/main/api/v3/graph_router.py

@@ -1,22 +1,33 @@
 import logging
 import textwrap
-from typing import Optional
+from typing import Optional, Union


Nit, but just use |

NolanTrem · 2024-11-15T05:00:52Z

py/core/pipes/retrieval/kg_search_pipe.py

@@ -110,7 +110,7 @@ async def local_search(

            # entity search
            search_type = "__Entity__"
-            async for search_result in await self.database_provider.vector_query(  # type: ignore
+            async for search_result in await self.database_provider.graph_search(  # type: ignore


Nit, but lets stop passing kwargs and move to args whenever possible

NolanTrem · 2024-11-15T05:01:33Z

py/core/providers/database/kg.py

        vector_column_str = _decorate_vector_type(
            f"({self.dimension})", self.quantization_type
        )

        query = f"""
            CREATE TABLE IF NOT EXISTS {self._get_table_name("chunk_entity")} (
-            id SERIAL PRIMARY KEY,
+            id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
+            sid SERIAL NOT NULL,


Let's make sure this gets in the migration script

NolanTrem · 2024-11-15T05:04:05Z

py/core/providers/database/kg.py

+        if entity_level is None:
+            raise ValueError("Entity level is not set")
+
+        for entity in entities:


Let's move this to a generator?

if not all(entity.level == entity_level for entity in entities): raise ValueError("All entities must be of the same level")

NolanTrem · 2024-11-15T05:05:49Z

py/sdk/v2/kg.py

@@ -103,39 +103,39 @@ async def get_entities(

        return await self._make_request("GET", "entities", params=params)  # type: ignore

-    async def get_triples(
+    async def get_relationships(


Let's keep v2 method names, it's fine to change the internal naming conventions, though

NolanTrem · 2024-11-15T05:07:24Z

py/shared/api/models/kg/responses.py

-WrappedKGTriplesResponse = ResultsWrapper[KGTriplesResponse]
+
+# KG Entities
+WrappedKGEntityResponse = ResultsWrapper[KGEntitiesResponse]


Nit, but it should be singular KGEntitiesResponse and so on

…plmentations

ellipsis-dev

❌ Changes requested. Reviewed everything up to 1cb66cf in 2 minutes and 51 seconds

More details

Looked at 10379 lines of code in 57 files
Skipped 0 files when reviewing.
Skipped posting 2 drafted comments based on config settings.

1. py/core/main/api/v2/kg_router.py:312

Draft comment:
The parameter triple_ids should be renamed to relationship_ids to maintain consistency with the rest of the code changes.

            relationship_ids: Optional[list[str]] = Query(

Reason this comment was not posted:
Marked as duplicate.

2. py/tests/core/providers/kg/test_kg_logic.py:87

Draft comment:
Consider renaming triples_raw_list to relationships_raw_list for consistency with the updated terminology.
Reason this comment was not posted:
Confidence changes required: 50%
The PR changes the term 'triples' to 'relationships' across the codebase. However, in the test file, the variable name 'triples_raw_list' should also be updated to 'relationships_raw_list' for consistency.

Workflow ID: wflow_aL6oV2GYDiYT7UXk

Want Ellipsis to fix these issues? Tag @ellipsis-dev in a comment. You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

ellipsis-dev · 2024-11-15T21:05:22Z

py/core/main/api/v2/kg_router.py

@@ -308,9 +311,9 @@ async def get_triples(
                description="Number of items to return. Use -1 to return all items.",
            ),
            auth_user=Depends(self.service.providers.auth.auth_wrapper),


The parameter triple_ids should be renamed to relationship_ids to maintain consistency with the rest of the code changes.

Suggested change

auth_user=Depends(self.service.providers.auth.auth_wrapper),

relationship_ids: Optional[list[str]] = Query(

emrgnt-cmplxty and others added 18 commits November 6, 2024 12:02

complete simple tests, cleanup routers

d7791f2

up

370d947

up

6fdbd6c

checkin

547b9e9

up

a2dc253

Merge remote-tracking branch 'origin/feature/api-v3' into v3-graph-im…

1c2f877

…plmentations

up

e1b5443

response models

db2cd77

checkin

541215b

up

819d6dd

checkin

110e87d

up

257bda8

up

deedf56

up

d6044b5

up

4655043

up

314ad01

Merge remote-tracking branch 'origin/feature/api-v3' into v3-graph-im…

c576ec7

…plmentations

up

623de53

shreyaspimpalgaonkar changed the base branch from feature/api-v3 to main November 15, 2024 04:29

NolanTrem changed the base branch from main to feature/api-v3 November 15, 2024 04:42

NolanTrem requested changes Nov 15, 2024

View reviewed changes

shreyaspimpalgaonkar added 4 commits November 14, 2024 21:13

v2

3289e3d

up

ad1eada

up

68327e6

Merge remote-tracking branch 'origin/feature/api-v3' into v3-graph-im…

62252ed

…plmentations

shreyaspimpalgaonkar changed the base branch from feature/api-v3 to main November 15, 2024 19:30

shreyaspimpalgaonkar added 3 commits November 15, 2024 12:15

up

0918944

up

7636c9d

Merge remote-tracking branch 'origin/feature/api-v3' into v3-graph-im…

1cb66cf

…plmentations

shreyaspimpalgaonkar changed the base branch from main to feature/api-v3 November 15, 2024 21:02

shreyaspimpalgaonkar marked this pull request as ready for review November 15, 2024 21:02

shreyaspimpalgaonkar merged commit e819d2e into feature/api-v3 Nov 15, 2024
16 of 20 checks passed

shreyaspimpalgaonkar deleted the v3-graph-implmentations branch November 15, 2024 21:03

ellipsis-dev bot reviewed Nov 15, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

V3 graph implmentations #1593

V3 graph implmentations #1593

shreyaspimpalgaonkar commented Nov 15, 2024 •

edited by ellipsis-dev bot

Loading

NolanTrem left a comment

NolanTrem Nov 15, 2024

NolanTrem Nov 15, 2024

NolanTrem Nov 15, 2024

NolanTrem Nov 15, 2024

NolanTrem Nov 15, 2024

NolanTrem Nov 15, 2024

NolanTrem Nov 15, 2024

NolanTrem Nov 15, 2024

NolanTrem Nov 15, 2024

NolanTrem Nov 15, 2024

ellipsis-dev bot left a comment

ellipsis-dev bot Nov 15, 2024

	auth_user=Depends(self.service.providers.auth.auth_wrapper),
	relationship_ids: Optional[list[str]] = Query(

V3 graph implmentations #1593

V3 graph implmentations #1593

Conversation

shreyaspimpalgaonkar commented Nov 15, 2024 • edited by ellipsis-dev bot Loading

NolanTrem left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ellipsis-dev bot left a comment

Choose a reason for hiding this comment

ellipsis-dev bot Nov 15, 2024

Choose a reason for hiding this comment

shreyaspimpalgaonkar commented Nov 15, 2024 •

edited by ellipsis-dev bot

Loading