KG Response sync

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

@@ -18,34 +18,25 @@
     WrappedMetadataUpdateResponse,
     WrappedUpdateResponse,
 )
+
 from shared.api.models.kg.responses import (
+    Entity,
+    Relationship,
+    Community,
     KGCreationResponse,
     KGEnrichmentResponse,
+    KGTunePromptResponse,
     KGEntityDeduplicationResponse,
-    WrappedKGCommunitiesResponse,
+    WrappedEntityResponse,
+    WrappedEntitiesResponse,
+    WrappedRelationshipResponse,
+    WrappedRelationshipsResponse,
+    WrappedCommunityResponse,
+    WrappedCommunitiesResponse,
     WrappedKGCreationResponse,
     WrappedKGEnrichmentResponse,
-    WrappedKGEntitiesResponse,
-    WrappedKGEntityDeduplicationResponse,
-    WrappedKGRelationshipsResponse,
     WrappedKGTunePromptResponse,
-)
-
-
-from shared.api.models.kg.responses_v3 import (
-    WrappedKGEntitiesResponse as WrappedKGEntitiesResponseV3,
-    WrappedKGRelationshipsResponse as WrappedKGRelationshipsResponseV3,
-    WrappedKGCommunitiesResponse as WrappedKGCommunitiesResponseV3,
-    WrappedKGCreationResponse as WrappedKGCreationResponseV3,
-    WrappedKGEnrichmentResponse as WrappedKGEnrichmentResponseV3,
-    WrappedKGTunePromptResponse as WrappedKGTunePromptResponseV3,
-    WrappedKGEntityDeduplicationResponse as WrappedKGEntityDeduplicationResponseV3,
-    WrappedKGDeletionResponse as WrappedKGDeletionResponseV3,
-    KGCreationResponse as KGCreationResponseV3,
-    KGEnrichmentResponse as KGEnrichmentResponseV3,
-    KGEntityDeduplicationResponse as KGEntityDeduplicationResponseV3,
-    KGTunePromptResponse as KGTunePromptResponseV3,
-    KGDeletionResponse as KGDeletionResponseV3,
+    WrappedKGEntityDeduplicationResponse,
 )
 
 
@@ -112,29 +103,24 @@
     "WrappedMetadataUpdateResponse",
     "WrappedListVectorIndicesResponse",
     "UpdateResponse",
-    # Knowledge Graph Responses for V2
-    # will be removed eventually
+    # Knowledge Graph Responses
+    "Entity",
+    "Relationship",
+    "Community",
     "KGCreationResponse",
-    "WrappedKGCreationResponse",
     "KGEnrichmentResponse",
-    "WrappedKGEnrichmentResponse",
+    "KGTunePromptResponse",
     "KGEntityDeduplicationResponse",
-    "WrappedKGEntityDeduplicationResponse",
+    "WrappedEntityResponse",
+    "WrappedEntitiesResponse",
+    "WrappedRelationshipResponse",
+    "WrappedRelationshipsResponse",
+    "WrappedCommunityResponse",
+    "WrappedCommunitiesResponse",
+    "WrappedKGCreationResponse",
+    "WrappedKGEnrichmentResponse",
     "WrappedKGTunePromptResponse",
-    # Knowledge Graph Responses for V3
-    "WrappedKGEntitiesResponseV3",
-    "WrappedKGRelationshipsResponseV3",
-    "WrappedKGCommunitiesResponseV3",
-    "WrappedKGCreationResponseV3",
-    "WrappedKGEnrichmentResponseV3",
-    "WrappedKGTunePromptResponseV3",
-    "WrappedKGEntityDeduplicationResponseV3",
-    "KGCreationResponseV3",
-    "KGEnrichmentResponseV3",
-    "KGEntityDeduplicationResponseV3",
-    "KGTunePromptResponseV3",
-    "WrappedKGDeletionResponseV3",
-    "KGDeletionResponseV3",
+    "WrappedKGEntityDeduplicationResponse",
     # Management Responses
     "PromptResponse",
     "ServerStats",

