add collection extraction

py/cli/command_group.py

@@ -1,11 +1,11 @@
+from functools import wraps
+
 import asyncclick as click
 from asyncclick import pass_context
 from asyncclick.exceptions import Exit
 
 from sdk import R2RAsyncClient
 
-from functools import wraps
-
 
 def deprecated_command(new_name):
     def decorator(f):

py/cli/commands/collections.py

@@ -3,8 +3,8 @@
 import asyncclick as click
 from asyncclick import pass_context
 
-from r2r import R2RAsyncClient
 from cli.utils.timer import timer
+from r2r import R2RAsyncClient
 
 
 @click.group()

py/cli/commands/conversations.py

@@ -3,8 +3,8 @@
 import asyncclick as click
 from asyncclick import pass_context
 
-from r2r import R2RAsyncClient
 from cli.utils.timer import timer
+from r2r import R2RAsyncClient
 
 
 @click.group()

py/cli/commands/indices.py

@@ -3,8 +3,8 @@
 import asyncclick as click
 from asyncclick import pass_context
 
-from r2r import R2RAsyncClient
 from cli.utils.timer import timer
+from r2r import R2RAsyncClient
 
 
 @click.group()

py/cli/commands/prompts.py

@@ -3,8 +3,8 @@
 import asyncclick as click
 from asyncclick import pass_context
 
-from r2r import R2RAsyncClient
 from cli.utils.timer import timer
+from r2r import R2RAsyncClient
 
 
 @click.group()

py/cli/commands/system.py

@@ -1,20 +1,14 @@
 import json
-
-import asyncclick as click
-from asyncclick import pass_context
-
-from r2r import R2RAsyncClient
-from cli.utils.timer import timer
-from cli.command_group import cli
-
-
 import os
 import platform
 import subprocess
 import sys
 from importlib.metadata import version as get_version
 
+import asyncclick as click
+from asyncclick import pass_context
 from dotenv import load_dotenv
+
 from cli.command_group import cli
 from cli.utils.docker_utils import (
     bring_down_docker_compose,
@@ -23,6 +17,8 @@
     run_local_serve,
     wait_for_container_health,
 )
+from cli.utils.timer import timer
+from r2r import R2RAsyncClient
 
 
 @click.group()

py/cli/commands/users.py

@@ -3,8 +3,8 @@
 import asyncclick as click
 from asyncclick import pass_context
 
-from r2r import R2RAsyncClient
 from cli.utils.timer import timer
+from r2r import R2RAsyncClient
 
 
 @click.group()

py/cli/commands/v2/kg.py

@@ -3,9 +3,9 @@
 import asyncclick as click
 from asyncclick import pass_context
 
-from r2r import R2RAsyncClient
 from cli.command_group import cli
 from cli.utils.timer import timer
+from r2r import R2RAsyncClient
 
 
 # TODO

py/cli/commands/v2/server.py

@@ -1,9 +1,9 @@
 import asyncclick as click
 from asyncclick import pass_context
 
-from r2r import R2RAsyncClient
 from cli.command_group import cli, deprecated_command
 from cli.utils.timer import timer
+from r2r import R2RAsyncClient
 
 
 @cli.command()

py/cli/main.py

@@ -1,25 +1,20 @@
 from cli.command_group import cli
-
-# TODO: Remove old commands in a later release
-from cli.commands.v2 import (
-    ingestion,
-    kg,
-    management,
-    server,
-    retrieval as v2_retrieval,
-)
-from cli.commands import (
-    database,
-    # V3 methods
+from cli.commands import (  # V3 methods
     collections,
     conversations,
+    database,
     documents,
     indices,
     prompts,
     retrieval,
     system,
     users,
 )
+
+# TODO: Remove old commands in a later release
+from cli.commands.v2 import ingestion, kg, management
+from cli.commands.v2 import retrieval as v2_retrieval
+from cli.commands.v2 import server
 from cli.utils.telemetry import posthog, telemetry
 
 

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

@@ -18,9 +18,10 @@
     WrappedMetadataUpdateResponse,
     WrappedUpdateResponse,
 )
