Coverage for aipyapp/llm/__init__.py: 87%

1#! /usr/bin/env python

2# -*- coding: utf-8 -*-

3

4

5 from .. import T

6 from .base import ChatMessage, BaseClient

7 from .base_openai import OpenAIBaseClient

8 from .client_claude import ClaudeClient

9 from .client_ollama import OllamaClient

10 from .client_oauth2 import OAuth2Client

11 from .models import ModelRegistry, ModelCapability

12

13 __all__ = ['ChatMessage', 'CLIENTS', 'ModelRegistry', 'ModelCapability']

14

15 class OpenAIClient(OpenAIBaseClient):

16 MODEL = 'gpt-4o'

17

18 class GeminiClient(OpenAIBaseClient):

19 BASE_URL = 'https://generativelanguage.googleapis.com/v1beta/'

20 MODEL = 'gemini-2.5-flash'

21

22 class DeepSeekClient(OpenAIBaseClient):

23 BASE_URL = 'https://api.deepseek.com'

24 MODEL = 'deepseek-chat'

25

26 class GrokClient(OpenAIBaseClient):

27 BASE_URL = 'https://api.x.ai/v1/'

28 MODEL = 'grok-3-mini'

29

30 class TrustClient(OpenAIBaseClient):

31 MODEL = 'auto'

32

33 def get_base_url(self):

34 return self.config.get("base_url") or T("https://sapi.trustoken.ai/v1")

35

36 class AzureOpenAIClient(OpenAIBaseClient):

37 MODEL = 'gpt-4o'

38

39 def __init__(self, config):

40 super().__init__(config)

41 self._end_point = config.get('endpoint')

42

43 def usable(self):

44 return super().usable() and self._end_point

45

46 def _get_client(self):

47 from openai import AzureOpenAI

48 return AzureOpenAI(azure_endpoint=self._end_point, api_key=self._api_key, api_version="2024-02-01")

49

50 class DoubaoClient(OpenAIBaseClient):

51 BASE_URL = 'https://ark.cn-beijing.volces.com/api/v3'

52 MODEL = 'doubao-seed-1.6-250615'

53

54 class MoonShotClient(OpenAIBaseClient):

55 BASE_URL = T('https://api.moonshot.ai/v1')

56 MODEL = 'kimi-latest'

57

58 class BigModelClient(OpenAIBaseClient):

59 BASE_URL = 'https://open.bigmodel.cn/api/paas/v4'

60 MODEL = 'glm-4.5-air'

61

62 class ZClient(OpenAIBaseClient):

63 BASE_URL = 'https://api.z.ai/api/paas/v4'

64 MODEL = 'glm-4.5-flash'

65

66 CLIENTS = {

67 "openai": OpenAIClient,

68 "ollama": OllamaClient,

69 "claude": ClaudeClient,

70 "gemini": GeminiClient,

71 "deepseek": DeepSeekClient,

72 'grok': GrokClient,

73 'trust': TrustClient,

74 'azure': AzureOpenAIClient,

75 'oauth2': OAuth2Client,

76 'doubao': DoubaoClient,

77 'kimi': MoonShotClient,

78 'bigmodel': BigModelClient,

79 'z': ZClient

80}

81