Merge pull request #283 from weaviate/modules/add-generative-xai

Add `generative-xai` when configuring collections and generate queries

Author: tsmith023

src/collections/config/types/generative.ts

@@ -93,6 +93,14 @@ export type GenerativeGoogleConfig = {
   topP?: number;
 };
 
+export type GenerativeXAIConfig = {
+  baseURL?: string;
+  maxTokens?: number;
+  model?: string;
+  temperature?: number;
+  topP?: number;
+};
+
 export type GenerativeConfig =
   | GenerativeAnthropicConfig
   | GenerativeAnyscaleConfig
@@ -106,6 +114,7 @@ export type GenerativeConfig =
   | GenerativeOllamaConfig
   | GenerativeOpenAIConfig
   | GenerativePaLMConfig
+  | GenerativeXAIConfig
   | Record<string, any>
   | undefined;
 
@@ -133,6 +142,8 @@ export type GenerativeConfigType<G> = G extends 'generative-anthropic'
   ? GenerativeOpenAIConfig
   : G extends GenerativePalm
   ? GenerativePaLMConfig
+  : G extends 'generative-xai'
+  ? GenerativeXAIConfig
   : G extends 'none'
   ? undefined
   : Record<string, any> | undefined;
@@ -153,5 +164,6 @@ export type GenerativeSearch =
   | 'generative-ollama'
   | 'generative-openai'
   | GenerativePalm
+  | 'generative-xai'
   | 'none'
   | string;

src/collections/configure/generative.ts

@@ -12,6 +12,7 @@ import {
   GenerativeOllamaConfig,
   GenerativeOpenAIConfig,
   GenerativePaLMConfig,
+  GenerativeXAIConfig,
   ModuleConfig,
 } from '../config/types/index.js';
 import {
@@ -27,6 +28,7 @@ import {
   GenerativeOllamaConfigCreate,
   GenerativeOpenAIConfigCreate,
   GenerativePaLMConfigCreate,
+  GenerativeXAIConfigCreate,
 } from '../index.js';
 
 export default {
@@ -263,4 +265,20 @@ export default {
       config,
     };
   },
+  /**
+   * Create a `ModuleConfig<'generative-xai', GenerativeXAIConfig | undefined>` object for use when performing AI generation using the `generative-xai` module.
+   *
+   * See the [documentation](https://weaviate.io/developers/weaviate/model-providers/xai/generative) for detailed usage.
+   *
+   * @param {GenerativeXAIConfigCreate} [config] The configuration for the `generative-xai` module.
+   * @returns {ModuleConfig<'generative-xai', GenerativeXAIConfig | undefined>} The configuration object.
+   */
+  xai: (
+    config?: GenerativeXAIConfigCreate
+  ): ModuleConfig<'generative-xai', GenerativeXAIConfig | undefined> => {
+    return {
+      name: 'generative-xai',
+      config,
+    };
+  },
 };

src/collections/configure/types/generative.ts

@@ -8,6 +8,7 @@ import {
   GenerativeNvidiaConfig,
   GenerativeOllamaConfig,
   GenerativePaLMConfig,
+  GenerativeXAIConfig,
 } from '../../index.js';
 
 export type GenerativeOpenAIConfigBaseCreate = {
@@ -55,6 +56,8 @@ export type GenerativeOpenAIConfigCreate = GenerativeOpenAIConfigBaseCreate & {
 
 export type GenerativePaLMConfigCreate = GenerativePaLMConfig;
 
+export type GenerativeXAIConfigCreate = GenerativeXAIConfig;
+
 export type GenerativeConfigCreate =
   | GenerativeAnthropicConfigCreate
   | GenerativeAnyscaleConfigCreate
@@ -68,6 +71,7 @@ export type GenerativeConfigCreate =
   | GenerativeOllamaConfigCreate
   | GenerativeOpenAIConfigCreate
   | GenerativePaLMConfigCreate
+  | GenerativeXAIConfigCreate
   | Record<string, any>
   | undefined;
 
@@ -93,6 +97,8 @@ export type GenerativeConfigCreateType<G> = G extends 'generative-anthropic'
   ? GenerativeOpenAIConfigCreate
   : G extends 'generative-palm'
   ? GenerativePaLMConfigCreate
+  : G extends 'generative-xai'
+  ? GenerativeXAIConfigCreate
   : G extends 'none'
   ? undefined
   : Record<string, any> | undefined;

src/collections/configure/unit.test.ts

@@ -10,6 +10,7 @@ import {
   GenerativeMistralConfig,
   GenerativeOllamaConfig,
   GenerativeOpenAIConfig,
+  GenerativeXAIConfig,
   ModuleConfig,
   VectorConfigCreate,
 } from '../types/index.js';
@@ -1816,6 +1817,34 @@ describe('Unit testing of the generative factory class', () => {
     });
   });
 
