lobe-chat

1
import { ListResponse, Ollama as OllamaBrowser, ProgressResponse } from 'ollama/browser';
2

3
import { ModelProvider } from '@/libs/agent-runtime';
4
import { useUserStore } from '@/store/user';
5
import { keyVaultsConfigSelectors } from '@/store/user/selectors';
6
import { ChatErrorType } from '@/types/fetch';
7
import { createErrorResponse } from '@/utils/errorResponse';
8
import { getMessageError } from '@/utils/fetch';
9

10
const DEFAULT_BASE_URL = 'http://127.0.0.1:11434';
11

12
interface OllamaServiceParams {
13
  fetch?: typeof fetch;
14
}
15

16
export class OllamaService {
17
  private _host: string;
18
  private _client: OllamaBrowser;
19
  private _fetch?: typeof fetch;
20

21
  constructor(params: OllamaServiceParams = {}) {
22
    this._host = this.getHost();
23
    this._fetch = params.fetch;
24
    this._client = new OllamaBrowser({ fetch: params?.fetch, host: this._host });
25
  }
26

27
  getHost = (): string => {
28
    const config = keyVaultsConfigSelectors.ollamaConfig(useUserStore.getState());
29

30
    return config.baseURL || DEFAULT_BASE_URL;
31
  };
32

33
  getOllamaClient = () => {
34
    if (this.getHost() !== this._host) {
35
      this._host = this.getHost();
36
      this._client = new OllamaBrowser({ fetch: this._fetch, host: this.getHost() });
37
    }
38
    return this._client;
39
  };
40

41
  abort = () => {
42
    this._client.abort();
43
  };
44

45
  pullModel = async (model: string): Promise<AsyncIterable<ProgressResponse>> => {
46
    let response: Response | AsyncIterable<ProgressResponse>;
47
    try {
48
      response = await this.getOllamaClient().pull({ insecure: true, model, stream: true });
49
      return response;
50
    } catch {
51
      response = createErrorResponse(ChatErrorType.OllamaServiceUnavailable, {
52
        host: this.getHost(),
53
        message: 'please check whether your ollama service is available or set the CORS rules',
54
        provider: ModelProvider.Ollama,
55
      });
56
    }
57

58
    if (!response.ok) {
59
      throw await getMessageError(response);
60
    }
61
    return response.json();
62
  };
63

64
  getModels = async (): Promise<ListResponse> => {
65
    let response: Response | ListResponse;
66
    try {
67
      return await this.getOllamaClient().list();
68
    } catch {
69
      response = createErrorResponse(ChatErrorType.OllamaServiceUnavailable, {
70
        host: this.getHost(),
71
        message: 'please check whether your ollama service is available or set the CORS rules',
72
        provider: ModelProvider.Ollama,
73
      });
74
    }
75

76
    if (!response.ok) {
77
      throw await getMessageError(response);
78
    }
79
    return response.json();
80
  };
81
}
82

83
export const ollamaService = new OllamaService();
84