py/core/base/providers/database.py

@@ -4,12 +4,10 @@
 from io import BytesIO
 from typing import (
     Any,
-    AsyncGenerator,
     BinaryIO,
     Optional,
     Sequence,
     Tuple,
-    Union,
 )
 from uuid import UUID
 
@@ -20,11 +18,9 @@
     CommunityInfo,
     Entity,
     Graph,
-    KGExtraction,
     Message,
     Relationship,
     VectorEntry,
-    EntityLevel,
 )
 from core.base.abstractions import (
     DocumentResponse,
@@ -38,7 +34,6 @@
     SearchSettings,
     UserStats,
     VectorEntry,
-    VectorQuantizationType,
     VectorSearchResult,
     VectorTableName,
 )
@@ -160,7 +155,7 @@ class DatabaseConnectionManager(ABC):
     def execute_query(
         self,
         query: str,
-        params: Optional[Union[dict[str, Any], Sequence[Any]]] = None,
+        params: Optional[dict[str, Any] | Sequence[Any]] = None,
         isolation_level: Optional[str] = None,
     ):
         pass
@@ -173,15 +168,15 @@ async def execute_many(self, query, params=None, batch_size=1000):
     def fetch_query(
         self,
         query: str,
-        params: Optional[Union[dict[str, Any], Sequence[Any]]] = None,
+        params: Optional[dict[str, Any] | Sequence[Any]] = None,
     ):
         pass
 
     @abstractmethod
     def fetchrow_query(
         self,
         query: str,
-        params: Optional[Union[dict[str, Any], Sequence[Any]]] = None,
+        params: Optional[dict[str, Any] | Sequence[Any]] = None,
     ):
         pass
 
@@ -192,7 +187,9 @@ async def initialize(self, pool: Any):
 
 class Handler(ABC):
     def __init__(
-        self, project_name: str, connection_manager: DatabaseConnectionManager
+        self,
+        project_name: str,
+        connection_manager: DatabaseConnectionManager,
     ):
         self.project_name = project_name
         self.connection_manager = connection_manager
@@ -210,13 +207,15 @@ class DocumentHandler(Handler):
     @abstractmethod
     async def upsert_documents_overview(
         self,
-        documents_overview: Union[DocumentResponse, list[DocumentResponse]],
+        documents_overview: DocumentResponse | list[DocumentResponse],
     ) -> None:
         pass
 
     @abstractmethod
     async def delete_from_documents_overview(
-        self, document_id: UUID, version: Optional[str] = None
+        self,
+        document_id: UUID,
+        version: Optional[str] = None,
     ) -> None:
         pass
 
