refactor imports

py/cli/commands/ingestion.py

@@ -11,11 +11,7 @@
 from cli.command_group import cli
 from cli.utils.param_types import JSON
 from cli.utils.timer import timer
-from shared.abstractions.vector import (
-    IndexMeasure,
-    IndexMethod,
-    VectorTableName,
-)
+from core.base.abstractions import IndexMeasure, IndexMethod, VectorTableName
 
 
 async def ingest_files_from_urls(client, urls):

py/core/base/abstractions/__init__.py

@@ -22,11 +22,16 @@
     Community,
     CommunityReport,
     Entity,
+    EntityLevel,
     EntityType,
     KGExtraction,
     RelationshipType,
     Triple,
 )
+from shared.abstractions.ingestion import (
+    ChunkEnrichmentSettings,
+    ChunkEnrichmentStrategy,
+)
 from shared.abstractions.kg import (
     KGCreationSettings,
     KGEnrichmentSettings,
@@ -59,10 +64,16 @@
 )
 from shared.abstractions.user import Token, TokenData, UserStats
 from shared.abstractions.vector import (
+    IndexArgsHNSW,
+    IndexArgsIVFFlat,
+    IndexMeasure,
+    IndexMethod,
     StorageResult,
     Vector,
     VectorEntry,
+    VectorQuantizationSettings,
     VectorQuantizationType,
+    VectorTableName,
     VectorType,
 )
 
@@ -97,6 +108,7 @@
     "CommunityReport",
     "KGExtraction",
     "Triple",
+    "EntityLevel",
     # LLM abstractions
     "GenerationConfig",
     "LLMChatCompletion",
@@ -132,6 +144,14 @@
     "Vector",
     "VectorEntry",
     "VectorType",
+    "IndexMeasure",
+    "IndexMethod",
+    "VectorTableName",
+    "IndexArgsHNSW",
+    "IndexArgsIVFFlat",
+    "VectorQuantizationSettings",
     "VectorQuantizationType",
     "StorageResult",
+    "ChunkEnrichmentSettings",
+    "ChunkEnrichmentStrategy",
 ]

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

@@ -17,7 +17,10 @@
     WrappedUpdateResponse,
 )
 from shared.api.models.kg.responses import (
+    KGCreationEstimationResponse,
     KGCreationResponse,
+    KGDeduplicationEstimationResponse,
+    KGEnrichmentEstimationResponse,
     KGEnrichmentResponse,
     KGEntityDeduplicationResponse,
     WrappedKGCommunitiesResponse,
@@ -95,6 +98,9 @@
     "KGEntityDeduplicationResponse",
     "WrappedKGEntityDeduplicationResponse",
     "WrappedKGTunePromptResponse",
+    "KGCreationEstimationResponse",
+    "KGDeduplicationEstimationResponse",
+    "KGEnrichmentEstimationResponse",
     # Management Responses
     "PromptResponse",
     "ServerStats",

py/core/base/providers/database.py

@@ -16,35 +16,29 @@
 )
 from core.base.abstractions import (
     DocumentInfo,
+    IndexArgsHNSW,
+    IndexArgsIVFFlat,
+    IndexMeasure,
+    IndexMethod,
+    KGCreationSettings,
+    KGEnrichmentSettings,
+    KGEntityDeduplicationSettings,
     UserStats,
     VectorEntry,
+    VectorQuantizationType,
     VectorSearchResult,
     VectorSearchSettings,
+    VectorTableName,
 )
 from core.base.api.models import (
     CollectionOverviewResponse,
     CollectionResponse,
-    UserResponse,
-)
-from shared.abstractions import (
-    KGCreationSettings,
-    KGEnrichmentSettings,
-    KGEntityDeduplicationSettings,
-)
-from shared.abstractions.vector import (
-    IndexArgsHNSW,
-    IndexArgsIVFFlat,
-    IndexMeasure,
-    IndexMethod,
-    VectorQuantizationType,
-    VectorTableName,
-)
-from shared.api.models.kg.responses import (
     KGCreationEstimationResponse,
     KGDeduplicationEstimationResponse,
     KGEnrichmentEstimationResponse,
+    UserResponse,
 )
-from shared.utils import _decorate_vector_type
+from core.base.utils import _decorate_vector_type
 
 from .base import Provider, ProviderConfig
 

py/core/base/providers/embedding.py

