Skip to content

Commit

Permalink
🦙 feat: Fetch list of Ollama Models (danny-avila#2565)
Browse files Browse the repository at this point in the history
* 🦙 feat: Fetch list of Ollama Models

* style: better Tag text styling for light mode
  • Loading branch information
danny-avila authored Apr 27, 2024
1 parent 8a78500 commit 63ef15a
Show file tree
Hide file tree
Showing 5 changed files with 210 additions and 4 deletions.
55 changes: 55 additions & 0 deletions api/server/services/ModelService.js
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,59 @@ const { HttpsProxyAgent } = require('https-proxy-agent');
const { EModelEndpoint, defaultModels, CacheKeys } = require('librechat-data-provider');
const { extractBaseURL, inputSchema, processModelData, logAxiosError } = require('~/utils');
const getLogStores = require('~/cache/getLogStores');
const { logger } = require('~/config');

const { openAIApiKey, userProvidedOpenAI } = require('./Config/EndpointService').config;

/**
* Extracts the base URL from the provided URL.
* @param {string} fullURL - The full URL.
* @returns {string} The base URL.
*/
function deriveBaseURL(fullURL) {
try {
const parsedUrl = new URL(fullURL);
const protocol = parsedUrl.protocol;
const hostname = parsedUrl.hostname;
const port = parsedUrl.port;

// Check if the parsed URL components are meaningful
if (!protocol || !hostname) {
return fullURL;
}

// Reconstruct the base URL
return `${protocol}//${hostname}${port ? `:${port}` : ''}`;
} catch (error) {
logger.error('Failed to derive base URL', error);
return fullURL; // Return the original URL in case of any exception
}
}

/**
* Fetches Ollama models from the specified base API path.
* @param {string} baseURL
* @returns {Promise<string[]>} The Ollama models.
*/
const fetchOllamaModels = async (baseURL) => {
let models = [];
if (!baseURL) {
return models;
}
try {
const ollamaEndpoint = deriveBaseURL(baseURL);
/** @type {Promise<AxiosResponse<OllamaListResponse>>} */
const response = await axios.get(`${ollamaEndpoint}/api/tags`);
models = response.data.models.map((tag) => tag.name);
return models;
} catch (error) {
const logMessage =
'Failed to fetch models from Ollama API. If you are not using Ollama directly, and instead, through some aggregator or reverse proxy that handles fetching via OpenAI spec, ensure the name of the endpoint doesn\'t start with `ollama` (case-insensitive).';
logger.error(logMessage, error);
return [];
}
};

/**
* Fetches OpenAI models from the specified base API path or Azure, based on the provided configuration.
*
Expand Down Expand Up @@ -41,6 +91,10 @@ const fetchModels = async ({
return models;
}

if (name && name.toLowerCase().startsWith('ollama')) {
return await fetchOllamaModels(baseURL);
}

try {
const options = {
headers: {
Expand Down Expand Up @@ -227,6 +281,7 @@ const getGoogleModels = () => {

module.exports = {
fetchModels,
deriveBaseURL,
getOpenAIModels,
getChatGPTBrowserModels,
getAnthropicModels,
Expand Down
119 changes: 118 additions & 1 deletion api/server/services/ModelService.spec.js
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
const axios = require('axios');
const { logger } = require('~/config');

const { fetchModels, getOpenAIModels } = require('./ModelService');
const { fetchModels, getOpenAIModels, deriveBaseURL } = require('./ModelService');
jest.mock('~/utils', () => {
const originalUtils = jest.requireActual('~/utils');
return {
Expand Down Expand Up @@ -256,3 +257,119 @@ describe('getOpenAIModels sorting behavior', () => {
jest.clearAllMocks();
});
});

describe('fetchModels with Ollama specific logic', () => {
const mockOllamaData = {
data: {
models: [{ name: 'Ollama-Base' }, { name: 'Ollama-Advanced' }],
},
};

beforeEach(() => {
axios.get.mockResolvedValue(mockOllamaData);
});

afterEach(() => {
jest.clearAllMocks();
});

it('should fetch Ollama models when name starts with "ollama"', async () => {
const models = await fetchModels({
user: 'user789',
apiKey: 'testApiKey',
baseURL: 'https://api.ollama.test.com',
name: 'OllamaAPI',
});

expect(models).toEqual(['Ollama-Base', 'Ollama-Advanced']);
expect(axios.get).toHaveBeenCalledWith('https://api.ollama.test.com/api/tags'); // Adjusted to expect only one argument if no options are passed
});

it('should handle errors gracefully when fetching Ollama models fails', async () => {
axios.get.mockRejectedValue(new Error('Network error'));
const models = await fetchModels({
user: 'user789',
apiKey: 'testApiKey',
baseURL: 'https://api.ollama.test.com',
name: 'OllamaAPI',
});

expect(models).toEqual([]);
expect(logger.error).toHaveBeenCalled();
});

it('should return an empty array if no baseURL is provided', async () => {
const models = await fetchModels({
user: 'user789',
apiKey: 'testApiKey',
name: 'OllamaAPI',
});
expect(models).toEqual([]);
});

it('should not fetch Ollama models if the name does not start with "ollama"', async () => {
// Mock axios to return a different set of models for non-Ollama API calls
axios.get.mockResolvedValue({
data: {
data: [{ id: 'model-1' }, { id: 'model-2' }],
},
});

const models = await fetchModels({
user: 'user789',
apiKey: 'testApiKey',
baseURL: 'https://api.test.com',
name: 'TestAPI',
});

expect(models).toEqual(['model-1', 'model-2']);
expect(axios.get).toHaveBeenCalledWith(
'https://api.test.com/models', // Ensure the correct API endpoint is called
expect.any(Object), // Ensuring some object (headers, etc.) is passed
);
});
});

describe('deriveBaseURL', () => {
it('should extract the base URL correctly from a full URL with a port', () => {
const fullURL = 'https://example.com:8080/path?query=123';
const baseURL = deriveBaseURL(fullURL);
expect(baseURL).toEqual('https://example.com:8080');
});

it('should extract the base URL correctly from a full URL without a port', () => {
const fullURL = 'https://example.com/path?query=123';
const baseURL = deriveBaseURL(fullURL);
expect(baseURL).toEqual('https://example.com');
});

it('should handle URLs using the HTTP protocol', () => {
const fullURL = 'http://example.com:3000/path?query=123';
const baseURL = deriveBaseURL(fullURL);
expect(baseURL).toEqual('http://example.com:3000');
});

it('should return only the protocol and hostname if no port is specified', () => {
const fullURL = 'http://example.com/path?query=123';
const baseURL = deriveBaseURL(fullURL);
expect(baseURL).toEqual('http://example.com');
});

it('should handle URLs with uncommon protocols', () => {
const fullURL = 'ftp://example.com:2121/path?query=123';
const baseURL = deriveBaseURL(fullURL);
expect(baseURL).toEqual('ftp://example.com:2121');
});

it('should handle edge case where URL ends with a slash', () => {
const fullURL = 'https://example.com/';
const baseURL = deriveBaseURL(fullURL);
expect(baseURL).toEqual('https://example.com');
});

it('should return the original URL if the URL is invalid', () => {
const invalidURL = 'htp:/example.com:8080';
const result = deriveBaseURL(invalidURL);
expect(result).toBe(invalidURL);
});
});
33 changes: 33 additions & 0 deletions api/typedefs.js
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,11 @@
* @typedef {import('openai').OpenAI} OpenAI
* @memberof typedefs
*/
/**
* @exports AxiosResponse
* @typedef {import('axios').AxiosResponse} AxiosResponse
* @memberof typedefs
*/

/**
* @exports Anthropic
Expand Down Expand Up @@ -1145,3 +1150,31 @@
* @param {onTokenProgress} params.onTokenProgress
* @memberof typedefs
*/

/**
* @typedef {Object} OllamaModelDetails
* @property {string} parent_model - The identifier for the parent model, if any.
* @property {string} format - The format of the model.
* @property {string} family - The primary family to which the model belongs.
* @property {string[]} families - An array of families that include the model.
* @property {string} parameter_size - The size of the parameters of the model.
* @property {string} quantization_level - The level of quantization of the model.
* @memberof typedefs
*/

/**
* @typedef {Object} OllamaModel
* @property {string} name - The name of the model, including version tag.
* @property {string} model - A redundant copy of the name, including version tag.
* @property {string} modified_at - The ISO string representing the last modification date.
* @property {number} size - The size of the model in bytes.
* @property {string} digest - The digest hash of the model.
* @property {OllamaModelDetails} details - Detailed information about the model.
* @memberof typedefs
*/

/**
* @typedef {Object} OllamaListResponse
* @property {OllamaModel[]} models - the list of models available.
* @memberof typedefs
*/
2 changes: 1 addition & 1 deletion client/src/components/ui/Tag.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ const TagPrimitiveRoot = React.forwardRef<HTMLDivElement, TagProps>(
ref={ref}
{...props}
className={cn(
'flex max-h-8 items-center overflow-y-hidden rounded rounded-3xl border-2 border-green-600 bg-green-600/20 text-sm text-xs text-white',
'flex max-h-8 items-center overflow-y-hidden rounded rounded-3xl border-2 border-green-600 bg-green-600/20 text-sm text-xs text-green-600 dark:text-white',
className,
)}
>
Expand Down
5 changes: 3 additions & 2 deletions docs/install/configuration/ai_endpoints.md
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,6 @@ Some of the endpoints are marked as **Known,** which means they might have speci
**Notes:**

- **Known:** icon provided.
- **Known issue:** fetching list of models is not supported. See [Pull Request 2728](https://github.com/ollama/ollama/pull/2728).
- Download models with ollama run command. See [Ollama Library](https://ollama.com/library)
- It's recommend to use the value "current_model" for the `titleModel` to avoid loading more than 1 model per conversation.
- Doing so will dynamically use the current conversation model for the title generation.
Expand All @@ -307,7 +306,9 @@ Some of the endpoints are marked as **Known,** which means they might have speci
"dolphin-mixtral",
"mistral-openorca"
]
fetch: false # fetching list of models is not supported
# fetching list of models is supported but the `name` field must start
# with `ollama` (case-insensitive), as it does in this example.
fetch: true
titleConvo: true
titleModel: "current_model"
summarize: false
Expand Down

0 comments on commit 63ef15a

Please sign in to comment.