Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/collections/config/types/generative.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ export type GenerativeAnthropicConfig = {
};

export type GenerativeAnyscaleConfig = {
baseURL?: string;
model?: string;
temperature?: number;
};
Expand Down Expand Up @@ -54,6 +55,7 @@ export type GenerativeFriendliAIConfig = {
};

export type GenerativeMistralConfig = {
baseURL?: string;
maxTokens?: number;
model?: string;
temperature?: number;
Expand Down
10 changes: 10 additions & 0 deletions src/collections/config/types/reranker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ export type RerankerCohereConfig = {
};

export type RerankerVoyageAIConfig = {
baseURL?: string;
model?: 'rerank-lite-1' | string;
};

Expand All @@ -18,9 +19,15 @@ export type RerankerJinaAIConfig = {
| string;
};

export type RerankerNvidiaConfig = {
baseURL?: string;
model?: 'nvidia/rerank-qa-mistral-4b' | string;
};

export type RerankerConfig =
| RerankerCohereConfig
| RerankerJinaAIConfig
| RerankerNvidiaConfig
| RerankerTransformersConfig
| RerankerVoyageAIConfig
| Record<string, any>
Expand All @@ -29,6 +36,7 @@ export type RerankerConfig =
export type Reranker =
| 'reranker-cohere'
| 'reranker-jinaai'
| 'reranker-nvidia'
| 'reranker-transformers'
| 'reranker-voyageai'
| 'none'
Expand All @@ -38,6 +46,8 @@ export type RerankerConfigType<R> = R extends 'reranker-cohere'
? RerankerCohereConfig
: R extends 'reranker-jinaai'
? RerankerJinaAIConfig
: R extends 'reranker-nvidia'
? RerankerNvidiaConfig
: R extends 'reranker-transformers'
? RerankerTransformersConfig
: R extends 'reranker-voyageai'
Expand Down
2 changes: 2 additions & 0 deletions src/collections/config/types/vectorizer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -390,6 +390,8 @@ export type Text2VecNvidiaConfig = {
* See the [documentation](https://weaviate.io/developers/weaviate/model-providers/mistral/embeddings) for detailed usage.
*/
export type Text2VecMistralConfig = {
/** The base URL to use where API requests should go. */
baseURL?: string;
/** The model to use. */
model?: 'mistral-embed' | string;
/** Whether to vectorize the collection name. */
Expand Down
17 changes: 17 additions & 0 deletions src/collections/configure/reranker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import {
ModuleConfig,
RerankerCohereConfig,
RerankerJinaAIConfig,
RerankerNvidiaConfig,
RerankerVoyageAIConfig,
} from '../config/types/index.js';

Expand Down Expand Up @@ -38,6 +39,22 @@ export default {
config: config,
};
},
/**
* Create a `ModuleConfig<'reranker-nvidia', RerankerNvidiaConfig>` object for use when reranking using the `reranker-nvidia` module.
*
* See the [documentation](https://weaviate.io/developers/weaviate/model-providers/nvidia/reranker) for detailed usage.
*
* @param {RerankerNvidiaConfig} [config] The configuration for the `reranker-nvidia` module.
* @returns {ModuleConfig<'reranker-nvidia', RerankerNvidiaConfig | undefined>} The configuration object.
*/
nvidia: (
config?: RerankerNvidiaConfig
): ModuleConfig<'reranker-nvidia', RerankerNvidiaConfig | undefined> => {
return {
name: 'reranker-nvidia',
config: config,
};
},
/**
* Create a `ModuleConfig<'reranker-transformers', Record<string, never>>` object for use when reranking using the `reranker-transformers` module.
*
Expand Down
105 changes: 105 additions & 0 deletions src/collections/configure/unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,11 @@ import {
GenerativeOpenAIConfig,
GenerativeXAIConfig,
ModuleConfig,
RerankerCohereConfig,
RerankerJinaAIConfig,
RerankerNvidiaConfig,
RerankerTransformersConfig,
RerankerVoyageAIConfig,
VectorConfigCreate,
} from '../types/index.js';
import { configure } from './index.js';
Expand Down Expand Up @@ -1220,6 +1225,7 @@ describe('Unit testing of the vectorizer factory class', () => {

it('should create the correct Text2VecMistralConfig type with all values', () => {
const config = configure.vectorizer.text2VecMistral({
baseURL: 'base-url',
name: 'test',
model: 'model',
vectorizeCollectionName: true,
Expand All @@ -1233,6 +1239,7 @@ describe('Unit testing of the vectorizer factory class', () => {
vectorizer: {
name: 'text2vec-mistral',
config: {
baseURL: 'base-url',
model: 'model',
vectorizeCollectionName: true,
},
Expand Down Expand Up @@ -1567,12 +1574,14 @@ describe('Unit testing of the generative factory class', () => {

it('should create the correct GenerativeAnyscaleConfig type with all values', () => {
const config = configure.generative.anyscale({
baseURL: 'base-url',
model: 'model',
temperature: 0.5,
});
expect(config).toEqual<ModuleConfig<'generative-anyscale', GenerativeAnyscaleConfig | undefined>>({
name: 'generative-anyscale',
config: {
baseURL: 'base-url',
model: 'model',
temperature: 0.5,
},
Expand Down Expand Up @@ -1749,13 +1758,15 @@ describe('Unit testing of the generative factory class', () => {

it('should create the correct GenerativeMistralConfig type with all values', () => {
const config = configure.generative.mistral({
baseURL: 'base-url',
maxTokens: 100,
model: 'model',
temperature: 0.5,
});
expect(config).toEqual<ModuleConfig<'generative-mistral', GenerativeMistralConfig | undefined>>({
name: 'generative-mistral',
config: {
baseURL: 'base-url',
maxTokens: 100,
model: 'model',
temperature: 0.5,
Expand Down Expand Up @@ -1909,3 +1920,97 @@ describe('Unit testing of the generative factory class', () => {
});
});
});

describe('Unit testing of the reranker factory class', () => {
it('should create the correct RerankerCohereConfig type using required & default values', () => {
const config = configure.reranker.cohere();
expect(config).toEqual<ModuleConfig<'reranker-cohere', RerankerCohereConfig | undefined>>({
name: 'reranker-cohere',
config: undefined,
});
});

it('should create the correct RerankerCohereConfig type with all values', () => {
const config = configure.reranker.cohere({
model: 'model',
});
expect(config).toEqual<ModuleConfig<'reranker-cohere', RerankerCohereConfig | undefined>>({
name: 'reranker-cohere',
config: {
model: 'model',
},
});
});

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

it('should create the correct RerankerJinaAIConfig type with all values', () => {
const config = configure.reranker.jinaai({
model: 'model',
});
expect(config).toEqual<ModuleConfig<'reranker-jinaai', RerankerJinaAIConfig | undefined>>({
name: 'reranker-jinaai',
config: {
model: 'model',
},
});
});

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

it('should create the correct RerankerNvidiaConfig type with all values', () => {
const config = configure.reranker.nvidia({
baseURL: 'base-url',
model: 'model',
});
expect(config).toEqual<ModuleConfig<'reranker-nvidia', RerankerNvidiaConfig | undefined>>({
name: 'reranker-nvidia',
config: {
baseURL: 'base-url',
model: 'model',
},
});
});

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

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

it('should create the correct RerankerVoyageAIConfig type with all values', () => {
const config = configure.reranker.voyageAI({
baseURL: 'base-url',
model: 'model',
});
expect(config).toEqual<ModuleConfig<'reranker-voyageai', RerankerVoyageAIConfig | undefined>>({
name: 'reranker-voyageai',
config: {
baseURL: 'base-url',
model: 'model',
},
});
});
});
Loading