@@ -7,7 +7,7 @@
 
 from litellm import AuthenticationError
 
-from shared.abstractions.vector import VectorQuantizationSettings
+from core.base.abstractions import VectorQuantizationSettings
 
 from ..abstractions import (
     EmbeddingPurpose,

py/core/base/providers/ingestion.py

@@ -2,7 +2,7 @@
 from abc import ABC
 from enum import Enum
 
-from shared.abstractions.ingestion import ChunkEnrichmentSettings
+from core.base.abstractions import ChunkEnrichmentSettings
 
 from .base import Provider, ProviderConfig
 

py/core/base/utils/__init__.py

@@ -1,6 +1,7 @@
 from shared.utils import (
     RecursiveCharacterTextSplitter,
     TextSplitter,
+    _decorate_vector_type,
     decrement_version,
     format_entity_types,
     format_relations,
@@ -14,6 +15,7 @@
     generate_run_id,
     generate_user_id,
     increment_version,
+    llm_cost_per_million_tokens,
     run_pipeline,
     to_async_generator,
     validate_uuid,
@@ -37,5 +39,7 @@
     "generate_default_prompt_id",
     "RecursiveCharacterTextSplitter",
     "TextSplitter",
+    "llm_cost_per_million_tokens",
     "validate_uuid",
+    "_decorate_vector_type",
 ]

py/core/main/api/ingestion_router.py

@@ -10,6 +10,13 @@
 from pydantic import Json
 
 from core.base import R2RException, RawChunk, generate_document_id
+from core.base.abstractions import (
+    IndexArgsHNSW,
+    IndexArgsIVFFlat,
+    IndexMeasure,
+    IndexMethod,
+    VectorTableName,
+)
 from core.base.api.models import (
     CreateVectorIndexResponse,
     WrappedCreateVectorIndexResponse,
@@ -20,13 +27,6 @@
     WrappedUpdateResponse,
 )
 from core.base.providers import OrchestrationProvider, Workflow
-from shared.abstractions.vector import (
-    IndexArgsHNSW,
-    IndexArgsIVFFlat,
-    IndexMeasure,
-    IndexMethod,
-    VectorTableName,
-)
 
 from ..services.ingestion_service import IngestionService
 from .base_router import BaseRouter, RunType

py/core/main/api/kg_router.py

@@ -7,6 +7,7 @@
 from fastapi import Body, Depends, Query
 
 from core.base import RunType
+from core.base.abstractions import EntityLevel, KGRunType
 from core.base.api.models import (
     WrappedKGCommunitiesResponse,
     WrappedKGCreationResponse,
@@ -17,10 +18,10 @@
     WrappedKGTunePromptResponse,
 )
 from core.base.providers import OrchestrationProvider, Workflow
-from core.utils import generate_default_user_collection_id
-from shared.abstractions.graph import EntityLevel
-from shared.abstractions.kg import KGRunType
-from shared.utils.base_utils import update_settings_from_dict
+from core.utils import (
+    generate_default_user_collection_id,
+    update_settings_from_dict,
+)
 
 from ..services.kg_service import KgService
 from .base_router import BaseRouter

py/core/main/api/management_router.py

@@ -33,7 +33,6 @@
 )
 from core.base.logging import AnalysisTypes, LogFilterCriteria
 from core.base.providers import OrchestrationProvider
-from shared.abstractions.kg import KGRunType
 
 from ..services.management_service import ManagementService
 from .base_router import BaseRouter, RunType

py/core/main/orchestration/hatchet/kg_workflow.py

@@ -9,7 +9,7 @@
 
 from core import GenerationConfig
 from core.base import OrchestrationProvider
-from shared.abstractions.document import KGExtractionStatus
+from core.base.abstractions import KGExtractionStatus
 
 from ...services import KgService
 

py/core/main/services/ingestion_service.py

@@ -21,17 +21,15 @@
     VectorType,
     decrement_version,
 )
-from core.base.api.models import UserResponse
-from core.telemetry.telemetry_decorator import telemetry_event
-from shared.abstractions.ingestion import (
+from core.base.abstractions import (
     ChunkEnrichmentSettings,
     ChunkEnrichmentStrategy,
-)
-from shared.abstractions.vector import (
     IndexMeasure,
     IndexMethod,
     VectorTableName,
 )
+from core.base.api.models import UserResponse
+from core.telemetry.telemetry_decorator import telemetry_event
 
 from ..abstractions import R2RAgents, R2RPipelines, R2RPipes, R2RProviders
 from ..config import R2RConfig

py/core/pipes/kg/community_summary.py

@@ -17,7 +17,7 @@
     PipeType,
     R2RLoggingProvider,
 )
