Coverage for aipyapp/llm/base_openai.py: 29%

1#! /usr/bin/env python3 

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

3 

4from collections import Counter 

5 

6import httpx 

7import openai 

8from loguru import logger 

9 

10from .. import T 

11from . import BaseClient, ChatMessage 

12 

13# https://platform.openai.com/docs/api-reference/chat/create 

14# https://api-docs.deepseek.com/api/create-chat-completion 

15class OpenAIBaseClient(BaseClient): 

16 """ OpenAI compatible client """ 

17 

18 def get_params(self): 

19 params = {'stream_options': {'include_usage': True}} 

20 params.update(super().get_params()) 

21 return params 

22 

23 def usable(self): 

24 return super().usable() and self._api_key 

25 

26 def _get_client(self): 

27 return openai.Client( 

28 api_key=self._api_key, 

29 base_url=self._base_url, 

30 timeout=self._timeout, 

31 http_client=httpx.Client( 

32 verify=self._tls_verify 

33 ) 

34 ) 

35 

36 def add_system_prompt(self, history, system_prompt): 

37 history.add("system", system_prompt) 

38 

39 def _parse_usage(self, usage): 

40 try: 

41 reasoning_tokens = int(usage.completion_tokens_details.reasoning_tokens) 

42 except Exception: 

43 reasoning_tokens = 0 

44 

45 usage = Counter({'total_tokens': usage.total_tokens, 

46 'input_tokens': usage.prompt_tokens, 

47 'output_tokens': usage.completion_tokens + reasoning_tokens}) 

48 return usage 

49 

50 def _parse_stream_response(self, response, stream_processor): 

51 usage = Counter() 

52 with stream_processor as lm: 

53 for chunk in response: 

54 #print(chunk) 

55 if hasattr(chunk, 'usage') and chunk.usage is not None: 

56 usage = self._parse_usage(chunk.usage) 

57 

58 if chunk.choices: 

59 content = None 

60 delta = chunk.choices[0].delta 

61 if delta.content: 

62 reason = False 

63 content = delta.content 

64 elif hasattr(delta, 'reasoning_content') and delta.reasoning_content: 

65 reason = True 

66 content = delta.reasoning_content 

67 if content: 

68 lm.process_chunk(content, reason=reason) 

69 

70 return ChatMessage(role="assistant", content=lm.content, reason=lm.reason, usage=usage) 

71 

72 def _parse_response(self, response): 

73 message = response.choices[0].message 

74 reason = getattr(message, "reasoning_content", None) 

75 return ChatMessage( 

76 role=message.role, 

77 content=message.content, 

78 reason=reason, 

79 usage=self._parse_usage(response.usage) 

80 ) 

81 

82 def get_completion(self, messages): 

83 if not self._client: 

84 self._client = self._get_client() 

85 

86 response = self._client.chat.completions.create( 

87 model = self._model, 

88 messages = messages, 

89 stream=self._stream, 

90 max_tokens = self.max_tokens, 

91 temperature = self._temperature, 

92 **self._params 

93 ) 

94 return response 

95