Feature/revive integration tests (SciPhi-AI#1343) (SciPhi-AI#1346)

* Feature/revive integration tests (SciPhi-AI#1343) * add postgres to integration * add postgres to integration * up * rename * hardcode * add back postgres * add back postgres * add pgvector * add pgvector * add pgvector * add pgvector * add pgvector * add pgvector * add pgvector * tweak config docs * fix integration suite * fix integration suite * fix integration suite * up * up * up * up * up * up * up * up * up * update integration test * final user tests * final user tests * Fix validation error on collection creation responses, remove unnecessary error on deletion (SciPhi-AI#1344) * Don't throw an error when deleting a collection with no documents, fix return on create collection, more js tests * Revert "Don't throw an error when deleting a collection with no documents, fix return on create collection, more js tests" This reverts commit f9f6ead. * Don't throw an error when deleting a collection with no documents, fix return on create collection, more js tests * more * Improve kg throughput (SciPhi-AI#1342) * try * up * up * space * add it in ingestion * rm ingestion * init * add semaphore * test * rm duplicates * kg_creation_settings * rm semaphores * increase conns * change it back * clean * up * up * up * up --------- Co-authored-by: --global=Shreyas Pimpalgaonkar <[email protected]> Co-authored-by: emrgnt-cmplxty <[email protected]> --------- Co-authored-by: Nolan Tremelling <[email protected]> Co-authored-by: Shreyas Pimpalgaonkar <[email protected]> Co-authored-by: --global=Shreyas Pimpalgaonkar <[email protected]>
FutureProofTech · Oct 7, 2024 · 94341a6 · 94341a6
1 parent f0d27a6
commit 94341a6
Show file tree

Hide file tree

Showing 18 changed files with 96 additions and 31 deletions.
diff --git a/js/sdk/__tests__/r2rClientIntegrationSuperUser.test.ts b/js/sdk/__tests__/r2rClientIntegrationSuperUser.test.ts
@@ -3,11 +3,13 @@ const fs = require("fs");
 import { describe, test, beforeAll, expect } from "@jest/globals";
 
 const baseUrl = "http://localhost:7272";
+let newCollectionId: string;
 
 /**
  * raskolnikov.txt should have an id of `f9f61fc8-079c-52d0-910a-c657958e385b`
  * karamozov.txt should have an id of `73749580-1ade-50c6-8fbe-a5e9e87783c8`
  * myshkin.txt should have an id of `2e05b285-2746-5778-9e4a-e293db92f3be`
+ * The default collection should have an id of `122fdf6a-e116-546b-a8f6-e4cb2e2c0a09`
  */
 
 /**
@@ -18,7 +20,7 @@ const baseUrl = "http://localhost:7272";
  *     X verifyEmail
  *     - login
  *     - logout
- *     X user
+ *     - user
  *     X updateUser
  *     - refreshAccessToken
  *     X changePassword
@@ -42,14 +44,14 @@ const baseUrl = "http://localhost:7272";
  *     - documentChunks
  *     X inspectKnowledgeGraph
  *     X collectionsOverview
- *     X createCollection
- *     X getCollection
- *     X updateCollection
- *     X deleteCollection
- *     X listCollections
+ *     - createCollection
+ *     - getCollection
+ *     - updateCollection
+ *     - deleteCollection
+ *     - listCollections
  *     X addUserToCollection
  *     X removeUserFromCollection
- *     X getUsersInCollection
+ *     - getUsersInCollection
  *     X getCollectionsForUser
  *     X assignDocumentToCollection
  *     X removeDocumentFromCollection
@@ -82,6 +84,10 @@ describe("r2rClient Integration Tests", () => {
     ).resolves.not.toThrow();
   });
 
+  test("User", async () => {
+    await expect(client.user()).resolves.not.toThrow();
+  });
+
   test("Server stats", async () => {
     await expect(client.serverStats()).resolves.not.toThrow();
   });
@@ -222,6 +228,53 @@ describe("r2rClient Integration Tests", () => {
     ).resolves.not.toThrow();
   });
 
+  test("Collections overview", async () => {
+    await expect(client.collectionsOverview()).resolves.not.toThrow();
+  });
+
+  test("Create collection", async () => {
+    const response = await client.createCollection("test_collection", "test_description");
+    newCollectionId = response.results.collection_id;
+
+    expect(newCollectionId).toBeDefined();
+  });
+
+  test("Get default collection", async () => {
+    await expect(client.getCollection("122fdf6a-e116-546b-a8f6-e4cb2e2c0a09")).resolves.not.toThrow();
+  });
+
+  test("Get newly created collection", async () => {
+    await expect(client.getCollection(newCollectionId)).resolves.not.toThrow();
+  });
+
+  test("Update collection", async () => {
+      await expect(
+        client.updateCollection(
+          newCollectionId,
+          "updated_test_collection",
+          "updated_test_description"
+        ),
+      ).resolves.not.toThrow();
+    });
+
+    test("List collections", async () => {
+      await expect(client.listCollections()).resolves.not.toThrow();
+    });
+
+    test("Delete collection", async () => {
+      await expect(
+        client.deleteCollection(newCollectionId),
+      ).resolves.not.toThrow();
+    });
+
+    test("Get users in collection", async () => {
+      await expect(client.getUsersInCollection("122fdf6a-e116-546b-a8f6-e4cb2e2c0a09")).resolves.not.toThrow();
+    });
+
+    test("Get users in collection with pagination", async () => {
+      await expect(client.getUsersInCollection("122fdf6a-e116-546b-a8f6-e4cb2e2c0a09", 10, 10)).resolves.not.toThrow();
+    });
+
   test("Clean up remaining documents", async () => {
     // Deletes karamozov.txt
     await expect(

diff --git a/py/core/base/providers/embedding.py b/py/core/base/providers/embedding.py
@@ -25,10 +25,10 @@ class EmbeddingConfig(ProviderConfig):
     batch_size: int = 1
     prefixes: Optional[dict[str, str]] = None
     add_title_as_prefix: bool = True
-    concurrent_request_limit: int = 16
-    max_retries: int = 2
-    initial_backoff: float = 1.0
-    max_backoff: float = 60.0
+    concurrent_request_limit: int = 256
+    max_retries: int = 8
+    initial_backoff: float = 1
+    max_backoff: float = 64.0
 
     def validate_config(self) -> None:
         if self.provider not in self.supported_providers:
@@ -63,6 +63,7 @@ async def _execute_with_backoff_async(self, task: dict[str, Any]):
             try:
                 async with self.semaphore:
                     return await self._execute_task(task)
+            # TODO: Capture different error types and handle them accordingly
             except Exception as e:
                 logger.warning(
                     f"Request failed (attempt {retries + 1}): {str(e)}"

diff --git a/py/core/base/providers/llm.py b/py/core/base/providers/llm.py
@@ -20,9 +20,9 @@ class CompletionConfig(ProviderConfig):
     provider: Optional[str] = None
     generation_config: GenerationConfig = GenerationConfig()
     concurrent_request_limit: int = 256
-    max_retries: int = 2
+    max_retries: int = 8
     initial_backoff: float = 1.0
-    max_backoff: float = 60.0
+    max_backoff: float = 64.0
 
     def validate_config(self) -> None:
         if not self.provider:

diff --git a/py/core/configs/r2r_aws_bedrock.toml b/py/core/configs/r2r_aws_bedrock.toml
@@ -18,7 +18,7 @@ overlap = 20
 
 [completion]
 provider = "litellm"
-concurrent_request_limit = 16
+concurrent_request_limit = 256
 
   [completion.generation_config]
   model = "bedrock/anthropic.claude-v2"

diff --git a/py/core/main/api/management_router.py b/py/core/main/api/management_router.py
@@ -436,9 +436,10 @@ async def create_collection_app(
             collection_id = await self.service.create_collection(
                 name, description
             )
-            return await self.service.add_user_to_collection(  # type: ignore
+            await self.service.add_user_to_collection(  # type: ignore
                 auth_user.id, collection_id.collection_id
             )
+            return collection_id
 
         @self.router.get("/get_collection/{collection_id}")
         @self.base_endpoint

diff --git a/py/core/pipes/kg/entity_description.py b/py/core/pipes/kg/entity_description.py
@@ -134,7 +134,7 @@ async def process_entity(
                                 ),
                             }
                         ],
-                        generation_config=self.kg_provider.config.kg_enrichment_settings.generation_config,
+                        generation_config=self.kg_provider.config.kg_creation_settings.generation_config,
                     )
                 )
                 .choices[0]

diff --git a/py/core/providers/database/relational.py b/py/core/providers/database/relational.py
@@ -1,6 +1,6 @@
 import logging
 from contextlib import asynccontextmanager
-
+import asyncio
 import asyncpg
 
 from core.base import RelationalDBProvider
@@ -35,13 +35,17 @@ def __init__(
         self.project_name = project_name
         self.pool = None
         self.postgres_configuration_settings = postgres_configuration_settings
+        self.semaphore = asyncio.Semaphore(
+            int(self.postgres_configuration_settings.max_connections * 0.9)
+        )
 
     async def initialize(self):
         try:
             self.pool = await asyncpg.create_pool(
                 self.connection_string,
                 max_size=self.postgres_configuration_settings.max_connections,
             )
+
             logger.info(
                 "Successfully connected to Postgres database and created connection pool."
             )
@@ -57,8 +61,9 @@ def _get_table_name(self, base_name: str) -> str:
 
     @asynccontextmanager
     async def get_connection(self):
-        async with self.pool.acquire() as conn:
-            yield conn
+        async with self.semaphore:
+            async with self.pool.acquire() as conn:
+                yield conn
 
     async def execute_query(self, query, params=None):
         async with self.get_connection() as conn:

diff --git a/py/core/providers/database/vector.py b/py/core/providers/database/vector.py
@@ -459,7 +459,7 @@ def delete_collection(self, collection_id: str) -> None:
                 sess.commit()
 
                 if affected_rows == 0:
-                    raise ValueError(
+                    logger.warning(
                         f"Collection {collection_id} not found in any documents."
                     )
             except NoResultFound:

diff --git a/py/core/providers/kg/postgres.py b/py/core/providers/kg/postgres.py
@@ -40,6 +40,7 @@ def __init__(
 
         self.db_provider = db_provider.relational
         self.embedding_provider = embedding_provider
+
         try:
             import networkx as nx
 
@@ -160,9 +161,6 @@ async def create_tables(self, project_name: str):
 
         await self.execute_query(query)
 
-        # TODO: Create another table for entity_embedding_collection
-        # entity embeddings at a collection level
-
         # communities table, result of the Leiden algorithm
         query = f"""
             CREATE TABLE IF NOT EXISTS {self._get_table_name("community")} (

diff --git a/py/core/providers/orchestration/hatchet.py b/py/core/providers/orchestration/hatchet.py
@@ -35,6 +35,7 @@ def get_worker(self, name: str, max_threads: Optional[int] = None) -> Any:
         self.worker = self.orchestrator.worker(name, max_threads)
         return self.worker
 
+
     def concurrency(self, *args, **kwargs) -> Callable:
         return self.orchestrator.concurrency(*args, **kwargs)
 

diff --git a/py/poetry.lock b/py/poetry.lock
diff --git a/py/pyproject.toml b/py/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "r2r"
 readme = "README.md"
-version = "3.2.5"
+version = "3.2.6"
 
 description = "SciPhi R2R"
 authors = ["Owen Colegrove <[email protected]>"]

diff --git a/py/r2r.toml b/py/r2r.toml
@@ -18,7 +18,7 @@ default_admin_password = "Sk7!PhnVrUC2zQgAvqeDBN"
 
 [completion]
 provider = "litellm"
-concurrent_request_limit = 16
+concurrent_request_limit = 256
 
   [completion.generation_config]
   model = "openai/gpt-4o"

diff --git a/py/tests/regression/observed_outputs/test_document_management.json b/py/tests/regression/observed_outputs/test_document_management.json
@@ -1111,4 +1111,4 @@
   "rerun_document_chunks_test": {
     "results": "{\"detail\":{\"message\":\"No chunks found for the given document ID.\",\"error_type\":\"R2RException\"}}"
   }
-}
+}
diff --git a/py/tests/regression/observed_outputs/test_group_management.json b/py/tests/regression/observed_outputs/test_group_management.json
@@ -17,4 +17,4 @@
   "cleanup": {
     "error": "'R2RClient' object has no attribute 'delete_group'"
   }
-}
+}
diff --git a/py/tests/regression/observed_outputs/test_observability.json b/py/tests/regression/observed_outputs/test_observability.json
@@ -177,4 +177,4 @@
       }
     }
   }
-}
+}
diff --git a/py/tests/regression/observed_outputs/test_retrieval.json b/py/tests/regression/observed_outputs/test_retrieval.json
@@ -764,4 +764,4 @@
       }
     }
   }
-}
+}
diff --git a/py/tests/regression/observed_outputs/test_user_management.json b/py/tests/regression/observed_outputs/test_user_management.json
@@ -128,4 +128,4 @@
       "message": "User account f7495bfb-58e1-539e-8b8d-6cb61f821c2b deleted successfully."
     }
   }
-}
+}
-Original file line number
+Diff line change
@@ Expand Up / @@ -134,7 +134,7 @@ async def process_entity( @@
                                     ),
                                 }
                             ],
-                            generation_config=self.kg_provider.config.kg_enrichment_settings.generation_config,
+                            generation_config=self.kg_provider.config.kg_creation_settings.generation_config,
                         )
                     )
                     .choices[0]
@@ Expand Down @@