Feature/cleanup refactor vector collection (#1432)

* Feature/include vectors option document chunks (#1419) * Fix async JSON parsing (#1408) * Fix async JSON parsing * Remove score completion from js * clean up js * lockfile * Feature/build custom logger (#1409) * building a custom logger for r2r * fix log * maintain bkwd compat * Feature/add kg description prompt (#1411) * add kg desc prompt * add kg desc prompt * add kg desc prompt * fix prompt name * separate test run freq * task_id check fix * add ingestion docs * updatet * add * rm old prompts * rm old prompots * rm old prompts * rm old prompts * add option to include vectors in document chunks * checkin * update vector --------- Co-authored-by: Nolan Tremelling <[email protected]> * Allow env var to set the default R2R deployment for the dashboard (#1417) * Feature/various documentation tweaks (#1422) * Fix async JSON parsing (#1408) * Fix async JSON parsing * Remove score completion from js * clean up js * lockfile * Feature/build custom logger (#1409) * building a custom logger for r2r * fix log * maintain bkwd compat * Feature/add kg description prompt (#1411) * add kg desc prompt * add kg desc prompt * add kg desc prompt * fix prompt name * separate test run freq * task_id check fix * add ingestion docs * updatet * add * rm old prompts * rm old prompots * rm old prompts * rm old prompts * add option to include vectors in document chunks * checkin * update vector * some various documentation tweaks * some various documentation tweaks --------- Co-authored-by: Nolan Tremelling <[email protected]> * Graphrag tests (#1418) * Fix async JSON parsing (#1408) * Fix async JSON parsing * Remove score completion from js * clean up js * lockfile * Feature/build custom logger (#1409) * building a custom logger for r2r * fix log * maintain bkwd compat * Feature/add kg description prompt (#1411) * add kg desc prompt * add kg desc prompt * add kg desc prompt * fix prompt name * separate test run freq (#1412) * separate test run freq * task_id check fix * add ingestion docs * updatet * add * rm old prompts * rm old prompots * rm old prompts * rm old prompts * Prod fixes + enhancements (#1407) * change default settings back to fp32 * add logging and cache triples * up * up * pre-commit and cleanups * making community summary prompt async * up * up * revert prompt changes * up * up * modify default * bump test timeout due to stricter concurrency limits * bump sleep * rm ubuntu from windows/mac workflows * up * add tests --------- Co-authored-by: Nolan Tremelling <[email protected]> Co-authored-by: emrgnt-cmplxty <[email protected]> Co-authored-by: emrgnt-cmplxty <[email protected]> * Modify graphrag tests timeouts (#1416) * Fix async JSON parsing (#1408) * Fix async JSON parsing * Remove score completion from js * clean up js * lockfile * Feature/build custom logger (#1409) * building a custom logger for r2r * fix log * maintain bkwd compat * Feature/add kg description prompt (#1411) * add kg desc prompt * add kg desc prompt * add kg desc prompt * fix prompt name * separate test run freq (#1412) * separate test run freq * task_id check fix * add ingestion docs * updatet * add * rm old prompts * rm old prompots * rm old prompts * rm old prompts * Prod fixes + enhancements (#1407) * change default settings back to fp32 * add logging and cache triples * up * up * pre-commit and cleanups * making community summary prompt async * up * up * revert prompt changes * up * up * modify default * bump test timeout due to stricter concurrency limits * bump sleep * rm ubuntu from windows/mac workflows * modify timeouts --------- Co-authored-by: Nolan Tremelling <[email protected]> Co-authored-by: emrgnt-cmplxty <[email protected]> Co-authored-by: emrgnt-cmplxty <[email protected]> * feat: Make prompt provider methods asynchronous (comments below) (#1415) * Fix async JSON parsing (#1408) * Fix async JSON parsing * Remove score completion from js * clean up js * lockfile * Feature/build custom logger (#1409) * building a custom logger for r2r * fix log * maintain bkwd compat * Feature/add kg description prompt (#1411) * add kg desc prompt * add kg desc prompt * add kg desc prompt * fix prompt name * separate test run freq (#1412) * separate test run freq * task_id check fix * add ingestion docs * updatet * add * rm old prompts * rm old prompots * rm old prompts * rm old prompts * Prod fixes + enhancements (#1407) * change default settings back to fp32 * add logging and cache triples * up * up * pre-commit and cleanups * making community summary prompt async * up * up * revert prompt changes * up * up * modify default * bump test timeout due to stricter concurrency limits * bump sleep * rm ubuntu from windows/mac workflows * feat: Make prompt provider methods asynchronous --------- Co-authored-by: Nolan Tremelling <[email protected]> Co-authored-by: emrgnt-cmplxty <[email protected]> Co-authored-by: emrgnt-cmplxty <[email protected]> * bump pyproject version * first commit * towards slimmer vector implementation logic * up * iterate * up * checkin * up * work doc chunks * working vector search * working full text search * remove asyncpg * passing vector tests * up * merge * rm pytest * up * up * fix delete * up * up --------- Co-authored-by: Nolan Tremelling <[email protected]> Co-authored-by: Shreyas Pimpalgaonkar <[email protected]>
SciPhi-AI · Oct 21, 2024 · 9e4a1cc · 9e4a1cc
1 parent c9b6549
commit 9e4a1cc
Show file tree

Hide file tree

Showing 68 changed files with 1,335 additions and 2,685 deletions.
diff --git a/.github/workflows/r2r-light-py-integration-tests.yml b/.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

diff --git a/py/cli/main.py b/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
 
 

diff --git a/py/core/__init__.py b/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",

diff --git a/py/core/base/__init__.py b/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",

diff --git a/py/core/base/abstractions/__init__.py b/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

diff --git a/py/core/base/api/models/__init__.py b/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,

diff --git a/py/core/base/providers/__init__.py b/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",

diff --git a/py/core/base/providers/database.py b/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
diff --git a/py/core/base/providers/ingestion.py b/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()
 
 

diff --git a/py/core/base/providers/kg.py b/py/core/base/providers/kg.py
@@ -10,11 +10,11 @@
     Entity,
     KGCreationSettings,
     KGEnrichmentSettings,
+    KGEntityDeduplicationSettings,
     KGExtraction,
     KGSearchSettings,
     RelationshipType,
     Triple,
-    KGEntityDeduplicationSettings,
 )
 from .base import ProviderConfig
 

diff --git a/py/core/examples/scripts/upload_hf_textbooks_ex.py b/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)

diff --git a/py/core/main/api/ingestion_router.py b/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,
                     },
                 },

diff --git a/py/core/main/api/kg_router.py b/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

diff --git a/py/core/main/api/management_router.py b/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

diff --git a/py/core/main/assembly/factory.py b/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,
             )