Skip to content

Commit

Permalink
Add support for generative-friendliai
Browse files Browse the repository at this point in the history
  • Loading branch information
tsmith023 committed Sep 16, 2024
1 parent 604f9dc commit b450b77
Show file tree
Hide file tree
Showing 6 changed files with 112 additions and 13 deletions.
33 changes: 31 additions & 2 deletions src/collections/config/types/generative.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,13 @@ export type GenerativeAnyscaleConfig = {
temperature?: number;
};

export type GenerativeFriendliAIConfig = {
baseURL?: string;
maxTokens?: number;
model?: string;
temperature?: number;
};

export type GenerativeMistralConfig = {
maxTokens?: number;
model?: string;
Expand Down Expand Up @@ -75,15 +82,36 @@ export type GenerativePaLMConfig = {
};

export type GenerativeConfig =
| GenerativeOpenAIConfig
| GenerativeAnthropicConfig
| GenerativeAnyscaleConfig
| GenerativeAWSConfig
| GenerativeAzureOpenAIConfig
| GenerativeCohereConfig
| GenerativeMistralConfig
| GenerativeOctoAIConfig
| GenerativeOllamaConfig
| GenerativeOpenAIConfig
| GenerativePaLMConfig
| Record<string, any>
| undefined;

export type GenerativeConfigType<G> = G extends 'generative-openai'
export type GenerativeConfigType<G> = G extends 'generative-anthropic'
? GenerativeAnthropicConfig
: G extends 'generative-anyscale'
? GenerativeAnyscaleConfig
: G extends 'generative-aws'
? GenerativeAWSConfig
: G extends 'generative-azure-openai'
? GenerativeOpenAIConfig
: G extends 'generative-cohere'
? GenerativeAzureOpenAIConfig
: 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
Expand All @@ -95,6 +123,7 @@ export type GenerativeSearch =
| 'generative-anthropic'
| 'generative-anyscale'
| 'generative-aws'
| 'generative-azure-openai'
| 'generative-mistral'
| 'generative-octoai'
| 'generative-ollama'
Expand Down
1 change: 1 addition & 0 deletions src/collections/config/types/reranker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ export type RerankerJinaAIConfig = {

export type RerankerConfig =
| RerankerCohereConfig
| RerankerJinaAIConfig
| RerankerTransformersConfig
| RerankerVoyageAIConfig
| Record<string, any>
Expand Down
15 changes: 15 additions & 0 deletions src/collections/configure/generative.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import {
GenerativeAnyscaleConfig,
GenerativeAzureOpenAIConfig,
GenerativeCohereConfig,
GenerativeFriendliAIConfig,
GenerativeMistralConfig,
GenerativeOctoAIConfig,
GenerativeOllamaConfig,
Expand All @@ -17,6 +18,7 @@ import {
GenerativeAnyscaleConfigCreate,
GenerativeAzureOpenAIConfigCreate,
GenerativeCohereConfigCreate,
GenerativeFriendliAIConfigCreate,
GenerativeMistralConfigCreate,
GenerativeOctoAIConfigCreate,
GenerativeOllamaConfigCreate,
Expand Down Expand Up @@ -121,6 +123,19 @@ export default {
: undefined,
};
},
/**
* Create a `ModuleConfig<'generative-friendliai', GenerativeFriendliAIConfig | undefined>` object for use when performing AI generation using the `generative-friendliai` module.
*
* See the [documentation](https://weaviate.io/developers/weaviate/model-providers/friendliai/generative) for detailed usage.
*/
friendliai(
config?: GenerativeFriendliAIConfigCreate
): ModuleConfig<'generative-friendliai', GenerativeFriendliAIConfig | undefined> {
return {
name: 'generative-friendliai',
config,
};
},
/**
* Create a `ModuleConfig<'generative-mistral', GenerativeMistralConfig | undefined>` object for use when performing AI generation using the `generative-mistral` module.
*
Expand Down
45 changes: 35 additions & 10 deletions src/collections/configure/types/generative.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import {
GenerativeAWSConfig,
GenerativeAnthropicConfig,
GenerativeAnyscaleConfig,
GenerativeFriendliAIConfig,
GenerativeMistralConfig,
GenerativeOctoAIConfig,
GenerativeOllamaConfig,
Expand All @@ -17,12 +18,12 @@ export type GenerativeOpenAIConfigBaseCreate = {
topP?: number;
};

export type GenerativeOpenAIConfigCreate = GenerativeOpenAIConfigBaseCreate & {
model?: string;
};

export type GenerativeAnthropicConfigCreate = GenerativeAnthropicConfig;

export type GenerativeAnyscaleConfigCreate = GenerativeAnyscaleConfig;

export type GenerativeAWSConfigCreate = GenerativeAWSConfig;

export type GenerativeAzureOpenAIConfigCreate = GenerativeOpenAIConfigBaseCreate & {
resourceName: string;
deploymentId: string;
Expand All @@ -37,29 +38,53 @@ export type GenerativeCohereConfigCreate = {
temperature?: number;
};

export type GenerativeAnyscaleConfigCreate = GenerativeAnyscaleConfig;

export type GenerativeAWSConfigCreate = GenerativeAWSConfig;
export type GenerativeFriendliAIConfigCreate = GenerativeFriendliAIConfig;

export type GenerativeMistralConfigCreate = GenerativeMistralConfig;

export type GenerativeOctoAIConfigCreate = GenerativeOctoAIConfig;

export type GenerativeOllamaConfigCreate = GenerativeOllamaConfig;

export type GenerativeOpenAIConfigCreate = GenerativeOpenAIConfigBaseCreate & {
model?: string;
};

export type GenerativePaLMConfigCreate = GenerativePaLMConfig;

export type GenerativeConfigCreate =
| GenerativeOpenAIConfigCreate
| GenerativeAnthropicConfigCreate
| GenerativeAnyscaleConfigCreate
| GenerativeAWSConfigCreate
| GenerativeAzureOpenAIConfigCreate
| GenerativeCohereConfigCreate
| GenerativeFriendliAIConfigCreate
| GenerativeMistralConfigCreate
| GenerativeOctoAIConfigCreate
| GenerativeOllamaConfigCreate
| GenerativeOpenAIConfigCreate
| GenerativePaLMConfigCreate
| Record<string, any>
| undefined;

export type GenerativeConfigCreateType<G> = G extends 'generative-openai'
? GenerativeOpenAIConfigCreate
export type GenerativeConfigCreateType<G> = G extends 'generative-anthropic'
? GenerativeAnthropicConfigCreate
: G extends 'generative-aws'
? GenerativeAWSConfigCreate
: G extends 'generative-azure-openai'
? GenerativeAzureOpenAIConfigCreate
: G extends 'generative-cohere'
? GenerativeCohereConfigCreate
: G extends 'generative-friendliai'
? GenerativeFriendliAIConfigCreate
: G extends 'generative-mistral'
? GenerativeMistralConfigCreate
: G extends 'generative-octoai'
? GenerativeOctoAIConfigCreate
: G extends 'generative-ollama'
? GenerativeOllamaConfigCreate
: G extends 'generative-openai'
? GenerativeOpenAIConfigCreate
: G extends 'generative-palm'
? GenerativePaLMConfigCreate
: G extends 'none'
Expand Down
27 changes: 27 additions & 0 deletions src/collections/configure/unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import {
GenerativeAnyscaleConfig,
GenerativeAzureOpenAIConfig,
GenerativeCohereConfig,
GenerativeFriendliAIConfig,
GenerativeMistralConfig,
GenerativeOctoAIConfig,
GenerativeOllamaConfig,
Expand Down Expand Up @@ -1258,6 +1259,32 @@ describe('Unit testing of the generative factory class', () => {
});
});

it('should create the correct GenerativeFriendliAIConfig type with required & default values', () => {
const config = configure.generative.friendliai();
expect(config).toEqual<ModuleConfig<'generative-friendliai', GenerativeFriendliAIConfig | undefined>>({
name: 'generative-friendliai',
config: undefined,
});
});

it('should create the correct GenerativeFriendliAIConfig type with all values', () => {
const config = configure.generative.friendliai({
baseURL: 'base-url',
maxTokens: 100,
model: 'model',
temperature: 0.5,
});
expect(config).toEqual<ModuleConfig<'generative-friendliai', GenerativeFriendliAIConfig | undefined>>({
name: 'generative-friendliai',
config: {
baseURL: 'base-url',
maxTokens: 100,
model: 'model',
temperature: 0.5,
},
});
});

it('should create the correct GenerativeMistralConfig type with required & default values', () => {
const config = configure.generative.mistral();
expect(config).toEqual<ModuleConfig<'generative-mistral', GenerativeMistralConfig | undefined>>({
Expand Down
4 changes: 3 additions & 1 deletion src/collections/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,9 @@ const collections = (connection: Connection, dbVersionSupport: DbVersionSupport)

const moduleConfig: any = {};
if (config.generative) {
moduleConfig[config.generative.name] = config.generative.config ? config.generative.config : {};
const generative =
config.generative.name === 'generative-azure-openai' ? 'generative-openai' : config.generative.name;
moduleConfig[generative] = config.generative.config ? config.generative.config : {};
}
if (config.reranker) {
moduleConfig[config.reranker.name] = config.reranker.config ? config.reranker.config : {};
Expand Down

0 comments on commit b450b77

Please sign in to comment.