@@ -233,13 +232,18 @@ async def get_documents_overview(
 
     @abstractmethod
     async def get_workflow_status(
-        self, id: Union[UUID, list[UUID]], status_type: str
+        self,
+        id: UUID | list[UUID],
+        status_type: str,
     ):
         pass
 
     @abstractmethod
     async def set_workflow_status(
-        self, id: Union[UUID, list[UUID]], status_type: str, status: str
+        self,
+        id: UUID | list[UUID],
+        status_type: str,
+        status: str,
     ):
         pass
 
@@ -433,7 +437,7 @@ async def remove_user_from_collection(
     @abstractmethod
     async def get_users_in_collection(
         self, collection_id: UUID, offset: int, limit: int
-    ) -> dict[str, Union[list[UserResponse], int]]:
+    ) -> dict[str, list[UserResponse] | int]:
         pass
 
     @abstractmethod
@@ -523,6 +527,10 @@ async def remove_document_from_collection_vector(
     ) -> None:
         pass
 
+    @abstractmethod
+    async def delete_node_via_document_id(self, document_id: UUID) -> None:
+        pass
+
     @abstractmethod
     async def delete_user_vector(self, user_id: UUID) -> None:
         pass
@@ -551,9 +559,7 @@ async def create_index(
         table_name: Optional[VectorTableName] = None,
         index_measure: IndexMeasure = IndexMeasure.cosine_distance,
         index_method: IndexMethod = IndexMethod.auto,
-        index_arguments: Optional[
-            Union[IndexArgsIVFFlat, IndexArgsHNSW]
-        ] = None,
+        index_arguments: Optional[IndexArgsIVFFlat | IndexArgsHNSW] = None,
         index_name: Optional[str] = None,
         index_column: Optional[str] = None,
         concurrently: bool = True,
@@ -1285,21 +1291,34 @@ async def delete(
         return await self.vector_handler.delete(filters)
 
     async def assign_document_to_collection_vector(
-        self, document_id: UUID, collection_id: UUID
+        self,
+        document_id: UUID,
+        collection_id: UUID,
     ) -> None:
         return await self.vector_handler.assign_document_to_collection_vector(
-            document_id, collection_id
+            document_id=document_id,
+            collection_id=collection_id,
         )
 
     async def remove_document_from_collection_vector(
-        self, document_id: UUID, collection_id: UUID
+        self,
+        document_id: UUID,
+        collection_id: UUID,
     ) -> None:
         return (
             await self.vector_handler.remove_document_from_collection_vector(
                 document_id, collection_id
             )
         )
 
+    async def delete_node_via_document_id(
+        self,
+        document_id: UUID,
+    ) -> None:
+        return await self.vector_handler.delete_node_via_document_id(
+            document_id=document_id,
+        )
+
     async def delete_user_vector(self, user_id: UUID) -> None:
         return await self.vector_handler.delete_user_vector(user_id)
 

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

@@ -9,12 +9,12 @@
 from core.base import Workflow
 from core.base.abstractions import EntityLevel, KGRunType
 from core.base.api.models import (
-    WrappedKGCommunitiesResponse,
+    WrappedCommunitiesResponse,
     WrappedKGCreationResponse,
     WrappedKGEnrichmentResponse,
-    WrappedKGEntitiesResponse,
+    WrappedEntitiesResponse,
     WrappedKGEntityDeduplicationResponse,
-    WrappedKGRelationshipsResponse,
+    WrappedRelationshipsResponse,
     WrappedKGTunePromptResponse,
 )
 
@@ -262,7 +262,7 @@ async def get_entities(
                 description="Number of items to return. Use -1 to return all items.",
             ),
             auth_user=Depends(self.service.providers.auth.auth_wrapper),
-        ) -> WrappedKGEntitiesResponse:
+        ) -> WrappedEntitiesResponse:
             """
             Retrieve entities from the knowledge graph.
             """
@@ -311,7 +311,7 @@ async def get_relationships(
                 description="Number of items to return. Use -1 to return all items.",
             ),
             auth_user=Depends(self.service.providers.auth.auth_wrapper),
-        ) -> WrappedKGRelationshipsResponse:
+        ) -> WrappedRelationshipsResponse:
             """
             Retrieve relationships from the knowledge graph.
             """
@@ -350,7 +350,7 @@ async def get_communities(
                 description="Number of items to return. Use -1 to return all items.",
             ),
             auth_user=Depends(self.service.providers.auth.auth_wrapper),
-        ) -> WrappedKGCommunitiesResponse:
+        ) -> WrappedCommunitiesResponse:
             """
             Retrieve communities from the knowledge graph.
             """

py/core/main/api/v3/chunks_router.py

@@ -6,7 +6,7 @@
 from typing import Any, Optional
 from uuid import UUID
 
-from fastapi import Body, Depends, Path, Query, Form
+from fastapi import Body, Depends, Path, Query
 
 from core.base import (
     KGSearchSettings,
@@ -486,7 +486,7 @@ async def update_chunk(
             )
             await simple_ingestor["update-chunk"](workflow_input)
 
-            return ChunkResponse(
+            return ChunkResponse(  # type: ignore
                 id=chunk_update.id,
                 document_id=existing_chunk["document_id"],
                 user_id=existing_chunk["user_id"],