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

1#! /usr/bin/env python 

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

3 

4 

5from .. import T 

6from .base import ChatMessage, BaseClient 

7from .base_openai import OpenAIBaseClient 

8from .client_claude import ClaudeClient 

9from .client_ollama import OllamaClient 

10from .client_oauth2 import OAuth2Client 

11from .models import ModelRegistry, ModelCapability 

12 

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

14 

15class OpenAIClient(OpenAIBaseClient): 

16 MODEL = 'gpt-4o' 

17 

18class GeminiClient(OpenAIBaseClient): 

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

20 MODEL = 'gemini-2.5-flash' 

21 

22class DeepSeekClient(OpenAIBaseClient): 

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

24 MODEL = 'deepseek-chat' 

25 

26class GrokClient(OpenAIBaseClient): 

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

28 MODEL = 'grok-3-mini' 

29 

30class 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 

36class 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 

50class DoubaoClient(OpenAIBaseClient): 

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

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

53 

54class MoonShotClient(OpenAIBaseClient): 

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

56 MODEL = 'kimi-latest' 

57 

58class BigModelClient(OpenAIBaseClient): 

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

60 MODEL = 'glm-4.5-air' 

61 

62class ZClient(OpenAIBaseClient): 

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

64 MODEL = 'glm-4.5-flash' 

65 

66CLIENTS = { 

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