Implementation:NVIDIA NeMo Curator OpenAIClient

Knowledge Sources	NVIDIA NeMo Curator
Domains	Machine Learning, LLM Integration, API Clients
Last Updated	2026-02-14 00:00 GMT

Overview

Provides synchronous and asynchronous OpenAI API client implementations for querying language models through the NeMo Curator LLM client interface.

Description

The openai_client module contains two classes that wrap the OpenAI Python SDK:

OpenAIClient extends LLMClient and provides synchronous access to the OpenAI Chat Completions API. On setup(), it instantiates an openai.OpenAI client with any extra keyword arguments passed to the constructor. The query_model method accepts messages, a model name, an optional ConversationFormatter (ignored with a warning), and an optional GenerationConfig (or dict). It calls client.chat.completions.create with the full generation parameters and returns a list of generated text strings. The top_k parameter is not supported by OpenAI and triggers a warning if set.

AsyncOpenAIClient extends AsyncLLMClient and provides the same functionality asynchronously. It adds concurrency control (max_concurrent_requests), retry logic (max_retries), and exponential backoff (base_delay) inherited from AsyncLLMClient. The _query_model_impl method calls the async completions API via await self.client.chat.completions.create.

Both clients support a configurable timeout (default 120 seconds) and pass through all extra kwargs to the underlying OpenAI client constructor, making them compatible with OpenAI-compatible endpoints such as local vLLM inference servers.

Usage

Use OpenAIClient for simple synchronous LLM queries and AsyncOpenAIClient for high-throughput scenarios requiring concurrent requests with built-in rate limiting and retry handling. These clients are used in LLM-based data annotation tasks such as classification, labeling, and quality assessment within the curation pipeline.

Code Reference

Source Location

• Repository: NeMo-Curator
• File: nemo_curator/models/client/openai_client.py
• Lines: 1-139

Signature

class OpenAIClient(LLMClient):
    def __init__(self, **kwargs) -> None: ...
    def setup(self) -> None: ...
    def query_model(
        self,
        *,
        messages: Iterable,
        model: str,
        conversation_formatter: ConversationFormatter | None = None,
        generation_config: GenerationConfig | dict | None = None,
    ) -> list[str]: ...

class AsyncOpenAIClient(AsyncLLMClient):
    def __init__(
        self,
        max_concurrent_requests: int = 5,
        max_retries: int = 3,
        base_delay: float = 1.0,
        **kwargs,
    ) -> None: ...
    def setup(self) -> None: ...
    async def _query_model_impl(
        self,
        *,
        messages: Iterable,
        model: str,
        conversation_formatter: ConversationFormatter | None = None,
        generation_config: GenerationConfig | dict | None = None,
    ) -> list[str]: ...

Import

from nemo_curator.models.client.openai_client import OpenAIClient, AsyncOpenAIClient

I/O Contract

Inputs (OpenAIClient Constructor)

Name	Type	Required	Description
timeout	int	No	Request timeout in seconds (default: 120)
**kwargs	dict	No	Additional keyword arguments passed to the openai.OpenAI constructor (e.g., api_key, base_url)

Inputs (AsyncOpenAIClient Constructor)

Name	Type	Required	Description
max_concurrent_requests	int	No	Maximum number of concurrent requests (default: 5)
max_retries	int	No	Maximum retry attempts for rate-limited requests (default: 3)
base_delay	float	No	Base delay for exponential backoff in seconds (default: 1.0)
timeout	int	No	Request timeout in seconds (default: 120)
**kwargs	dict	No	Additional keyword arguments passed to the openai.AsyncOpenAI constructor

Inputs (query_model / _query_model_impl)

Name	Type	Required	Description
messages	Iterable	Yes	Chat messages in OpenAI format (list of dicts with role and content)
model	str	Yes	Model identifier to query (e.g., "gpt-4", or a local model name)
conversation_formatter	ConversationFormatter or None	No	Not used; triggers a warning if provided
generation_config	GenerationConfig or dict or None	No	Generation parameters (max_tokens, temperature, top_p, etc.). Defaults to GenerationConfig() if None.

Outputs

Name	Type	Description
responses	list[str]	List of generated text strings, one per completion choice

Usage Examples

Basic Synchronous Usage

from nemo_curator.models.client.openai_client import OpenAIClient

client = OpenAIClient(api_key="your-api-key")
client.setup()

messages = [{"role": "user", "content": "Classify this text: ..."}]
responses = client.query_model(messages=messages, model="gpt-4")
print(responses[0])

Async Usage with Custom Config

import asyncio
from nemo_curator.models.client.openai_client import AsyncOpenAIClient
from nemo_curator.models.client.llm_client import GenerationConfig

client = AsyncOpenAIClient(
    max_concurrent_requests=10,
    max_retries=5,
    api_key="your-api-key",
    base_url="http://localhost:8000/v1",  # Local vLLM server
)
client.setup()

config = GenerationConfig(max_tokens=256, temperature=0.7)
messages = [{"role": "user", "content": "Summarize: ..."}]

responses = asyncio.run(
    client.query_model(messages=messages, model="local-model", generation_config=config)
)

Related Pages

• Environment:NVIDIA_NeMo_Curator_Python_Linux_Base