wrap-openai 0.3.0

Wrap custom generate function as an OpenAI SDK compatible API service.

Wrap OpenAI

Wrap a research-oriented custom generate function as an OpenAI Chat Completions compatible API service.

Experimental Package: This package is designed for model research and prototype validation, not production inference serving.

1. Features

• OpenAI Chat Completions request and response format
• Streaming and non-streaming generation
• Raw messages input, including structured multimodal content
• Explicit fixed, OpenAI, and custom parameter groups
• Custom client parameters through OpenAI SDK extra_body
• API Key management, CORS, and health check endpoints

2. Installation

pip install wrap-openai

Install from source:

git clone https://github.com/WKQ9411/wrap-openai.git
cd wrap-openai
uv sync

Install the Qwen demo dependencies:

uv sync --extra qwen

3. Generate Function Contract

The registered generate function must accept the OpenAI messages list as its first positional argument. The parameter name is not enforced, although messages is recommended.

def generate(messages, model, tokenizer, temperature=0.7):
    ...

The messages structure is preserved:

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "What is this?"},
            {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}},
        ],
    },
]

The generate function is responsible for applying the model-specific chat template:

prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)

Return a string when support_stream=False, or yield string chunks when support_stream=True.

4. Register a Generate Function

from wrap_openai import register_generate, run_server


def generate(
    messages,
    model,
    tokenizer,
    temperature=0.7,
    max_tokens=512,
    top_p=0.9,
    top_k=50,
    draft_steps=4,
):
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        temperature=temperature,
        max_new_tokens=max_tokens,
        top_p=top_p,
        top_k=top_k,
        draft_steps=draft_steps,
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)


register_generate(
    generate_func=generate,
    support_stream=False,
    model_id="research-model-v1",
    fixed_kwargs={
        "model": model,
        "tokenizer": tokenizer,
    },
    openai_kwargs={
        "temperature": 0.7,
        "max_tokens": 512,
        "top_p": 0.9,
    },
    custom_kwargs={
        "top_k": 50,
        "draft_steps": 4,
    },
)

run_server(host="0.0.0.0", port=8000)

register_generate accepts three flat keyword dictionaries:

• fixed_kwargs: server-only objects and values. Clients cannot override them.
• openai_kwargs: enabled OpenAI parameters and server defaults. Standard request fields override them.
• custom_kwargs: custom parameters and server defaults. extra_body fields override them.

All three dictionaries are flattened when calling the function:

generate_func(
    messages,
    **fixed_kwargs,
    **effective_openai_kwargs,
    **effective_custom_kwargs,
)

The following OpenAI parameters can be enabled through openai_kwargs:

• temperature
• max_tokens
• top_p
• presence_penalty
• frequency_penalty
• n
• stop
• seed

Non-OpenAI parameters such as top_k and experimental decoding controls belong in custom_kwargs.

Registration validates that:

• the function accepts messages as its first positional argument;
• registered keyword names are accepted by the function or **kwargs;
• the three keyword groups do not overlap;
• OpenAI parameters are placed in openai_kwargs;
• reserved fields are not exposed as custom parameters.

5. OpenAI SDK Client

from openai import OpenAI


client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="sk-dummy",
)

response = client.chat.completions.create(
    model="research-model-v1",
    messages=[{"role": "user", "content": "Hello"}],
    temperature=0.2,
    max_tokens=256,
    extra_body={
        "top_k": 20,
        "draft_steps": 8,
    },
)

print(response.choices[0].message.content)

extra_body values are merged into the JSON request body by the OpenAI SDK. Only fields declared in custom_kwargs are accepted. Unknown custom fields and attempts to override fixed values return HTTP 422.

The request model must match the model_id passed to register_generate.

6. Streaming

A streaming generate function yields string chunks:

def stream_generate(messages, model, tokenizer, temperature=0.7):
    for text_chunk in custom_model_stream(messages, model, tokenizer, temperature):
        yield text_chunk


register_generate(
    generate_func=stream_generate,
    support_stream=True,
    model_id="research-model-v1",
    fixed_kwargs={
        "model": model,
        "tokenizer": tokenizer,
    },
    openai_kwargs={
        "temperature": 0.7,
    },
)

When the client requests stream=False, wrap-openai collects the chunks into one response. When a non-streaming function receives a stream=True request, the complete result is returned in one SSE chunk with a warning chunk.

Client example:

stream = client.chat.completions.create(
    model="research-model-v1",
    messages=[{"role": "user", "content": "Hello"}],
    stream=True,
)

for chunk in stream:
    content = chunk.choices[0].delta.content
    if content:
        print(content, end="", flush=True)

7. Server Configuration

run_server(
    host="0.0.0.0",
    port=8000,
    require_api_key=False,
    allow_remote_api_key_management=False,
    enable_cors=True,
    cors_origins="*",
)

Health check:

GET /health

8. API Key Management

wrap-openai --generate --name "my-key"
wrap-openai --list
wrap-openai --revoke <api_key>

Configure the storage path from Python:

from wrap_openai import set_api_keys_path

set_api_keys_path("/custom/path/to/keys")

9. Examples

• demo/run_server.py: lightweight messages, streaming, and custom parameter example
• demo/server_demo.py: Qwen model deployment example
• demo/run_client.py: OpenAI SDK client examples
• demo/chat_demo.py: CLI chat application
• demo/manage_api_keys.py: API Key management over HTTP

10. License

MIT License