From 8f35761e88b3c41ef1d79141e5bb13fba7d8fbe6 Mon Sep 17 00:00:00 2001 From: Tommy Smith Date: Mon, 16 Sep 2024 12:20:55 +0100 Subject: [PATCH] Add support for `generative-databricks` --- src/collections/config/types/generative.ts | 36 ++++++++++++------- src/collections/configure/generative.ts | 20 ++++++++++- src/collections/configure/types/generative.ts | 6 ++++ src/collections/configure/unit.test.ts | 33 +++++++++++++++++ 4 files changed, 82 insertions(+), 13 deletions(-) diff --git a/src/collections/config/types/generative.ts b/src/collections/config/types/generative.ts index a8f2df1..4ad0dc4 100644 --- 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,6 +113,10 @@ export type GenerativeConfigType = 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' @@ -112,8 +124,6 @@ export type GenerativeConfigType = G extends 'generative-anthropic' : 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 index 7d88e9c..ebe8b15 100644 --- 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 index 74acacf..59df7f8 100644 --- 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 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 index 49d879e..9bcadd8 100644 --- 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>({ + 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>({ + 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>({