Coverage for haystack/token_counters/openai_counter.py: 98%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import os 

6from typing import Any 

7 

8from openai import OpenAI 

9 

10from haystack.components.generators.chat.openai_responses import _convert_chat_message_to_responses_api_format 

11from haystack.core.serialization import default_to_dict 

12from haystack.dataclasses import ChatMessage 

13from haystack.token_counters.types import TokenCounter 

14from haystack.tools import ToolsType, flatten_tools_or_toolsets 

15from haystack.utils import Secret 

16from haystack.utils.http_client import init_http_client 

17 

18 

19class OpenAITokenCounter(TokenCounter): 

20 """ 

21 Counts tokens with OpenAI's input token counting API. 

22 

23 Unlike local token counters, this counter sends the input to OpenAI's 

24 `POST /v1/responses/input_tokens` endpoint. The returned count includes the model-specific formatting used for 

25 messages and tool schemas, as well as supported non-text content such as images and files. 

26 

27 ## Usage Example: 

28 ```python 

29 from haystack.dataclasses import ChatMessage 

30 from haystack.token_counters import OpenAITokenCounter 

31 

32 counter = OpenAITokenCounter("gpt-5-mini") 

33 messages = [ChatMessage.from_user("Hello, how are you?")] 

34 token_count = counter.count(messages) 

35 print(f"Token count: {token_count}") 

36 ``` 

37 """ 

38 

39 def __init__( 

40 self, 

41 model: str, 

42 *, 

43 api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"), 

44 api_base_url: str | None = None, 

45 organization: str | None = None, 

46 timeout: float | None = None, 

47 max_retries: int | None = None, 

48 http_client_kwargs: dict[str, Any] | None = None, 

49 ) -> None: 

50 """ 

51 Initialize the counter. 

52 

53 :param model: The model whose tokenization should be used. 

54 :param api_key: The OpenAI API key. You can set it with the `OPENAI_API_KEY` environment variable or pass it 

55 explicitly. 

56 :param api_base_url: An optional base URL for the OpenAI API. 

57 :param organization: Your OpenAI organization ID. 

58 :param timeout: Timeout for OpenAI client calls. If unset, uses `OPENAI_TIMEOUT` or 30 seconds. 

59 :param max_retries: Maximum retries for OpenAI client calls. If unset, uses `OPENAI_MAX_RETRIES` or 5. 

60 :param http_client_kwargs: Keyword arguments used to configure the underlying HTTPX client. 

61 """ 

62 self.api_key = api_key 

63 self.model = model 

64 self.api_base_url = api_base_url 

65 self.organization = organization 

66 self.timeout = timeout 

67 self.max_retries = max_retries 

68 self.http_client_kwargs = http_client_kwargs 

69 

70 self.client: OpenAI | None = None 

71 

72 def warm_up(self) -> None: 

73 """Initialize the OpenAI client.""" 

74 if self.client is not None: 

75 return 

76 

77 timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0")) 

78 max_retries = ( 

79 self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5")) 

80 ) 

81 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. 

82 # https://github.com/openai/openai-python/blob/main/httpx2.md 

83 http_client = init_http_client(self.http_client_kwargs, async_client=False) 

84 self.client = OpenAI( 

85 api_key=self.api_key.resolve_value(), 

86 organization=self.organization, 

87 base_url=self.api_base_url, 

88 timeout=timeout, 

89 max_retries=max_retries, 

90 http_client=http_client, # type: ignore[arg-type] 

91 ) 

92 

93 def count(self, messages: list[ChatMessage], tools: ToolsType | None = None) -> int: 

94 """ 

95 Return the exact number of input tokens OpenAI will use for the given messages and tools. 

96 

97 :param messages: The messages to measure. 

98 :param tools: Tools whose schemas are sent alongside the messages, and so consume tokens too. 

99 :returns: The token count, or `0` when there is nothing to measure. 

100 """ 

101 if not messages and not tools: 

102 return 0 

103 

104 self.warm_up() 

105 client = self.client 

106 if client is None: 

107 raise RuntimeError("The OpenAI client was not initialized.") 

108 

109 openai_input: list[dict[str, Any]] = [] 

110 for message in messages: 

111 openai_input.extend(_convert_chat_message_to_responses_api_format(message=message)) 

112 

113 request: dict[str, Any] = {"model": self.model, "input": openai_input} 

114 if tools: 

115 request["tools"] = [ 

116 {"type": "function", **tool.tool_spec} for tool in flatten_tools_or_toolsets(tools=tools) 

117 ] 

118 

119 response = client.responses.input_tokens.count(**request) 

120 return response.input_tokens 

121 

122 def close(self) -> None: 

123 """Close the OpenAI client and its underlying HTTP resources.""" 

124 if self.client is not None: 

125 self.client.close() 

126 self.client = None 

127 

128 def to_dict(self) -> dict[str, Any]: 

129 """ 

130 Serialize the counter. 

131 

132 :returns: A dictionary representation of the counter. 

133 """ 

134 return default_to_dict( 

135 self, 

136 api_key=self.api_key, 

137 model=self.model, 

138 api_base_url=self.api_base_url, 

139 organization=self.organization, 

140 timeout=self.timeout, 

141 max_retries=self.max_retries, 

142 http_client_kwargs=self.http_client_kwargs, 

143 )