Fix GraphRAG tests (#1579)

js/sdk/src/r2rClient.ts

@@ -7,6 +7,7 @@ import { ChunksClient } from "./v3/clients/chunks";
 import { CollectionsClient } from "./v3/clients/collections";
 import { ConversationsClient } from "./v3/clients/conversations";
 import { DocumentsClient } from "./v3/clients/documents";
+import { IndiciesClient } from "./v3/clients/indices";
 import { UsersClient } from "./v3/clients/users";
 import { PromptsClient } from "./v3/clients/prompts";
 
@@ -38,6 +39,7 @@ export class r2rClient extends BaseClient {
   public readonly collections: CollectionsClient;
   public readonly conversations: ConversationsClient;
   public readonly documents: DocumentsClient;
+  public readonly indices: IndiciesClient;
   public readonly users: UsersClient;
   public readonly prompts: PromptsClient;
 
@@ -48,6 +50,7 @@ export class r2rClient extends BaseClient {
     this.collections = new CollectionsClient(this);
     this.conversations = new ConversationsClient(this);
     this.documents = new DocumentsClient(this);
+    this.indices = new IndiciesClient(this);
     this.users = new UsersClient(this);
     this.prompts = new PromptsClient(this);
 
@@ -656,6 +659,7 @@ export class r2rClient extends BaseClient {
    * Create a vector index for similarity search.
    * @param options The options for creating the vector index
    * @returns Promise resolving to the creation response
+   * @deprecated Use `client.indices.create` instead.
    */
   @feature("createVectorIndex")
   async createVectorIndex(options: {
@@ -693,6 +697,7 @@ export class r2rClient extends BaseClient {
    * List existing vector indices for a table.
    * @param options The options for listing vector indices
    * @returns Promise resolving to the list of indices
+   * @deprecated Use `client.indices.list` instead.
    */
   @feature("listVectorIndices")
   async listVectorIndices(options: {
@@ -712,6 +717,7 @@ export class r2rClient extends BaseClient {
    * Delete a vector index from a table.
    * @param options The options for deleting the vector index
    * @returns Promise resolving to the deletion response
+   * @deprecated Use `client.indices.delete` instead.
    */
   @feature("deleteVectorIndex")
   async deleteVectorIndex(options: {

js/sdk/src/types.ts

@@ -32,6 +32,18 @@ export type WrappedCollectionsResponse = PaginatedResultsWrapper<
   CollectionResponse[]
 >;
 
+// Index types
+export interface IndexConfig {
+  name?: string;
+  table_name?: string;
+  index_method?: string;
+  index_measure?: string;
+  index_arguments?: string;
+  index_name?: string;
+  index_column?: string;
+  concurrently?: boolean;
+}
+
 // User types
 export interface UserResponse {
   id: string;

js/sdk/src/v3/clients/conversations.ts

@@ -1,4 +1,3 @@
-import { off } from "process";
 import { r2rClient } from "../../r2rClient";
 
 export class ConversationsClient {

js/sdk/src/v3/clients/indices.ts

@@ -0,0 +1,86 @@
+import { r2rClient } from "../../r2rClient";
+import { IndexConfig } from "../../types";
+
+export class IndiciesClient {
+  constructor(private client: r2rClient) {}
+
+  /**
+   * Create a new vector similarity search index in the database.
+   * @param config Configuration for the vector index.
+   * @param run_with_orchestration Whether to run index creation as an orchestrated task.
+   * @returns
+   */
+  async create(options: {
+    config: IndexConfig;
+    run_with_orchestration?: boolean;
+  }): Promise<any> {
+    const data = {
+      config: options.config,
+      ...(options.run_with_orchestration && {
+        run_with_orchestration: options.run_with_orchestration,
+      }),
+    };
+
+    return this.client.makeRequest("POST", `indices`, {
+      data,
+    });
+  }
+
+  /**
+   * List existing vector similarity search indices with pagination support.
+   * @param filters Filter criteria for indices.
+   * @param offset Specifies the number of objects to skip. Defaults to 0.
+   * @param limit Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.
+   * @returns
+   */
+  async list(options?: {
+    filters?: Record<string, any>;
+    offset?: number;
+    limit?: number;
+  }): Promise<any> {
+    const params: Record<string, any> = {
+      offset: options?.offset ?? 0,
+      limit: options?.limit ?? 100,
+    };
+
+    if (options?.filters) {
+      params.filters = options.filters;
+    }
+
+    return this.client.makeRequest("GET", `indices`, {
+      params,
+    });
+  }
+
+  /**
+   * Get detailed information about a specific vector index.
+   * @param index_name The name of the index to retrieve.
+   * @param table_name The name of the table where the index is stored.
+   * @returns
+   */
+  async retrieve(options: {
+    table_name: string;
+    index_name: string;
+  }): Promise<any> {
+    return this.client.makeRequest(
+      "GET",
+      `indices/${options.index_name}/${options.table_name}`,
+    );
+  }
+
+  /**
+   * Delete an existing vector index.
+   * @param index_name The name of the index to delete.
+   * @param table_name The name of the table where the index is stored.
+   * @returns
+   */
+  async delete(options: {
+    table_name: string;
+    index_name: string;
+  }): Promise<any> {
+    return this.client.makeRequest(
+      "DELETE",
+      `indices/${options.index_name}/${options.table_name}`,
+    );
+  }
+}

py/core/__init__.py

@@ -89,6 +89,7 @@
     "Vector",
     "VectorEntry",
     "VectorType",
+    "IndexConfig",
     ## AGENT
     # Agent abstractions
     "Agent",

py/core/base/__init__.py

@@ -62,6 +62,7 @@
     "VectorEntry",
     "VectorType",
     "StorageResult",
+    "IndexConfig",
     ## AGENT
     # Agent abstractions
     "Agent",

py/core/base/abstractions/__init__.py

@@ -70,6 +70,7 @@
 from shared.abstractions.vector import (
     IndexArgsHNSW,
     IndexArgsIVFFlat,
+    IndexConfig,
     IndexMeasure,
     IndexMethod,
     StorageResult,
@@ -115,6 +116,8 @@
     "KGExtraction",
     "Triple",
     "EntityLevel",
+    # Index abstractions
+    "IndexConfig",
     # LLM abstractions
     "GenerationConfig",
     "LLMChatCompletion",

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

@@ -13,15 +13,11 @@
     WrappedUserResponse,
 )
 from shared.api.models.ingestion.responses import (
-    CreateVectorIndexResponse,
     IngestionResponse,
     UpdateResponse,
-    WrappedCreateVectorIndexResponse,
-    WrappedDeleteVectorIndexResponse,
     WrappedIngestionResponse,
     WrappedListVectorIndicesResponse,
     WrappedMetadataUpdateResponse,
-    WrappedSelectVectorIndexResponse,
     WrappedUpdateResponse,
 )
 from shared.api.models.kg.responses import (
@@ -107,11 +103,7 @@
     "WrappedIngestionResponse",
     "WrappedUpdateResponse",
     "WrappedMetadataUpdateResponse",
-    "CreateVectorIndexResponse",
-    "WrappedCreateVectorIndexResponse",
     "WrappedListVectorIndicesResponse",
-    "WrappedDeleteVectorIndexResponse",
-    "WrappedSelectVectorIndexResponse",
     "UpdateResponse",
     # Knowledge Graph Responses
     "KGCreationResponse",

py/core/main/api/v2/ingestion_router.py

@@ -27,8 +27,8 @@
     VectorTableName,
 )
 from core.base.api.models import (
-    WrappedCreateVectorIndexResponse,
-    WrappedDeleteVectorIndexResponse,
+    GenericMessageResponse,
+    WrappedGenericMessageResponse,
     WrappedIngestionResponse,
     WrappedListVectorIndicesResponse,
     WrappedMetadataUpdateResponse,
@@ -610,7 +610,7 @@ async def create_vector_index_app(
                 description=create_vector_descriptions.get("concurrently"),
             ),
             auth_user=Depends(self.service.providers.auth.auth_wrapper),
-        ) -> WrappedCreateVectorIndexResponse:
+        ) -> WrappedGenericMessageResponse:
             """
             Create a vector index for a given table.
 
@@ -638,7 +638,7 @@ async def create_vector_index_app(
                 },
             )
 
-            return raw_message  # type: ignore
+            return GenericMessageResponse(message=raw_message)
 
         list_vector_indices_extras = self.openapi_extras.get(
             "create_vector_index", {}
@@ -708,7 +708,7 @@ async def delete_vector_index_app(
                 ),
             ),
             auth_user=Depends(self.service.providers.auth.auth_wrapper),
-        ) -> WrappedDeleteVectorIndexResponse:
+        ) -> WrappedGenericMessageResponse:
             logger.info(
                 f"Deleting vector index {index_name} from table {table_name}"
             )
@@ -727,7 +727,7 @@ async def delete_vector_index_app(
                 },
             )
 
-            return raw_message  # type: ignore
+            return GenericMessageResponse(message=raw_message)
 
     @staticmethod
     async def _process_files(files):

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

@@ -282,11 +282,11 @@ async def get_entities(
                 entity_table_name = "collection_entity"
 
             return await self.service.get_entities(
-                collection_id,
-                entity_ids,
-                entity_table_name,
-                offset,
-                limit,
+                collection_id=collection_id,
+                entity_ids=entity_ids,
+                entity_table_name=entity_table_name,
+                offset=offset,
+                limit=limit,
             )
 
         @self.router.get("/triples")
@@ -321,11 +321,11 @@ async def get_triples(
                 )
 
             return await self.service.get_triples(
-                collection_id,
-                entity_names,
-                triple_ids,
-                offset,
-                limit,
+                offset=offset,
+                limit=limit,
+                collection_id=collection_id,
+                entity_names=entity_names,
+                triple_ids=triple_ids,
             )
 
         @self.router.get("/communities")
@@ -360,11 +360,11 @@ async def get_communities(
                 )
 
             return await self.service.get_communities(
-                collection_id,
-                levels,
-                community_numbers,
-                offset,
-                limit,
+                offset=offset,
+                limit=limit,
+                collection_id=collection_id,
+                levels=levels,
+                community_numbers=community_numbers,
             )
 
         @self.router.post("/deduplicate_entities")