-from shared.api.models.kg.responses import (
+from shared.api.models.kg.responses import (  # TODO: Need to review anything above this
     Community,
     Entity,
+    GraphResponse,
     KGCreationResponse,
     KGEnrichmentResponse,
     KGEntityDeduplicationResponse,
@@ -30,16 +31,14 @@
     WrappedCommunityResponse,
     WrappedEntitiesResponse,
     WrappedEntityResponse,
+    WrappedGraphResponse,
+    WrappedGraphsResponse,
     WrappedKGCreationResponse,
     WrappedKGEnrichmentResponse,
     WrappedKGEntityDeduplicationResponse,
     WrappedKGTunePromptResponse,
     WrappedRelationshipResponse,
     WrappedRelationshipsResponse,
-    # TODO: Need to review anything above this
-    GraphResponse,
-    WrappedGraphResponse,
-    WrappedGraphsResponse,
 )
 from shared.api.models.management.responses import (  # Document Responses; Prompt Responses; Chunk Responses; Conversation Responses; User Responses; TODO: anything below this hasn't been reviewed
     AnalyticsResponse,

py/core/base/utils/__init__.py

@@ -2,6 +2,7 @@
     RecursiveCharacterTextSplitter,
     TextSplitter,
     _decorate_vector_type,
+    _get_str_estimation_output,
     decrement_version,
     format_entity_types,
     format_relations,
@@ -18,7 +19,6 @@
     run_pipeline,
     to_async_generator,
     validate_uuid,
-    _get_str_estimation_output,
 )
 
 __all__ = [

py/core/examples/scripts/test_v3_sdk/test_v3_sdk_graph.py

@@ -1,6 +1,7 @@
-from r2r import R2RClient
-import uuid
 import time
+import uuid
+
+from r2r import R2RClient
 
 # Initialize client
 client = R2RClient("http://localhost:7276", prefix="/v3")

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

@@ -119,7 +119,9 @@ async def create_graph(
                     auth_user.id
                 )
 
-            logger.info(f"Running extract-triples on collection {collection_id}")
+            logger.info(
+                f"Running extract-triples on collection {collection_id}"
+            )
 
             # If no run type is provided, default to estimate
             if not run_type:
@@ -157,7 +159,9 @@ async def create_graph(
                 else:
                     from core.main.orchestration import simple_kg_factory
 
-                    logger.info("Running extract-triples without orchestration.")
+                    logger.info(
+                        "Running extract-triples without orchestration."
+                    )
                     simple_kg = simple_kg_factory(self.service)
                     await simple_kg["extract-triples"](workflow_input)
                     return {
@@ -234,7 +238,9 @@ async def enrich_graph(
                 else:
                     from core.main.orchestration import simple_kg_factory
 
-                    logger.info("Running build-communities without orchestration.")
+                    logger.info(
+                        "Running build-communities without orchestration."
+                    )
                     simple_kg = simple_kg_factory(self.service)
                     await simple_kg["build-communities"](workflow_input)
                     return {

py/core/main/api/v2/management_router.py

@@ -14,7 +14,6 @@
 from core.base.api.models import (
     GenericMessageResponse,
     WrappedAnalyticsResponse,
-    WrappedSettingsResponse,
     WrappedBooleanResponse,
     WrappedChunksResponse,
     WrappedCollectionResponse,
@@ -26,6 +25,7 @@
     WrappedLogsResponse,
     WrappedPromptsResponse,
     WrappedServerStatsResponse,
+    WrappedSettingsResponse,
     WrappedUsersResponse,
 )
 from core.base.logger import AnalysisTypes, LogFilterCriteria

py/core/main/api/v3/collections_router.py

@@ -5,7 +5,7 @@
 
 from fastapi import Body, Depends, Path, Query
 
-from core.base import R2RException, RunType, KGRunType, KGCreationSettings
+from core.base import KGCreationSettings, KGRunType, R2RException, RunType
 from core.base.api.models import (
     GenericBooleanResponse,
     GenericMessageResponse,
@@ -14,8 +14,8 @@
     WrappedCollectionsResponse,
     WrappedDocumentsResponse,
     WrappedGenericMessageResponse,
+    WrappedKGCreationResponse,
     WrappedUsersResponse,
-    WrappedKGCreationResponse
 )
 from core.providers import (
     HatchetOrchestrationProvider,
@@ -1023,7 +1023,6 @@ async def remove_user_from_collection(
                 user_id, id
             )
             return GenericBooleanResponse(success=True)  # type: ignore
-
 
         @self.router.post(
             "/collections/{id}/extract",

py/core/main/api/v3/conversations_router.py

@@ -1,6 +1,5 @@
 import logging
 import textwrap
-
 from typing import Optional
 from uuid import UUID
 

py/core/main/api/v3/documents_router.py

@@ -562,7 +562,6 @@ async def get_documents(
                 False,
                 description="Specifies whether or not to include embeddings of each document summary.",
             ),
-
             auth_user=Depends(self.providers.auth.auth_wrapper),
         ) -> WrappedDocumentsResponse:
             """
@@ -1321,7 +1320,9 @@ async def extract(
                 else:
                     from core.main.orchestration import simple_kg_factory
 
-                    logger.info("Running extract-triples without orchestration.")
+                    logger.info(
+                        "Running extract-triples without orchestration."
+                    )
                     simple_kg = simple_kg_factory(self.services["kg"])
                     await simple_kg["extract-triples"](workflow_input)  # type: ignore
                     return {  # type: ignore

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

@@ -512,7 +512,6 @@ async def build_communities(
                 "task_id": None,
             }
 
-
         @self.router.post(
             "/graphs/{collection_id}/reset",
             summary="Reset a graph back to the initial state.",
@@ -982,7 +981,6 @@ async def delete_relationship(
             )
             return {"success": True}
 
-
         @self.router.post(
             "/graphs/{collection_id}/communities",
             summary="Create a new community",
@@ -1488,4 +1486,3 @@ async def remove_document(
             )
 
             return GenericBooleanResponse(success=success)
-

py/core/main/api/v3/indices_router.py

@@ -4,15 +4,12 @@
 
 import logging
 import textwrap
-
 from typing import Optional
 
 from fastapi import Body, Depends, Path, Query
 
-from core.base import R2RException, RunType, IndexConfig
-from core.base.abstractions import (
-    VectorTableName,
-)
+from core.base import IndexConfig, R2RException, RunType
+from core.base.abstractions import VectorTableName
 from core.base.api.models import (
     GenericMessageResponse,
     WrappedGenericMessageResponse,

py/core/main/api/v3/system_router.py

@@ -1,19 +1,17 @@
-import psutil
 import textwrap
-
 from datetime import datetime, timezone
-
 from typing import Optional
 
+import psutil
 from fastapi import Depends, Query
 
 from core.base import R2RException, RunType
 from core.base.api.models import (
     GenericMessageResponse,
-    WrappedSettingsResponse,
     WrappedGenericMessageResponse,
     WrappedLogsResponse,
     WrappedServerStatsResponse,
+    WrappedSettingsResponse,
 )
 from core.providers import (
     HatchetOrchestrationProvider,