Skip to content

Add Codex-backed dataset generation client #47

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kdunee merged 1 commit into from
May 24, 2026
Merged

Add Codex-backed dataset generation client #47

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 61 additions & 9 deletions ai_research/dataset_generation/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@

from ai_research.dataset_generation.domain.category import get_random_category
from ai_research.dataset_generation.domain.llm_response import parse_examples
from ai_research.dataset_generation.infrastructure.openai_codex_client import (
OpenAICodexTextClient,
)
from ai_research.dataset_generation.infrastructure.openai_client import OpenAITextClient
from ai_research.dataset_generation.infrastructure.prompts import (
get_filtering_prompt,
Expand All @@ -18,28 +21,64 @@

logger = logging.getLogger(__name__)

LOG_FORMAT = "%(asctime)s %(levelname)s %(name)s: %(message)s"
LOG_DATE_FORMAT = "%Y-%m-%dT%H:%M:%S%z"
PROMPT_ARTIFACT_MARKERS = (
"[Assertion]",
"[Code]",
"[Thinking]",
"[Explanation]",
"Example Set",
)
PROMPT_ARTIFACT_MARKERS_LOWER = tuple(
marker.lower() for marker in PROMPT_ARTIFACT_MARKERS
)


def _add_llm_arguments(parser: argparse.ArgumentParser) -> None:
parser.add_argument("--model", default="gpt-4o-mini")
parser.add_argument("--max-tokens", type=int, default=32768)
parser.add_argument("--temperature", type=float, default=1.0)
parser.add_argument("--api-key", required=True)
parser.add_argument("--api-key", default="")
parser.add_argument("--base-url")


def _build_client(args: argparse.Namespace) -> OpenAITextClient:
return OpenAITextClient(
api_key=args.api_key,
def _build_client(args: argparse.Namespace) -> OpenAITextClient | OpenAICodexTextClient:
if args.api_key:
return OpenAITextClient(
api_key=args.api_key,
base_url=args.base_url,
model=args.model,
temperature=args.temperature,
max_tokens=args.max_tokens,
)

return OpenAICodexTextClient(
api_key="",
base_url=args.base_url,
model=args.model,
temperature=args.temperature,
max_tokens=args.max_tokens,
)


def _contains_prompt_artifact(value: object) -> bool:
if isinstance(value, str):
lowered = value.lower()
return any(marker in lowered for marker in PROMPT_ARTIFACT_MARKERS_LOWER)
if isinstance(value, dict):
return any(_contains_prompt_artifact(item) for item in value.values())
if isinstance(value, list):
return any(_contains_prompt_artifact(item) for item in value)
return False


TextClient = OpenAITextClient | OpenAICodexTextClient


def run_generate(
args: argparse.Namespace,
client_factory: Callable[[argparse.Namespace], OpenAITextClient] = _build_client,
client_factory: Callable[[argparse.Namespace], TextClient] = _build_client,
category_picker: Callable[[], str] = get_random_category,
) -> int:
template = get_generation_prompt()
Expand All @@ -59,7 +98,7 @@ def run_generate(

def run_filter(
args: argparse.Namespace,
client_factory: Callable[[argparse.Namespace], OpenAITextClient] = _build_client,
client_factory: Callable[[argparse.Namespace], TextClient] = _build_client,
) -> int:
template = get_filtering_prompt()
client = client_factory(args)
Expand All @@ -84,13 +123,21 @@ def run_filter(

def run_transform(args: argparse.Namespace) -> int:
with Path(args.input).open() as in_file, Path(args.output).open("w") as out_file:
for raw_line in in_file:
for line_number, raw_line in enumerate(in_file, start=1):
line = raw_line.strip()
if not line:
continue

output_line = json.loads(line)
for example in output_line.get("examples", []):
for example_index, example in enumerate(output_line.get("examples", []), start=1):
if _contains_prompt_artifact(example):
logger.warning(
"Rejected mis-parsed generated example at line %s, example %s",
line_number,
example_index,
)
continue

json.dump(example, out_file)
out_file.write("\n")

Expand Down Expand Up @@ -131,5 +178,10 @@ def build_parser() -> argparse.ArgumentParser:
def main(argv: Optional[Sequence[str]] = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
logging.basicConfig(level=logging.INFO)
logging.basicConfig(
level=logging.INFO,
format=LOG_FORMAT,
datefmt=LOG_DATE_FORMAT,
)
logger.info("Starting dataset generation with command: %s", args.command)
return args.func(args)
26 changes: 4 additions & 22 deletions ai_research/dataset_generation/infrastructure/openai_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,27 +64,16 @@ def __init__(
def infer(self, template: str, data: Mapping[str, object]) -> str:
prompt = render_prompt_template(template, data)
delay = self.retry_delay_seconds
use_max_completion_tokens = False

for attempt in range(self.max_retries + 1):
try:
messages: list[ChatCompletionUserMessageParam] = [
{"role": "user", "content": prompt}
]
if use_max_completion_tokens:
response = self._client.chat.completions.create(
model=self.model,
messages=messages,
temperature=self.temperature,
max_completion_tokens=self.max_tokens,
)
else:
response = self._client.chat.completions.create(
model=self.model,
messages=messages,
temperature=self.temperature,
max_tokens=self.max_tokens,
)
response = self._client.chat.completions.create(
model=self.model,
messages=messages,
)
Comment on lines +73 to +76

@coderabbitai coderabbitai Bot May 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🌐 Web query:

For openai==1.109.1, what are the currently supported chat.completions.createparameters forgpt-4o-miniandgpt-5-mini(specificallytemperature, max_tokens, and max_completion_tokens)?

💡 Result:

For openai==1.109.1, the currently supported chat.completions.create parameters for gpt-4o-mini and gpt-5-mini are: gpt-4o-mini (chat.completions.create) - temperature: supported (Optional[float], between 0 and 2). [1] - max_tokens: supported but deprecated (and not compatible with o-series models). [2] - max_completion_tokens: supported (upper bound on tokens generated for a completion, including visible output and reasoning tokens). [2] gpt-5-mini (chat.completions.create) - temperature: not supported (gpt-5 series requires default temperature; user reports indicate temperature not supported/only certain defaults allowed). [3] - max_tokens: not supported (use max_completion_tokens instead). [4] - max_completion_tokens: supported (it is the replacement for max_tokens; request type supports it). [2][4] Notes on what “supported” means here - In the openai Python SDK, chat.completions.create exposes both temperature, max_tokens, and max_completion_tokens as request parameters at the type level; however, model-specific support still applies at runtime (e.g., max_tokens not compatible/not supported for some model families, and temperature not supported for GPT-5 mini per reported behavior). [2]

Citations:

Forward generation parameters to OpenAI chat.completions.create

In ai_research/dataset_generation/infrastructure/openai_client.py (lines 73-76), the request only passes model and messages, so temperature and any token limits aren’t forwarded to OpenAI. That makes generation-related CLI options ineffective and can change output/cost behavior.

Map and pass the correct parameters per model (openai==1.109.1): gpt-4o-mini supports temperature and max_completion_tokens (while max_tokens is deprecated/incompatible for o-series), and gpt-5-mini doesn’t support temperature/max_tokens but does support max_completion_tokens.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ai_research/dataset_generation/infrastructure/openai_client.py` around lines
73 - 76, The chat completion call in openai_client.py currently only passes
model and messages; update the call in the method that uses
self._client.chat.completions.create to forward generation args from the client
(e.g., self.temperature, self.max_completion_tokens, self.max_tokens) and
conditionalize them by model: for "gpt-4o-mini" include temperature (if set) and
max_completion_tokens; for "gpt-5-mini" do not pass temperature or legacy
max_tokens but do pass max_completion_tokens when present; ensure you only pass
parameters supported by openai==1.109.1 to avoid incompatible fields.

return _extract_response_text(response)
except (
RateLimitError,
Expand All @@ -95,13 +84,6 @@ def infer(self, template: str, data: Mapping[str, object]) -> str:
if attempt == self.max_retries:
raise
except APIStatusError as exc:
if (
not use_max_completion_tokens
and _should_retry_with_max_completion_tokens(exc)
):
use_max_completion_tokens = True
continue

status_code = exc.status_code or 0
if status_code not in {408, 409, 429} and status_code < 500:
raise
Expand Down
Loading
Loading