Add support for generative-databricks

weaviate · Sep 16, 2024 · 8f35761 · 8f35761
1 parent b450b77
commit 8f35761
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 13 deletions.
diff --git a/src/collections/config/types/generative.ts b/src/collections/config/types/generative.ts
@@ -28,6 +28,23 @@ export type GenerativeAnyscaleConfig = {
   temperature?: number;
 };
 
+export type GenerativeCohereConfig = {
+  kProperty?: number;
+  model?: string;
+  maxTokensProperty?: number;
+  returnLikelihoodsProperty?: string;
+  stopSequencesProperty?: string[];
+  temperatureProperty?: number;
+};
+
+export type GenerativeDatabricksConfig = {
+  endpoint: string;
+  maxTokens?: number;
+  temperature?: number;
+  topK?: number;
+  topP?: number;
+};
+
 export type GenerativeFriendliAIConfig = {
   baseURL?: string;
   maxTokens?: number;
@@ -62,15 +79,6 @@ export type GenerativeAzureOpenAIConfig = GenerativeOpenAIConfigBase & {
   deploymentId: string;
 };
 
-export type GenerativeCohereConfig = {
-  kProperty?: number;
-  model?: string;
-  maxTokensProperty?: number;
-  returnLikelihoodsProperty?: string;
-  stopSequencesProperty?: string[];
-  temperatureProperty?: number;
-};
-
 export type GenerativePaLMConfig = {
   apiEndpoint?: string;
   maxOutputTokens?: number;
@@ -105,15 +113,17 @@ export type GenerativeConfigType<G> = G extends 'generative-anthropic'
   ? GenerativeOpenAIConfig
   : G extends 'generative-cohere'
   ? GenerativeAzureOpenAIConfig
+  : G extends 'generative-databricks'
+  ? GenerativeDatabricksConfig
+  : G extends 'generative-friendliai'
+  ? GenerativeFriendliAIConfig
   : G extends 'generative-mistral'
   ? GenerativeMistralConfig
   : G extends 'generative-octoai'
   ? GenerativeOctoAIConfig
   : G extends 'generative-ollama'
   ? GenerativeOllamaConfig
   : G extends 'generative-openai'
-  ? GenerativeCohereConfig
-  : G extends 'generative-palm'
   ? GenerativePaLMConfig
   : G extends 'none'
   ? undefined
@@ -124,11 +134,13 @@ export type GenerativeSearch =
   | 'generative-anyscale'
   | 'generative-aws'
   | 'generative-azure-openai'
+  | 'generative-cohere'
+  | 'generative-databricks'
+  | 'generative-friendliai'
   | 'generative-mistral'
   | 'generative-octoai'
   | 'generative-ollama'
   | 'generative-openai'
-  | 'generative-cohere'
   | 'generative-palm'
   | 'none'
   | string;
diff --git a/src/collections/configure/generative.ts b/src/collections/configure/generative.ts
@@ -4,6 +4,7 @@ import {
   GenerativeAnyscaleConfig,
   GenerativeAzureOpenAIConfig,
   GenerativeCohereConfig,
+  GenerativeDatabricksConfig,
   GenerativeFriendliAIConfig,
   GenerativeMistralConfig,
   GenerativeOctoAIConfig,
@@ -18,6 +19,7 @@ import {
   GenerativeAnyscaleConfigCreate,
   GenerativeAzureOpenAIConfigCreate,
   GenerativeCohereConfigCreate,
+  GenerativeDatabricksConfigCreate,
   GenerativeFriendliAIConfigCreate,
   GenerativeMistralConfigCreate,
   GenerativeOctoAIConfigCreate,
@@ -104,7 +106,7 @@ export default {
    * See the [documentation](https://weaviate.io/developers/weaviate/model-providers/cohere/generative) for detailed usage.
    *
    * @param {GenerativeCohereConfigCreate} [config] The configuration for the `generative-cohere` module.
-   * @returns {ModuleConfig<'generative-cohere', GenerativeCohereConfig>} The configuration object.
+   * @returns {ModuleConfig<'generative-cohere', GenerativeCohereConfig | undefined>} The configuration object.
    */
   cohere: (
     config?: GenerativeCohereConfigCreate
@@ -123,6 +125,22 @@ export default {
         : undefined,
     };
   },
+  /**
+   * Create a `ModuleConfig<'generative-databricks', GenerativeDatabricksConfig>` object for use when performing AI generation using the `generative-databricks` module.
+   *
+   * See the [documentation](https://weaviate.io/developers/weaviate/model-providers/databricks/generative) for detailed usage.
+   *
+   * @param {GenerativeDatabricksConfigCreate} config The configuration for the `generative-databricks` module.
+   * @returns {ModuleConfig<'generative-databricks', GenerativeDatabricksConfig>} The configuration object.
+   */
+  databricks: (
+    config: GenerativeDatabricksConfigCreate
+  ): ModuleConfig<'generative-databricks', GenerativeDatabricksConfig> => {
+    return {
+      name: 'generative-databricks',
+      config,
+    };
+  },
   /**
    * Create a `ModuleConfig<'generative-friendliai', GenerativeFriendliAIConfig | undefined>` object for use when performing AI generation using the `generative-friendliai` module.
    *

diff --git a/src/collections/configure/types/generative.ts b/src/collections/configure/types/generative.ts
@@ -2,6 +2,7 @@ import {
   GenerativeAWSConfig,
   GenerativeAnthropicConfig,
   GenerativeAnyscaleConfig,
+  GenerativeDatabricksConfig,
   GenerativeFriendliAIConfig,
   GenerativeMistralConfig,
   GenerativeOctoAIConfig,
@@ -38,6 +39,8 @@ export type GenerativeCohereConfigCreate = {
   temperature?: number;
 };
 
+export type GenerativeDatabricksConfigCreate = GenerativeDatabricksConfig;
+
 export type GenerativeFriendliAIConfigCreate = GenerativeFriendliAIConfig;
 
 export type GenerativeMistralConfigCreate = GenerativeMistralConfig;
@@ -58,6 +61,7 @@ export type GenerativeConfigCreate =
   | GenerativeAWSConfigCreate
   | GenerativeAzureOpenAIConfigCreate
   | GenerativeCohereConfigCreate
+  | GenerativeDatabricksConfigCreate
   | GenerativeFriendliAIConfigCreate
   | GenerativeMistralConfigCreate
   | GenerativeOctoAIConfigCreate
@@ -75,6 +79,8 @@ export type GenerativeConfigCreateType<G> = G extends 'generative-anthropic'
   ? GenerativeAzureOpenAIConfigCreate
   : G extends 'generative-cohere'
   ? GenerativeCohereConfigCreate
+  : G extends 'generative-databricks'
+  ? GenerativeDatabricksConfigCreate
   : G extends 'generative-friendliai'
   ? GenerativeFriendliAIConfigCreate
   : G extends 'generative-mistral'

diff --git a/src/collections/configure/unit.test.ts b/src/collections/configure/unit.test.ts
@@ -4,6 +4,7 @@ import {
   GenerativeAnyscaleConfig,
   GenerativeAzureOpenAIConfig,
   GenerativeCohereConfig,
+  GenerativeDatabricksConfig,
   GenerativeFriendliAIConfig,
   GenerativeMistralConfig,
   GenerativeOctoAIConfig,
@@ -1259,6 +1260,38 @@ describe('Unit testing of the generative factory class', () => {
     });
   });
 
+  it('should create the correct GenerativeDatabricksConfig type with required & default values', () => {
+    const config = configure.generative.databricks({
+      endpoint: 'endpoint',
+    });
+    expect(config).toEqual<ModuleConfig<'generative-databricks', GenerativeDatabricksConfig>>({
+      name: 'generative-databricks',
+      config: {
+        endpoint: 'endpoint',
+      },
+    });
+  });
+
+  it('should create the correct GenerativeDatabricksConfig type with all values', () => {
+    const config = configure.generative.databricks({
+      endpoint: 'endpoint',
+      maxTokens: 100,
+      temperature: 0.5,
+      topK: 10,
+      topP: 0.8,
+    });
+    expect(config).toEqual<ModuleConfig<'generative-databricks', GenerativeDatabricksConfig>>({
+      name: 'generative-databricks',
+      config: {
+        endpoint: 'endpoint',
+        maxTokens: 100,
+        temperature: 0.5,
+        topK: 10,
+        topP: 0.8,
+      },
+    });
+  });
+
   it('should create the correct GenerativeFriendliAIConfig type with required & default values', () => {
     const config = configure.generative.friendliai();
     expect(config).toEqual<ModuleConfig<'generative-friendliai', GenerativeFriendliAIConfig | undefined>>({