-from shared.abstractions.graph import Entity, Triple
+from core.base.abstractions import Entity, Triple
 
 logger = logging.getLogger()
 

py/core/pipes/kg/deduplication.py

@@ -3,15 +3,14 @@
 from uuid import UUID
 
 from core.base import AsyncState, R2RException
+from core.base.abstractions import Entity, KGEntityDeduplicationType
 from core.base.logging import R2RLoggingProvider
 from core.base.pipes import AsyncPipe, PipeType
 from core.base.providers import (
     CompletionProvider,
     DatabaseProvider,
     EmbeddingProvider,
 )
-from shared.abstractions.graph import Entity
-from shared.abstractions.kg import KGEntityDeduplicationType
 
 logger = logging.getLogger()
 

py/core/pipes/kg/deduplication_summary.py

@@ -4,14 +4,14 @@
 from uuid import UUID
 
 from core.base import AsyncState
+from core.base.abstractions import Entity, GenerationConfig
 from core.base.logging import R2RLoggingProvider
 from core.base.pipes import AsyncPipe, PipeType
 from core.base.providers import (
     CompletionProvider,
     DatabaseProvider,
     EmbeddingProvider,
 )
-from shared.abstractions import Entity, GenerationConfig
 
 logger = logging.getLogger()
 

py/core/providers/database/kg.py

@@ -16,19 +16,19 @@
     R2RException,
     Triple,
 )
-from shared.abstractions import (
+from core.base.abstractions import (
+    EntityLevel,
     KGCreationSettings,
     KGEnrichmentSettings,
     KGEntityDeduplicationSettings,
+    VectorQuantizationType,
 )
-from shared.abstractions.graph import EntityLevel
-from shared.abstractions.vector import VectorQuantizationType
-from shared.api.models.kg.responses import (
+from core.base.api.models import (
     KGCreationEstimationResponse,
     KGDeduplicationEstimationResponse,
     KGEnrichmentEstimationResponse,
 )
-from shared.utils import _decorate_vector_type, llm_cost_per_million_tokens
+from core.base.utils import _decorate_vector_type, llm_cost_per_million_tokens
 
 from .base import PostgresConnectionManager
 from .collection import PostgresCollectionHandler

py/core/providers/database/postgres.py

@@ -21,7 +21,6 @@
 from core.providers.database.tokens import PostgresTokenHandler
 from core.providers.database.user import PostgresUserHandler
 from core.providers.database.vector import PostgresVectorHandler
-from shared.abstractions.vector import VectorQuantizationType
 
 from .base import SemaphoreConnectionPool
 

py/core/providers/database/vector.py

@@ -7,25 +7,23 @@
 from uuid import UUID
 
 from core.base import (
-    VectorEntry,
-    VectorHandler,
-    VectorQuantizationType,
-    VectorSearchResult,
-)
-from core.base.abstractions import VectorSearchSettings
-from shared.abstractions.vector import (
     IndexArgsHNSW,
     IndexArgsIVFFlat,
     IndexMeasure,
     IndexMethod,
+    VectorEntry,
+    VectorHandler,
+    VectorQuantizationType,
+    VectorSearchResult,
+    VectorSearchSettings,
     VectorTableName,
 )
 
 from .base import PostgresConnectionManager
 from .vecs.exc import ArgError, FilterError
 
 logger = logging.getLogger()
-from shared.utils import _decorate_vector_type
+from core.base.utils import _decorate_vector_type
 
 
 def index_measure_to_ops(

py/core/utils/__init__.py

@@ -13,6 +13,7 @@
     increment_version,
     run_pipeline,
     to_async_generator,
+    update_settings_from_dict,
     validate_uuid,
 )
 from shared.utils.splitter.text import (
@@ -36,6 +37,7 @@
     "to_async_generator",
     "generate_default_user_collection_id",
     "validate_uuid",
+    "update_settings_from_dict",
     # Text splitter
     "RecursiveCharacterTextSplitter",
     "TextSplitter",