| |
| import logging |
| from copy import deepcopy |
| from openai import AzureOpenAI, DefaultHttpxClient |
| from openai.types.responses import Response |
| from openai.types.chat import ChatCompletion |
| from collections import defaultdict |
| from typing import Any, Generator, List, Dict |
| from .base_api import BaseAPI, ProgressCallback |
| from ..core import APICallSample, ToolCall |
|
|
| _logger = logging.getLogger(f"sr_harness.{__name__}") |
|
|
|
|
| @BaseAPI.register("openai") |
| class OpenAIAPI(BaseAPI): |
| """Azure OpenAI adapter for the gpt-4o-mini and gpt-5-mini workflows.""" |
|
|
| def __init__(self, model: str = "gpt-5-mini", **kwargs: Any) -> None: |
| super().__init__(model=model, **kwargs) |
| self.dummy_message = 'Please proofread the message above and say "I have received it."' |
|
|
| def cancel(self) -> str: |
| """Report that the current synchronous SDK request cannot be preempted.""" |
| return ( |
| "OpenAIAPI does not support force-interrupting an active non-streaming request. " |
| "The Agent will pause at its next safe boundary." |
| ) |
|
|
| def format_response( |
| self, message: Any, tool_calls: list[ToolCall], content: str | None = None |
| ) -> APICallSample: |
| """Convert one OpenAI or Azure message into one API-call sample.""" |
| if hasattr(message, "to_dict"): |
| message = message.to_dict() |
| elif hasattr(message, "model_dump"): |
| message = message.model_dump() |
| elif not isinstance(message, dict): |
| raise TypeError(f"Unsupported OpenAI message type: {type(message).__name__}") |
| return self._build_response(message, tool_calls, content) |
|
|
| @property |
| def model_list(self) -> list[str]: |
| """Return the OpenAI and Azure models explicitly supported by this adapter.""" |
| return ["gpt-4o-mini", "gpt-5-mini"] |
|
|
| def _request_samples( |
| self, |
| messages: List[Dict[str, str]], |
| n=1, |
| max_tokens=None, |
| temperature=1.0, |
| top_p=1.0, |
| use_chat_completions=False, |
| callback: ProgressCallback | None = None, |
| ) -> Generator[APICallSample, None, dict[str, Any]]: |
| """ |
| Keyword Args: |
| model (str): The model name to use. Default is 'gpt-4o-mini'. |
| max_tokens (int): The maximum number of tokens to generate. Default is 1024. |
| temperature (float): Sampling temperature. Default is 1.0. |
| top_p (float): Nucleus sampling probability. Default is 1.0. |
| n (int): Number of completions to generate. Default is 1. |
| """ |
| if self.model == "gpt-4o-mini" and not use_chat_completions: |
| use_chat_completions = True |
| _logger.warning( |
| "Turn to use create_chat_completions for gpt-4o-mini model to save token cost." |
| ) |
|
|
| if use_chat_completions: |
| results = yield from self.create_chat_completions( |
| messages, |
| n=n, |
| max_tokens=max_tokens, |
| temperature=temperature, |
| top_p=top_p, |
| callback=callback, |
| ) |
| return results |
| else: |
| results = yield from self.create_responses( |
| messages, |
| n=n, |
| max_tokens=max_tokens, |
| temperature=temperature, |
| top_p=top_p, |
| callback=callback, |
| ) |
| return results |
|
|
| def build_native_tool_description(self, use_chat_completions: bool = False) -> List[Dict]: |
| """Build native tool description. |
| |
| Args: |
| use_chat_completions: The use chat completions value. |
| |
| Returns: |
| List[Dict]: The operation result. |
| """ |
| tools = self.tool_box.format_tools() |
| if not isinstance(tools, list): |
| raise TypeError("Native tool parser must return JSON tool descriptions") |
| if use_chat_completions: |
| return tools |
| return [ |
| { |
| "type": "function", |
| "name": tool["function"]["name"], |
| "description": tool["function"].get("description", ""), |
| "parameters": tool["function"].get( |
| "parameters", {"type": "object", "properties": {}} |
| ), |
| } |
| for tool in tools |
| ] |
|
|
| def create_responses( |
| self, |
| messages: List[Dict[str, str]], |
| n: int = 1, |
| max_tokens: int | None = None, |
| temperature: float = 1.0, |
| top_p: float = 1.0, |
| callback: ProgressCallback | None = None, |
| ) -> Generator[APICallSample, None, dict[str, Any]]: |
| """Create responses. |
| |
| Args: |
| messages: Conversation messages in provider-compatible order. |
| n: The n value. |
| max_tokens: The max tokens value. |
| temperature: The temperature value. |
| top_p: The top p value. |
| |
| Returns: |
| Generator yielding call samples and returning aggregate metadata. |
| """ |
| |
| yield from [] |
| client = AzureOpenAI( |
| api_version=self.get_env("OPENAI_API_VERSION", required=True), |
| azure_endpoint=(endpoint := self.get_env("OPENAI_ENDPOINT", required=True)), |
| api_key=self.get_env("OPENAI_API_KEY", required=True), |
| http_client=DefaultHttpxClient( |
| proxy=self.proxy_url(endpoint), |
| trust_env=False, |
| ), |
| ) |
| payload = { |
| "input": deepcopy(messages), |
| "reasoning": {"summary": "auto"}, |
| "model": self.model, |
| "top_p": top_p, |
| **({"max_output_tokens": max_tokens} if max_tokens is not None else {}), |
| "temperature": temperature, |
| |
| |
| } |
| if self.tool_box is None: |
| pass |
| elif not self.tool_box.native_tool_calls: |
| description = ( |
| "Use the following tools when a tool call is needed. " |
| "Return tool calls in the specified format.\n\n" |
| f"{self.tool_box.format_tools()}" |
| ) |
| if payload["input"][0]["role"] in {"system", "developer"}: |
| first = payload["input"][0] |
| payload["input"][0] = { |
| **first, |
| "content": f"{first['content']}\n\n{description}", |
| } |
| else: |
| payload["input"].insert(0, {"role": "system", "content": description}) |
| else: |
| payload["tools"] = self.build_native_tool_description() |
| payload["tool_choice"] = "auto" |
| if payload["input"][0]["role"] == "system": |
| |
| |
| |
| payload["input"][0]["role"] = "developer" |
|
|
| if n == 1: |
| self.report_progress(callback, {"type": "start", "sample": 1}) |
| try: |
| response = client.responses.create(**payload) |
| response_dict = response.to_dict() |
| usage = self.parse_usage(response) |
| message_object = response.choices[0].message |
| message = ( |
| message_object.to_dict() |
| if hasattr(message_object, "to_dict") |
| else message_object.model_dump() |
| ) |
| content = response.output_text |
| tool_call = self.parse_tool_calls(message, content) |
| sample_response = self.format_response(message, tool_call, content) |
| if callback is not None: |
| self.report_progress(callback, { |
| "type": "complete", |
| "sample": 1, |
| "content": sample_response.content, |
| "reasoning": sample_response.reasoning, |
| "tool_calls": [call.raw for call in sample_response.tool_calls], |
| }) |
| yield sample_response |
| results = { |
| "ok": True, |
| "error": None, |
| "token_usage": usage.get("token", {}), |
| "money_usage": usage.get("price", {}), |
| "extra": {"responses": [response_dict]}, |
| } |
| except InterruptedError: |
| raise |
| except Exception as e: |
| public_error = self.error_message(e) |
| self.report_progress(callback, { |
| "type": "failed", |
| "sample": 1, |
| "error": public_error, |
| }) |
| return { |
| "ok": False, |
| "error": public_error, |
| "token_usage": {}, |
| "money_usage": {}, |
| "extra": {}, |
| } |
| return results |
| else: |
| self.report_progress(callback, {"type": "start", "sample": 1}) |
| try: |
| responses = [] |
| details = [] |
| |
| assert payload["input"][-1]["role"] == "user", ( |
| "When n > 1, the last message must be from user." |
| ) |
| last_message = payload["input"].pop(-1) |
| payload["input"].append({"role": "user", "content": self.dummy_message}) |
| response = client.responses.create(**payload) |
| usage = self.parse_usage(response) |
| responses.append(response) |
| |
| child_payload = deepcopy(payload) |
| child_payload["previous_response_id"] = response.id |
| child_payload["input"] = [last_message] |
| for idx in range(1, n + 1): |
| child_response = client.responses.create(**child_payload) |
| responses.append(child_response) |
| child_response_dict = child_response.to_dict() |
| message = child_response.choices[0].message.to_dict() |
| content = message["content"] |
| tool_call = self.parse_tool_calls(message, content) |
| sample_response = self.format_response(message, tool_call, content) |
| if callback is not None: |
| self.report_progress(callback, { |
| "type": "complete", |
| "sample": idx, |
| "content": sample_response.content, |
| "reasoning": sample_response.reasoning, |
| "tool_calls": [call.raw for call in sample_response.tool_calls], |
| }) |
| if idx < n: |
| self.report_progress(callback, { |
| "type": "start", |
| "sample": idx + 1, |
| }) |
| yield sample_response |
| new_usage = self.parse_usage(child_response) |
| for k, v in new_usage["token"].items(): |
| usage["token"][k] = usage["token"].get(k, 0) + v |
| for k, v in new_usage["price"].items(): |
| usage["price"][k] = usage["price"].get(k, 0) + v |
| details.append( |
| { |
| "content": content, |
| "tool_call": tool_call, |
| "token_usage": dict(new_usage.get("token", {})), |
| "price_usage": dict(new_usage.get("price", {})), |
| "response": child_response_dict, |
| } |
| ) |
| results = { |
| "ok": True, |
| "error": None, |
| "token_usage": usage.get("token", {}), |
| "money_usage": usage.get("price", {}), |
| "extra": {"responses": [resp.to_dict() for resp in responses]}, |
| } |
| except InterruptedError: |
| raise |
| except Exception as e: |
| public_error = self.error_message(e) |
| self.report_progress(callback, { |
| "type": "failed", |
| "sample": min(len(details) + 1, n), |
| "error": public_error, |
| }) |
| return { |
| "ok": False, |
| "error": public_error, |
| "token_usage": {}, |
| "money_usage": {}, |
| "extra": {}, |
| } |
| return results |
|
|
| def create_chat_completions( |
| self, |
| messages: List[Dict[str, str]], |
| n: int = 1, |
| max_tokens: int | None = None, |
| temperature: float = 1.0, |
| top_p: float = 1.0, |
| callback: ProgressCallback | None = None, |
| ) -> Generator[APICallSample, None, dict[str, Any]]: |
| """Create chat completions. |
| |
| Args: |
| messages: Conversation messages in provider-compatible order. |
| n: The n value. |
| max_tokens: The max tokens value. |
| temperature: The temperature value. |
| top_p: The top p value. |
| |
| Returns: |
| Generator yielding call samples and returning aggregate metadata. |
| """ |
| |
| yield from [] |
| client = AzureOpenAI( |
| api_version=self.get_env("OPENAI_API_VERSION", required=True), |
| azure_endpoint=(endpoint := self.get_env("OPENAI_OLDTIME_ENDPOINT", required=True)), |
| api_key=self.get_env("OPENAI_API_KEY", required=True), |
| http_client=DefaultHttpxClient( |
| proxy=self.proxy_url(endpoint), |
| trust_env=False, |
| ), |
| ) |
| payload = { |
| "messages": deepcopy(messages), |
| "n": n, |
| "model": (model := self.model), |
| "top_p": top_p, |
| **({"max_tokens": max_tokens} if max_tokens is not None else {}), |
| "temperature": temperature, |
| } |
| if self.tool_box is None: |
| pass |
| elif not self.tool_box.native_tool_calls: |
| description = ( |
| "Use the following tools when a tool call is needed. " |
| "Return tool calls in the specified format.\n\n" |
| f"{self.tool_box.format_tools()}" |
| ) |
| if payload["messages"][0]["role"] in {"system", "developer"}: |
| first = payload["messages"][0] |
| payload["messages"][0] = { |
| **first, |
| "content": f"{first['content']}\n\n{description}", |
| } |
| else: |
| payload["messages"].insert(0, {"role": "system", "content": description}) |
| else: |
| tools = self.tool_box.format_tools() |
| if not isinstance(tools, list): |
| raise TypeError("Native tool parser must return JSON tool descriptions") |
| payload["tools"] = tools |
| payload["tool_choice"] = "auto" |
| assert model == "gpt-4o-mini", "Only gpt-4o-mini model is supported in chat_completions." |
|
|
| self.report_progress(callback, {"type": "start", "sample": 1}) |
| try: |
| response = client.chat.completions.create(**payload) |
| details = [] |
| for idx, choice in enumerate(response.choices, 1): |
| response_dict = {"choices": [{"message": choice.message.to_dict()}]} |
| message = choice.message.to_dict() |
| content = message["content"] |
| truncated = getattr(choice, "finish_reason", None) == "length" |
| if truncated: |
| message.pop("tool_calls", None) |
| tool_call = [] if truncated else self.parse_tool_calls(message, content) |
| details.append( |
| { |
| "content": content, |
| "tool_call": tool_call, |
| "token_usage": {}, |
| "price_usage": {}, |
| "response": response_dict, |
| } |
| ) |
| sample_response = self.format_response(message, tool_call, content) |
| sample_response.finish_reason = getattr(choice, "finish_reason", None) |
| if callback is not None: |
| self.report_progress(callback, { |
| "type": "complete", |
| "sample": idx, |
| "content": sample_response.content, |
| "reasoning": sample_response.reasoning, |
| "tool_calls": [call.raw for call in sample_response.tool_calls], |
| }) |
| if idx < len(response.choices): |
| self.report_progress(callback, { |
| "type": "start", |
| "sample": idx + 1, |
| }) |
| yield sample_response |
| usage = self.parse_chat_completions_usage(response) |
| for detail in details: |
| detail["token_usage"] = dict(usage.get("token", {})) |
| detail["price_usage"] = dict(usage.get("price", {})) |
| results = { |
| "ok": True, |
| "error": None, |
| "token_usage": usage.get("token", {}), |
| "money_usage": usage.get("price", {}), |
| "extra": {"responses": [response.to_dict()]}, |
| } |
| except InterruptedError: |
| raise |
| except Exception as e: |
| public_error = self.error_message(e) |
| self.report_progress(callback, { |
| "type": "failed", |
| "sample": 1, |
| "error": public_error, |
| }) |
| return { |
| "ok": False, |
| "error": public_error, |
| "token_usage": {}, |
| "money_usage": {}, |
| "extra": {}, |
| } |
| return results |
|
|
| def parse_usage(self, response: Response) -> Dict: |
| """Parse usage. |
| |
| Args: |
| response: Provider response object. |
| |
| Returns: |
| Dict: The operation result. |
| """ |
| usage = {"token": defaultdict(float), "price": defaultdict(float)} |
| if response.model.startswith("gpt-5-mini"): |
| usage["token"]["cached"] = ( |
| cached_tokens := response.usage.input_tokens_details.cached_tokens |
| ) |
| usage["token"]["prompt"] = ( |
| prompt_tokens := response.usage.input_tokens - cached_tokens |
| ) |
| usage["token"]["reason"] = ( |
| reason_tokens := response.usage.output_tokens_details.reasoning_tokens |
| ) |
| usage["token"]["answer"] = ( |
| answer_tokens := response.usage.output_tokens - reason_tokens |
| ) |
| if ( |
| other := response.usage.total_tokens |
| - cached_tokens |
| - prompt_tokens |
| - reason_tokens |
| - answer_tokens |
| ) != 0: |
| usage["token"]["other"] = other |
| usage["price"]["cached"] = 0.025 * cached_tokens / 1e6 |
| usage["price"]["prompt"] = 0.25 * prompt_tokens / 1e6 |
| usage["price"]["reason"] = 2.0 * reason_tokens / 1e6 |
| usage["price"]["answer"] = 2.0 * answer_tokens / 1e6 |
| elif response.model.startswith("gpt-4o-mini"): |
| cached_tokens = 0 |
| usage["token"]["prompt"] += (prompt_tokens := response.usage.prompt_tokens) |
| usage["token"]["reason"] += ( |
| reason_tokens := response.usage.completion_tokens_details.reasoning_tokens |
| ) |
| usage["token"]["answer"] += ( |
| answer_tokens := response.usage.completion_tokens - reason_tokens |
| ) |
| if ( |
| other := response.usage.total_tokens |
| - cached_tokens |
| - prompt_tokens |
| - reason_tokens |
| - answer_tokens |
| ) != 0: |
| usage["token"]["other"] = other |
| usage["price"]["prompt"] += 0.15 * prompt_tokens / 1e6 |
| usage["price"]["reason"] += 0.60 * reason_tokens / 1e6 |
| usage["price"]["answer"] += 0.60 * answer_tokens / 1e6 |
| else: |
| raise NotImplementedError( |
| f"Usage parsing for model {response.model} is not implemented." |
| ) |
| return usage |
|
|
| def parse_chat_completions_usage(self, response: ChatCompletion) -> Dict: |
| """Parse chat completions usage. |
| |
| Args: |
| response: Provider response object. |
| |
| Returns: |
| Dict: The operation result. |
| """ |
| usage = {"token": defaultdict(float), "price": defaultdict(float)} |
| if response.model.startswith("gpt-4o-mini"): |
| usage["token"]["prompt"] += (prompt_tokens := response.usage.prompt_tokens) |
| usage["token"]["reason"] += ( |
| reason_tokens := response.usage.completion_tokens_details.reasoning_tokens |
| ) |
| usage["token"]["answer"] += ( |
| answer_tokens := response.usage.completion_tokens - reason_tokens |
| ) |
| if ( |
| other := response.usage.total_tokens - prompt_tokens - reason_tokens - answer_tokens |
| ) > 0: |
| usage["token"]["other"] = other |
| usage["price"]["prompt"] += 0.15 * prompt_tokens / 1e6 |
| usage["price"]["reason"] += 0.60 * reason_tokens / 1e6 |
| usage["price"]["answer"] += 0.60 * answer_tokens / 1e6 |
| else: |
| raise NotImplementedError( |
| f"Usage parsing for model {response.model} is not implemented." |
| ) |
| return usage |
|
|
