Feature/cleanup refactor vector collection

.github/workflows/r2r-light-py-integration-tests.yml

@@ -19,7 +19,7 @@ jobs:
 
     strategy:
       matrix:
-        os: [windows-latest]
+        os: [ubuntu-latest]
         test_category:
           - cli-ingestion
           - cli-retrieval

py/cli/main.py

@@ -1,12 +1,5 @@
 from cli.command_group import cli
-from cli.commands import (
-    auth,
-    ingestion,
-    kg,
-    management,
-    retrieval,
-    server,
-)
+from cli.commands import auth, ingestion, kg, management, retrieval, server
 from cli.utils.telemetry import posthog, telemetry
 
 

py/core/__init__.py

@@ -79,6 +79,7 @@
     "KGSearchSettings",
     "VectorSearchResult",
     "VectorSearchSettings",
+    "HybridSearchSettings",
     # User abstractions
     "Token",
     "TokenData",
@@ -139,8 +140,6 @@
     # Database providers
     "DatabaseConfig",
     "DatabaseProvider",
-    "RelationalDBProvider",
-    "VectorDBProvider",
     # Embedding provider
     "EmbeddingConfig",
     "EmbeddingProvider",

py/core/base/__init__.py

@@ -47,6 +47,7 @@
     "KGSearchSettings",
     "VectorSearchResult",
     "VectorSearchSettings",
+    "HybridSearchSettings",
     # KG abstractions
     "KGCreationSettings",
     "KGEnrichmentSettings",
@@ -112,8 +113,6 @@
     # Database providers
     "DatabaseConfig",
     "DatabaseProvider",
-    "RelationalDBProvider",
-    "VectorDBProvider",
     "PostgresConfigurationSettings",
     # Embedding provider
     "EmbeddingConfig",

py/core/base/abstractions/__init__.py

@@ -1,5 +1,4 @@
 from shared.abstractions.base import AsyncSyncMeta, R2RSerializable, syncable
-from shared.abstractions.llm import MessageType
 from shared.abstractions.document import (
     DataType,
     Document,
@@ -40,6 +39,7 @@
     LLMChatCompletion,
     LLMChatCompletionChunk,
     Message,
+    MessageType,
     RAGCompletion,
 )
 from shared.abstractions.prompt import Prompt

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

@@ -21,29 +21,30 @@
     WrappedKGCreationResponse,
     WrappedKGEnrichmentResponse,
     WrappedKGEntitiesResponse,
-    WrappedKGTriplesResponse,
     WrappedKGEntityDeduplicationResponse,
+    WrappedKGTriplesResponse,
 )
 from shared.api.models.management.responses import (
     AnalyticsResponse,
     AppSettingsResponse,
     CollectionOverviewResponse,
     CollectionResponse,
+    ConversationOverviewResponse,
     DocumentChunkResponse,
     DocumentOverviewResponse,
     LogResponse,
     PromptResponse,
     ScoreCompletionResponse,
     ServerStats,
     UserOverviewResponse,
-    ConversationOverviewResponse,
     WrappedAddUserResponse,
     WrappedAnalyticsResponse,
     WrappedAppSettingsResponse,
     WrappedCollectionListResponse,
     WrappedCollectionOverviewResponse,
     WrappedCollectionResponse,
     WrappedConversationResponse,
+    WrappedConversationsOverviewResponse,
     WrappedDeleteResponse,
     WrappedDocumentChunkResponse,
     WrappedDocumentOverviewResponse,
@@ -54,7 +55,6 @@
     WrappedUserCollectionResponse,
     WrappedUserOverviewResponse,
     WrappedUsersInCollectionResponse,
-    WrappedConversationsOverviewResponse,
 )
 from shared.api.models.retrieval.responses import (
     RAGAgentResponse,

py/core/base/providers/__init__.py

@@ -5,8 +5,7 @@
     DatabaseConfig,
     DatabaseProvider,
     PostgresConfigurationSettings,
-    RelationalDBProvider,
-    VectorDBProvider,
+    VectorQuantizationType,
 )
 from .embedding import EmbeddingConfig, EmbeddingProvider
 from .file import FileConfig, FileProvider
