up (#1442)

SciPhi-AI · Oct 22, 2024 · 99b65b9 · 99b65b9
1 parent 44fd5b4
commit 99b65b9
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 5 deletions.
diff --git a/.github/actions/run-sdk-graphrag-tests/action.yml b/.github/actions/run-sdk-graphrag-tests/action.yml
@@ -13,7 +13,7 @@ runs:
       shell: bash
       run: poetry run python tests/integration/runner_sdk.py test_kg_create_graph_sample_file_sdk
 
-    - name: Deduplicate entities (CLI)
+    - name: Deduplicate entities (SDK)
       working-directory: ./py
       shell: bash
       run: poetry run python tests/integration/runner_sdk.py test_kg_deduplicate_entities_sample_file_sdk

diff --git a/py/shared/abstractions/document.py b/py/shared/abstractions/document.py
@@ -128,6 +128,8 @@ class KGExtractionStatus(str, Enum):
     SUCCESS = "success"
     FAILED = "failed"
 
+    def __str__(self):
+        return self.value
 
 class KGEnrichmentStatus(str, Enum):
     """Status of KG Enrichment per collection."""
@@ -137,6 +139,8 @@ class KGEnrichmentStatus(str, Enum):
     SUCCESS = "success"
     FAILED = "failed"
 
+    def __str__(self):
+        return self.value
 
 class DocumentInfo(R2RSerializable):
     """Base class for document information handling."""

diff --git a/py/tests/integration/runner_cli.py b/py/tests/integration/runner_cli.py
@@ -254,7 +254,7 @@ def test_kg_create_graph_sample_file_cli():
     output = run_command("poetry run r2r create-graph --run")
 
     if "queued" in output:
-        time.sleep(120)
+        time.sleep(60)
 
     response = requests.get(
         "http://localhost:7272/v2/entities/",
@@ -291,7 +291,7 @@ def test_kg_deduplicate_entities_sample_file_cli():
     print(output)
 
     if "queued" in output:
-        time.sleep(60)
+        time.sleep(45)
 
     response = requests.get(
         "http://localhost:7272/v2/entities",

diff --git a/py/tests/integration/runner_sdk.py b/py/tests/integration/runner_sdk.py
@@ -689,7 +689,7 @@ def test_kg_create_graph_sample_file_sdk():
     print(create_graph_result)
 
     if "queued" in create_graph_result["results"]["message"]:
-        time.sleep(120)
+        time.sleep(60)
 
     result = client.get_entities(
         collection_id="122fdf6a-e116-546b-a8f6-e4cb2e2c0a09",
@@ -716,7 +716,7 @@ def test_kg_deduplicate_entities_sample_file_sdk():
     )
 
     if "queued" in entities_deduplication_result["results"]["message"]:
-        time.sleep(30)
+        time.sleep(45)
 
     response = client.get_entities(
         collection_id="122fdf6a-e116-546b-a8f6-e4cb2e2c0a09",