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

1#! /usr/bin/env python3

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

3

4 from collections import Counter

5

6 import httpx

7 import openai

8 from loguru import logger

9

10 from .. import T

11 from . import BaseClient, ChatMessage

12

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

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

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