@@ -35,8 +34,7 @@
     "DatabaseConfig",
     "PostgresConfigurationSettings",
     "DatabaseProvider",
-    "RelationalDBProvider",
-    "VectorDBProvider",
+    "VectorQuantizationType",
     # Embedding provider
     "EmbeddingConfig",
     "EmbeddingProvider",

py/core/base/providers/database.py

@@ -1,9 +1,9 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, Optional
+from typing import Any, Optional, Union, Sequence
 
 from pydantic import BaseModel
-
+from uuid import UUID
 from shared.abstractions.vector import VectorQuantizationType
 
 from .base import Provider, ProviderConfig
@@ -53,6 +53,7 @@ class DatabaseConfig(ProviderConfig):
     ] = None
     default_collection_name: str = "Default"
     default_collection_description: str = "Your default collection."
+    enable_fts: bool = False
 
     def __post_init__(self):
         self.validate_config()
@@ -69,41 +70,41 @@ def supported_providers(self) -> list[str]:
         return ["postgres"]
 
 
-class VectorDBProvider(Provider, ABC):
-    @abstractmethod
-    def _initialize_vector_db(
-        self, dimension: int, quantization_type: VectorQuantizationType
-    ) -> None:
-        pass
-
-
-class RelationalDBProvider(Provider, ABC):
-    @abstractmethod
-    async def _initialize_relational_db(self) -> None:
-        pass
-
-
 class DatabaseProvider(Provider):
     def __init__(self, config: DatabaseConfig):
-        if not isinstance(config, DatabaseConfig):
-            raise ValueError(
-                "DatabaseProvider must be initialized with a `DatabaseConfig`."
-            )
         logger.info(f"Initializing DatabaseProvider with config {config}.")
+
         super().__init__(config)
 
-        # remove later to re-introduce typing...
-        self.vector: Any = None
-        self.relational: Any = None
+    @abstractmethod
+    def _get_table_name(self, base_name: str) -> str:
+        pass
 
     @abstractmethod
-    def _initialize_vector_db(self) -> VectorDBProvider:
+    def execute_query(
+        self,
+        query: str,
+        params: Optional[Union[dict[str, Any], Sequence[Any]]] = None,
+        isolation_level: Optional[str] = None,
+    ):
         pass
 
     @abstractmethod
-    async def _initialize_relational_db(self) -> RelationalDBProvider:
+    async def execute_many(self, query, params=None, batch_size=1000):
         pass
 
     @abstractmethod
-    def _get_table_name(self, base_name: str) -> str:
+    def fetch_query(
+        self,
+        query: str,
+        params: Optional[Union[dict[str, Any], Sequence[Any]]] = None,
+    ):
         pass
+
+    @abstractmethod
+    def fetchrow_query(
+        self,
+        query: str,
+        params: Optional[Union[dict[str, Any], Sequence[Any]]] = None,
+    ):
+        pass

py/core/base/providers/ingestion.py

@@ -1,9 +1,11 @@
 import logging
 from abc import ABC
 from enum import Enum
-from .base import Provider, ProviderConfig
+
 from shared.abstractions.ingestion import ChunkEnrichmentSettings
 
+from .base import Provider, ProviderConfig
+
 logger = logging.getLogger()
 
 

py/core/base/providers/kg.py

@@ -10,11 +10,11 @@
     Entity,
     KGCreationSettings,
     KGEnrichmentSettings,
+    KGEntityDeduplicationSettings,
     KGExtraction,
     KGSearchSettings,
     RelationshipType,
     Triple,
-    KGEntityDeduplicationSettings,
 )
 from .base import ProviderConfig
 

py/core/examples/scripts/upload_hf_textbooks_ex.py

