Feat: Added Contextual AI's Generative and Reranker Client

Author: Jinash Rouniyar

src/collections/config/integration.test.ts

@@ -10,6 +10,7 @@ import {
   PropertyConfig,
   RQConfig,
   RerankerCohereConfig,
+  RerankerContextualAIConfig,
   VectorIndexConfigDynamic,
   VectorIndexConfigHNSW,
 } from './types/index.js';
@@ -787,6 +788,44 @@ describe('Testing of the collection.config namespace', () => {
         model: 'model',
       },
     });
+
+    await collection.config.update({
+      reranker: weaviate.reconfigure.reranker.contextualai({
+        model: 'ctxl-rerank-v2-instruct-multilingual',
+      }),
+    });
+
+    config = await collection.config.get();
+    expect(config.reranker).toEqual<ModuleConfig<'reranker-contextualai', RerankerContextualAIConfig>>({
+      name: 'reranker-contextualai',
+      config: {
+        model: 'ctxl-rerank-v2-instruct-multilingual',
+      },
+    });
+
+    await collection.config.update({
+      generative: weaviate.reconfigure.generative.contextualai({
+        model: 'v2',
+        maxTokens: 100,
+        temperature: 0.7,
+        topP: 0.9,
+        systemPrompt: 'sys',
+        avoidCommentary: false,
+      }),
+    });
+
+    config = await collection.config.get();
+    expect(config.generative).toEqual<ModuleConfig<'generative-contextualai', any>>({
+      name: 'generative-contextualai',
+      config: {
+        model: 'v2',
+        maxTokensProperty: 100,
+        temperatureProperty: 0.7,
+        topPProperty: 0.9,
+        systemPromptProperty: 'sys',
+        avoidCommentaryProperty: false,
+      },
+    });
   });
 
   requireAtLeast(1, 31, 0).it(

src/collections/config/types/generative.ts

@@ -106,12 +106,22 @@ export type GenerativeXAIConfig = {
   topP?: number;
 };
 
+export type GenerativeContextualAIConfig = {
+  model?: string;
+  maxTokensProperty?: number;
+  temperatureProperty?: number;
+  topPProperty?: number;
+  systemPromptProperty?: string;
+  avoidCommentaryProperty?: boolean;
+};
+
 export type GenerativeConfig =
   | GenerativeAnthropicConfig
   | GenerativeAnyscaleConfig
   | GenerativeAWSConfig
   | GenerativeAzureOpenAIConfig
   | GenerativeCohereConfig
+  | GenerativeContextualAIConfig
   | GenerativeDatabricksConfig
   | GenerativeGoogleConfig
   | GenerativeFriendliAIConfig
@@ -133,6 +143,8 @@ export type GenerativeConfigType<G> = G extends 'generative-anthropic'
   ? GenerativeAzureOpenAIConfig
   : G extends 'generative-cohere'
   ? GenerativeCohereConfig
+  : G extends 'generative-contextualai'
+  ? GenerativeContextualAIConfig
   : G extends 'generative-databricks'
   ? GenerativeDatabricksConfig
   : G extends 'generative-google'
@@ -162,6 +174,7 @@ export type GenerativeSearch =
   | 'generative-aws'
   | 'generative-azure-openai'
   | 'generative-cohere'
+  | 'generative-contextualai'
   | 'generative-databricks'
   | 'generative-google'
   | 'generative-friendliai'

src/collections/config/types/reranker.ts

@@ -24,8 +24,20 @@ export type RerankerNvidiaConfig = {
   model?: 'nvidia/rerank-qa-mistral-4b' | string;
 };
 
+export type RerankerContextualAIConfig = {
+  baseURL?: string;
+  model?:
+    | 'ctxl-rerank-v2-instruct-multilingual'
+    | 'ctxl-rerank-v2-instruct-multilingual-mini'
+    | 'ctxl-rerank-v1-instruct'
+    | string;
+  instruction?: string;
+  topN?: number;
+};
+
 export type RerankerConfig =
   | RerankerCohereConfig
+  | RerankerContextualAIConfig
   | RerankerJinaAIConfig
   | RerankerNvidiaConfig
   | RerankerTransformersConfig
@@ -35,6 +47,7 @@ export type RerankerConfig =
 
 export type Reranker =
   | 'reranker-cohere'
+  | 'reranker-contextualai'
   | 'reranker-jinaai'
   | 'reranker-nvidia'
   | 'reranker-transformers'
@@ -48,6 +61,8 @@ export type RerankerConfigType<R> = R extends 'reranker-cohere'
   ? RerankerJinaAIConfig
   : R extends 'reranker-nvidia'
   ? RerankerNvidiaConfig
+  : R extends 'reranker-contextualai'
+  ? RerankerContextualAIConfig
   : R extends 'reranker-transformers'
   ? RerankerTransformersConfig
   : R extends 'reranker-voyageai'

src/collections/configure/generative.ts

@@ -4,6 +4,7 @@ import {
   GenerativeAnyscaleConfig,
   GenerativeAzureOpenAIConfig,
   GenerativeCohereConfig,
+  GenerativeContextualAIConfig,
   GenerativeDatabricksConfig,
   GenerativeFriendliAIConfig,
   GenerativeGoogleConfig,
@@ -21,6 +22,7 @@ import {
   GenerativeAnyscaleConfigCreate,
   GenerativeAzureOpenAIConfigCreate,
   GenerativeCohereConfigCreate,
+  GenerativeContextualAIConfigCreate,
   GenerativeDatabricksConfigCreate,
   GenerativeFriendliAIConfigCreate,
   GenerativeMistralConfigCreate,
@@ -48,6 +50,31 @@ export default {
       config,
     };
   },
+  /**
+   * Create a `ModuleConfig<'generative-contextualai', GenerativeContextualAIConfig | undefined>` object for use when performing AI generation using the `generative-contextualai` module.
+   *
+   * See the [documentation](https://weaviate.io/developers/weaviate/model-providers/contextualai/generative) for detailed usage.
+   *
+   * @param {GenerativeContextualAIConfigCreate} [config] The configuration for the `generative-contextualai` module.
+   * @returns {ModuleConfig<'generative-contextualai', GenerativeContextualAIConfig | undefined>} The configuration object.
+   */
+  contextualai: (
+    config?: GenerativeContextualAIConfigCreate
+  ): ModuleConfig<'generative-contextualai', GenerativeContextualAIConfig | undefined> => {
+    return {
+      name: 'generative-contextualai',
+      config: config
+        ? {
+            model: config.model,
+            maxTokensProperty: config.maxTokens,
+            temperatureProperty: config.temperature,
+            topPProperty: config.topP,
+            systemPromptProperty: config.systemPrompt,
+            avoidCommentaryProperty: config.avoidCommentary,
+          }
+        : undefined,
+    };
+  },
   /**
    * Create a `ModuleConfig<'generative-anyscale', GenerativeAnyscaleConfig | undefined>` object for use when performing AI generation using the `generative-anyscale` module.
    *

src/collections/configure/reranker.ts

@@ -1,6 +1,7 @@
 import {
   ModuleConfig,
   RerankerCohereConfig,
+  RerankerContextualAIConfig,
   RerankerJinaAIConfig,
   RerankerNvidiaConfig,
   RerankerVoyageAIConfig,
@@ -23,6 +24,22 @@ export default {
       config: config,
     };
   },
+  /**
+   * Create a `ModuleConfig<'reranker-contextualai', RerankerContextualAIConfig>` object for use when reranking using the `reranker-contextualai` module.
+   *
+   * See the [documentation](https://weaviate.io/developers/weaviate/model-providers/contextualai/reranker) for detailed usage.
+   *
+   * @param {RerankerContextualAIConfig} [config] The configuration for the `reranker-contextualai` module.
+   * @returns {ModuleConfig<'reranker-contextualai', RerankerContextualAIConfig | undefined>} The configuration object.
+   */
+  contextualai: (
+    config?: RerankerContextualAIConfig
+  ): ModuleConfig<'reranker-contextualai', RerankerContextualAIConfig | undefined> => {
+    return {
+      name: 'reranker-contextualai',
+      config: config,
+    };
+  },
   /**
    * Create a `ModuleConfig<'reranker-jinaai', RerankerJinaAIConfig>` object for use when reranking using the `reranker-jinaai` module.
    *

src/collections/configure/types/generative.ts

@@ -2,6 +2,7 @@ import {
   GenerativeAWSConfig,
   GenerativeAnthropicConfig,
   GenerativeAnyscaleConfig,
+  GenerativeContextualAIConfig,
   GenerativeDatabricksConfig,
   GenerativeFriendliAIConfig,
   GenerativeMistralConfig,
@@ -58,12 +59,22 @@ export type GenerativePaLMConfigCreate = GenerativePaLMConfig;
 
 export type GenerativeXAIConfigCreate = GenerativeXAIConfig;
 
+export type GenerativeContextualAIConfigCreate = {
+  model?: string;
+  maxTokens?: number;
+  temperature?: number;
+  topP?: number;
+  systemPrompt?: string;
+  avoidCommentary?: boolean;
+};
+
 export type GenerativeConfigCreate =
   | GenerativeAnthropicConfigCreate
   | GenerativeAnyscaleConfigCreate
   | GenerativeAWSConfigCreate
   | GenerativeAzureOpenAIConfigCreate
   | GenerativeCohereConfigCreate
+  | GenerativeContextualAIConfigCreate
   | GenerativeDatabricksConfigCreate
   | GenerativeFriendliAIConfigCreate
   | GenerativeMistralConfigCreate
@@ -83,6 +94,8 @@ export type GenerativeConfigCreateType<G> = G extends 'generative-anthropic'
   ? GenerativeAzureOpenAIConfigCreate
   : G extends 'generative-cohere'
   ? GenerativeCohereConfigCreate
+  : G extends 'generative-contextualai'
+  ? GenerativeContextualAIConfigCreate
   : G extends 'generative-databricks'
   ? GenerativeDatabricksConfigCreate
   : G extends 'generative-friendliai'

src/collections/configure/unit.test.ts

@@ -5,6 +5,7 @@ import {
   GenerativeAnyscaleConfig,
   GenerativeAzureOpenAIConfig,
   GenerativeCohereConfig,
+  GenerativeContextualAIConfig,
   GenerativeDatabricksConfig,
   GenerativeFriendliAIConfig,
   GenerativeGoogleConfig,
@@ -14,6 +15,7 @@ import {
   GenerativeXAIConfig,
   ModuleConfig,
   RerankerCohereConfig,
+  RerankerContextualAIConfig,
   RerankerJinaAIConfig,
   RerankerNvidiaConfig,
   RerankerTransformersConfig,
@@ -1931,6 +1933,38 @@ describe('Unit testing of the generative factory class', () => {
     });
   });
 
+  it('should create the correct GenerativeContextualAIConfig type with required & default values', () => {
+    const config = configure.generative.contextualai();
+    expect(config).toEqual<ModuleConfig<'generative-contextualai', GenerativeContextualAIConfig | undefined>>(
+      {
+        name: 'generative-contextualai',
+        config: undefined,
+      }
+    );
+  });
+
+  it('should create the correct GenerativeContextualAIConfig type with all values', () => {
+    const config = configure.generative.contextualai({
+      model: 'v2',
+      maxTokens: 100,
+      temperature: 0.7,
+      topP: 0.9,
+      systemPrompt: 'system',
+      avoidCommentary: false,
+    });
+    expect(config).toEqual<ModuleConfig<'generative-contextualai', GenerativeContextualAIConfig>>({
+      name: 'generative-contextualai',
+      config: {
+        model: 'v2',
+        maxTokensProperty: 100,
+        temperatureProperty: 0.7,
+        topPProperty: 0.9,
+        systemPromptProperty: 'system',
+        avoidCommentaryProperty: false,
+      },
+    });
+  });
+
   it('should create the correct GenerativeCohereConfig type with all values', () => {
     const config = configure.generative.cohere({
       k: 5,
@@ -2265,6 +2299,32 @@ describe('Unit testing of the reranker factory class', () => {
     });
   });
 
+  it('should create the correct RerankerContextualAIConfig type using required & default values', () => {
+    const config = configure.reranker.contextualai();
+    expect(config).toEqual<ModuleConfig<'reranker-contextualai', RerankerContextualAIConfig | undefined>>({
+      name: 'reranker-contextualai',
+      config: undefined,
+    });
+  });
+
+  it('should create the correct RerankerContextualAIConfig type with all values', () => {
+    const config = configure.reranker.contextualai({
+      baseURL: 'https://api.contextual.ai',
+      model: 'ctxl-rerank-v2-instruct-multilingual',
+      instruction: 'Custom reranking instruction',
+      topN: 10,
+    });
+    expect(config).toEqual<ModuleConfig<'reranker-contextualai', RerankerContextualAIConfig | undefined>>({
+      name: 'reranker-contextualai',
+      config: {
+        baseURL: 'https://api.contextual.ai',
+        model: 'ctxl-rerank-v2-instruct-multilingual',
+        instruction: 'Custom reranking instruction',
+        topN: 10,
+      },
+    });
+  });
+
   it('should create the correct RerankerVoyageAIConfig type with all values', () => {
     const config = configure.reranker.voyageAI({
       baseURL: 'base-url',

src/collections/generate/config.ts

@@ -6,6 +6,7 @@ import {
   GenerativeAnyscaleConfigRuntime,
   GenerativeCohereConfigRuntime,
   GenerativeConfigRuntimeType,
+  GenerativeContextualAIConfigRuntime,
   GenerativeDatabricksConfigRuntime,
   GenerativeFriendliAIConfigRuntime,
   GenerativeGoogleConfigRuntime,
@@ -297,4 +298,29 @@ export const generativeParameters = {
         : undefined,
     };
   },
+  /**
+   * Create a `ModuleConfig<'generative-contextualai', GenerativeConfigRuntimeType<'generative-contextualai'> | undefined>`
+   * object for use when performing runtime-specific AI generation using the `generative-contextualai` module.
+   */
+  contextualai(
+    config?: GenerativeContextualAIConfigRuntime
+  ): ModuleConfig<
+    'generative-contextualai',
+    GenerativeConfigRuntimeType<'generative-contextualai'> | undefined
+  > {
+    // Contextual AI does not require special GRPC wrappers; pass primitives directly
+    return {
+      name: 'generative-contextualai',
+      config: config
+        ? {
+            model: config.model,
+            maxTokens: config.maxTokens,
+            temperature: config.temperature,
+            topP: config.topP,
+            systemPrompt: config.systemPrompt,
+            avoidCommentary: config.avoidCommentary,
+          }
+        : undefined,
+    };
+  },
 };