All files / api/src/ai-providers/ollama api.ts

82.85% Statements 58/70
66.66% Branches 14/21
100% Functions 3/3
82.85% Lines 58/70

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 901x   1x   1x 1x 2x 2x 2x   2x 1x 1x   2x 2x 1x 3x 3x 2x     2x     2x         2x         2x 1x 1x     1x     2x 2x 2x   2x         2x 2x       2x 1x 3x 1x 6x 6x   6x 2x 1x 1x 1x 1x 2x   6x 6x 6x 3x 3x 6x 2x 2x 6x 1x 6x 6x 1x 1x 1x   1x  
import { OllamaCustomFieldsSchema } from '@api/ai-providers/ollama/types';
import type { InternalProviderAPIConfig } from '@shared/types/ai-providers/config';
import { FunctionName } from '@shared/types/api/request';
 
export const ollamaAPIConfig: InternalProviderAPIConfig = {
  headers: ({ saTarget }) => {
    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
    };
 
    if (saTarget.api_key) {
      headers['x-ollama-api-key'] = saTarget.api_key;
    }
 
    return headers;
  },
  getBaseURL: ({ saTarget }) => {
    const customHost = saTarget.custom_host;
    if (customHost) {
      try {
        // SECURITY: Comprehensive URL validation for custom_host
        // This prevents SSRF attacks and ensures only safe URLs are used
        const url = new URL(customHost);
 
        // Validate protocol
        if (!['http:', 'https:'].includes(url.protocol)) {
          throw new Error('Only HTTP and HTTPS protocols are allowed');
        }
 
        // Validate hostname is not empty
        if (!url.hostname) {
          throw new Error('Hostname is required');
        }
 
        // Validate port if specified
        if (url.port) {
          const portNum = Number.parseInt(url.port, 10);
          if (Number.isNaN(portNum) || portNum < 1 || portNum > 65535) {
            throw new Error('Invalid port number');
          }
        }
 
        // Prevent path traversal in URL
        if (
          url.pathname &&
          url.pathname !== '/' &&
          url.pathname.includes('..')
        ) {
          throw new Error('Path traversal not allowed');
        }
 
        // Return the sanitized URL without query params or hash
        return `${url.protocol}//${url.host}${url.pathname === '/' ? '' : url.pathname}`;
      } catch (error) {
        const message = error instanceof Error ? error.message : 'Invalid URL';
        throw new Error(`Invalid custom_host URL: ${message}`);
      }
    }
    return 'http://localhost:11434';
  },
  getEndpoint: ({ saRequestData, saTarget }) => {
    let mappedFn = saRequestData.functionName;
    const urlToFetch = saTarget.ollama_url_to_fetch;
 
    if (saRequestData.functionName === FunctionName.PROXY && urlToFetch) {
      if (urlToFetch.indexOf('/api/chat') > -1) {
        mappedFn = FunctionName.CHAT_COMPLETE;
      } else if (urlToFetch.indexOf('/embeddings') > -1) {
        mappedFn = FunctionName.EMBED;
      }
    }
 
    switch (mappedFn) {
      case FunctionName.CHAT_COMPLETE:
      case FunctionName.STREAM_CHAT_COMPLETE: {
        return '/v1/chat/completions'; // OpenAI-compatible endpoint
      }
      case FunctionName.EMBED: {
        return '/api/embeddings'; // Ollama-specific endpoint
      }
      default:
        return '';
    }
  },
  customFieldsSchema: OllamaCustomFieldsSchema,
  isAPIKeyRequired: false,
};
 
export default ollamaAPIConfig;