+  it('should create the correct GenerativeXAIConfig type with required & default values', () => {
+    const config = configure.generative.xai();
+    expect(config).toEqual<ModuleConfig<'generative-xai', GenerativeXAIConfig | undefined>>({
+      name: 'generative-xai',
+      config: undefined,
+    });
+  });
+
+  it('should create the correct GenerativeXAIConfig type with all values', () => {
+    const config = configure.generative.xai({
+      baseURL: 'base-url',
+      maxTokens: 100,
+      model: 'model',
+      temperature: 0.5,
+      topP: 0.8,
+    });
+    expect(config).toEqual<ModuleConfig<'generative-xai', GenerativeXAIConfig | undefined>>({
+      name: 'generative-xai',
+      config: {
+        baseURL: 'base-url',
+        maxTokens: 100,
+        model: 'model',
+        temperature: 0.5,
+        topP: 0.8,
+      },
+    });
+  });
+
   it('should create the correct GeneratGoogleConfig type with required & default values', () => {
     const config = configure.generative.google();
     expect(config).toEqual<ModuleConfig<'generative-google', undefined>>({

src/collections/generate/config.ts

@@ -13,6 +13,7 @@ import {
   GenerativeNvidiaConfigRuntime,
   GenerativeOllamaConfigRuntime,
   GenerativeOpenAIConfigRuntime,
+  GenerativeXAIConfigRuntime,
 } from '../index.js';
 
 export const generativeParameters = {
@@ -277,4 +278,23 @@ export const generativeParameters = {
         : undefined,
     };
   },
+  /**
+   * Create a `ModuleConfig<'generative-xai', GenerativeConfigRuntimeType<'generative-xai'> | undefined>` object for use when performing runtime-specific AI generation using the `generative-xai` module.
+   *
+   * See the [documentation](https://weaviate.io/developers/weaviate/model-providers/xai/generative) for detailed usage.
+   */
+  xai: (
+    config?: GenerativeXAIConfigRuntime
+  ): ModuleConfig<'generative-xai', GenerativeConfigRuntimeType<'generative-xai'> | undefined> => {
+    const { baseURL, ...rest } = config || {};
+    return {
+      name: 'generative-xai',
+      config: config
+        ? {
+            ...rest,
+            baseUrl: baseURL,
+          }
+        : undefined,
+    };
+  },
 };

src/collections/generate/mock.test.ts

@@ -158,6 +158,8 @@ describe('Mock testing of generate with runtime config', () => {
     generativeParameters.ollama(model),
     generativeParameters.openAI(),
     generativeParameters.openAI(model),
+    generativeParameters.xai(),
+    generativeParameters.xai(model),
   ];
 
   tests.forEach((conf) => {

src/collections/generate/unit.test.ts

@@ -277,4 +277,37 @@ describe('Unit testing of the generativeParameters factory methods', () => {
       });
     });
   });
+
+  describe('xai', () => {
+    it('with defaults', () => {
+      const config = generativeParameters.xai();
+      expect(config).toEqual<
+        ModuleConfig<'generative-xai', GenerativeConfigRuntimeType<'generative-xai'> | undefined>
+      >({
+        name: 'generative-xai',
+        config: undefined,
+      });
+    });
+    it('with values', () => {
+      const config = generativeParameters.xai({
+        baseURL: 'http://localhost:8080',
+        maxTokens: 100,
+        model: 'model',
+        temperature: 0.5,
+        topP: 0.9,
+      });
+      expect(config).toEqual<
+        ModuleConfig<'generative-xai', GenerativeConfigRuntimeType<'generative-xai'> | undefined>
+      >({
+        name: 'generative-xai',
+        config: {
+          baseUrl: 'http://localhost:8080',
+          maxTokens: 100,
+          model: 'model',
+          temperature: 0.5,
+          topP: 0.9,
+        },
+      });
+    });
+  });
 });

src/collections/types/generate.ts

@@ -24,8 +24,10 @@ import {
   GenerativeOllamaMetadata,
   GenerativeOpenAI as GenerativeOpenAIGRPC,
   GenerativeOpenAIMetadata,
+  GenerativeXAI as GenerativeXAIGRPC,
+  GenerativeXAIMetadata,
 } from '../../proto/v1/generative.js';
-import { ModuleConfig } from '../index.js';
+import { GenerativeXAIConfig, ModuleConfig } from '../index.js';
 import { GroupByObject, GroupByResult, WeaviateGenericObject, WeaviateNonGenericObject } from './query.js';
 
 export type GenerativeGenericObject<
@@ -134,7 +136,8 @@ export type GenerativeConfigRuntime =
   | ModuleConfig<'generative-mistral', GenerativeConfigRuntimeType<'generative-mistral'> | undefined>
   | ModuleConfig<'generative-nvidia', GenerativeConfigRuntimeType<'generative-nvidia'> | undefined>
   | ModuleConfig<'generative-ollama', GenerativeConfigRuntimeType<'generative-ollama'> | undefined>
-  | ModuleConfig<'generative-openai', GenerativeConfigRuntimeType<'generative-openai'>>;
+  | ModuleConfig<'generative-openai', GenerativeConfigRuntimeType<'generative-openai'>>
+  | ModuleConfig<'generative-xai', GenerativeConfigRuntimeType<'generative-xai'> | undefined>;
 
 export type GenerativeConfigRuntimeType<G> = G extends 'generative-anthropic'
   ? Omit<GenerativeAnthropicGRPC, omitFields>
@@ -160,6 +163,8 @@ export type GenerativeConfigRuntimeType<G> = G extends 'generative-anthropic'
   ? Omit<GenerativeOllamaGRPC, omitFields>
   : G extends 'generative-openai'
   ? Omit<GenerativeOpenAIGRPC, omitFields> & { isAzure?: false }
+  : G extends 'generative-xai'
+  ? Omit<GenerativeXAIGRPC, omitFields>
   : G extends 'none'
   ? undefined
   : Record<string, any> | undefined;
@@ -191,6 +196,8 @@ export type GenerativeMetadata<C extends GenerativeConfigRuntime | undefined> =
     ? GenerativeOllamaMetadata
     : R['name'] extends 'generative-openai'
     ? GenerativeOpenAIMetadata
+    : R['name'] extends 'generative-xai'
+    ? GenerativeXAIMetadata
     : never
   : never;
 
@@ -312,3 +319,5 @@ export type GenerativeOpenAIConfigRuntime = {
   resourceName?: string | undefined;
   deploymentId?: string | undefined;
 };
+
+export type GenerativeXAIConfigRuntime = GenerativeXAIConfig;