@@ -25,7 +25,6 @@ def remove_file(file_path):
 async def process_batch(client, batch):
     results = await client.ingest_files(batch)
     print(f"Submitted {len(results['results'])} files for processing")
-    print("results = ", results["results"])
     # Remove the processed files
     for file_path in batch:
         remove_file(file_path)

py/core/main/api/ingestion_router.py

@@ -336,7 +336,7 @@ async def ingest_chunks_app(
         @self.base_endpoint
         async def create_vector_index_app(
             table_name: Optional[VectorTableName] = Body(
-                default=VectorTableName.CHUNKS,
+                default=VectorTableName.RAW_CHUNKS,
                 description="The name of the vector table to create.",
             ),
             index_method: IndexMethod = Body(
@@ -353,9 +353,9 @@ async def create_vector_index_app(
                 None,
                 description="The arguments for the index method.",
             ),
-            replace: bool = Body(
-                default=True,
-                description="Whether to replace an existing index.",
+            index_name: Optional[str] = Body(
+                None,
+                description="The name of the index to create.",
             ),
             concurrently: bool = Body(
                 default=True,
@@ -365,7 +365,7 @@ async def create_vector_index_app(
         ) -> WrappedCreateVectorIndexResponse:
 
             logger.info(
-                f"Creating vector index for {table_name} with method {index_method}, measure {measure}, replace {replace}, concurrently {concurrently}"
+                f"Creating vector index for {table_name} with method {index_method}, measure {measure}, concurrently {concurrently}"
             )
 
             raw_message = await self.orchestration_provider.run_workflow(
@@ -376,7 +376,6 @@ async def create_vector_index_app(
                         "index_method": index_method,
                         "measure": measure,
                         "index_arguments": index_arguments,
-                        "replace": replace,
                         "concurrently": concurrently,
                     },
                 },

py/core/main/api/kg_router.py

@@ -12,8 +12,8 @@
     WrappedKGCreationResponse,
     WrappedKGEnrichmentResponse,
     WrappedKGEntitiesResponse,
-    WrappedKGTriplesResponse,
     WrappedKGEntityDeduplicationResponse,
+    WrappedKGTriplesResponse,
 )
 from core.base.providers import OrchestrationProvider, Workflow
 from core.utils import generate_default_user_collection_id

py/core/main/api/management_router.py

@@ -10,7 +10,7 @@
 from fastapi.responses import StreamingResponse
 from pydantic import Json
 
-from core.base import R2RException, Message
+from core.base import Message, R2RException
 from core.base.api.models import (
     WrappedAddUserResponse,
     WrappedAnalyticsResponse,
@@ -19,6 +19,7 @@
     WrappedCollectionOverviewResponse,
     WrappedCollectionResponse,
     WrappedConversationResponse,
+    WrappedConversationsOverviewResponse,
     WrappedDeleteResponse,
     WrappedDocumentChunkResponse,
     WrappedDocumentOverviewResponse,
@@ -29,7 +30,6 @@
     WrappedUserCollectionResponse,
     WrappedUserOverviewResponse,
     WrappedUsersInCollectionResponse,
-    WrappedConversationsOverviewResponse,
 )
 from core.base.logging import AnalysisTypes, LogFilterCriteria
 from core.base.providers import OrchestrationProvider

py/core/main/assembly/factory.py

@@ -147,7 +147,7 @@ async def create_database_provider(
                 "Embedding config must have a base dimension to initialize database."
             )
 
-        vector_db_dimension = self.config.embedding.base_dimension
+        dimension = self.config.embedding.base_dimension
         quantization_type = (
             self.config.embedding.quantization_settings.quantization_type
         )
@@ -156,7 +156,7 @@ async def create_database_provider(
 
             database_provider = PostgresDBProvider(
                 db_config,
-                vector_db_dimension,
+                dimension,
                 crypto_provider=crypto_provider,
                 quantization_type=quantization_type,
             )