Extending the API Payload Schema

Relevant source files

The following files were used as context for generating this wiki page:

• src/handler.py
• src/utils.py
• src/engine.py
• test_inputs/openai_completion.json
• test_inputs/openai_get_models.json
• test_inputs/default_runpodio.json

The RunPod Worker Ollama project provides a flexible mechanism for handling API requests by mapping incoming job payloads to internal processing logic. Extending the API payload schema involves modifying the way inputs are parsed and how they are routed to the underlying Ollama or OpenAI-compatible engines. This architecture allows the worker to support both simple prompt-based requests and complex OpenAI-formatted payloads.

The system relies on a centralized JobInput class to abstract the raw JSON input received from RunPod. This abstraction layer enables developers to add new fields or modify existing routing logic without significantly altering the core execution flow within the handler.

Architecture and Input Mapping

The input handling architecture is designed to distinguish between standard RunPod requests and those specifically targeting OpenAI-compatible endpoints. The handler.py script receives the initial payload and utilizes utils.JobInput to sanitize and categorize the data.

The JobInput Class

The JobInput class is the primary point for extending the schema. It extracts specific keys from the input object provided by the RunPod job.

Field	Source Key	Description
llm_input	messages or prompt	The primary text or chat history for the model.
stream	stream	Boolean flag indicating if the response should be streamed. Defaults to False.
openai_route	openai_route	The specific API endpoint targeted (e.g., /v1/chat/completions).
openai_input	openai_input	A nested dictionary containing OpenAI-specific parameters like model or max_tokens.

Sources: src/utils.py:1-7, src/handler.py:10-15

Data Flow for Payload Processing

The following diagram illustrates how a raw JSON payload is processed and routed to the appropriate engine.

The logic determines the engine based on the presence of openai_route. If present, the OllamaOpenAiEngine handles the request; otherwise, the OllamaEngine performs a mapping of simple prompt inputs into a structured OpenAI format.
Sources: src/handler.py:12-15, src/engine.py:22-48

Modifying Engine Logic for New Fields

When extending the payload schema, the OllamaEngine and OllamaOpenAiEngine in src/engine.py must be updated to handle the new data members of JobInput.

Internal Route Mapping

In OllamaEngine, the generate method creates a internal "synthetic" JobInput if the user provides a raw string prompt instead of a full OpenAI structure. This is where default behaviors for simplified payloads are defined.

if isinstance(job_input.llm_input, str):
    openAiJob = JobInput({
        "openai_route": "/v1/completions",
        "openai_input": {
            "model": model,
            "prompt": job_input.llm_input,
            "stream": job_input.stream
        }
    })

Sources: src/engine.py:25-34

OpenAI Compatibility Layer

The OllamaOpenAiEngine passes the openai_input dictionary directly to the OpenAI client using kwargs (**openai_input). This means any field added to the openai_input section of the JSON payload that is supported by the OpenAI SDK (and subsequently Ollama's OpenAI-compatible API) will be automatically passed through.

async def _handle_chat_or_completion_request(self, openai_input, chat=False):
    if chat:
        response = client.chat.completions.create(**openai_input)
    else:
        response = client.completions.create(**openai_input)

Sources: src/engine.py:82-87

Payload Examples and Testing

The project includes several test input files that demonstrate the current schema. These can be used as templates for extension.

Standard OpenAI Completion Schema

This schema uses the openai_route and openai_input keys to provide fine-grained control over the request.

{
    "input": {
        "openai_route": "/v1/completions",
        "openai_input": {
            "model": "llama3.2:1b",
            "prompt": "How are you?",
            "max_tokens": 60
        }
    }
}

Sources: test_inputs/openai_completion.json:1-10

Simplified Prompt Schema

A minimal schema that only requires a prompt key. The worker automatically wraps this in a completion request.

{
    "input": {
        "prompt": "How are you?"
    }
}

Sources: test_inputs/default_runpodio.json:1-5

Conclusion

Extending the API Payload Schema in this worker is a two-step process: defining the new field in the JobInput class in src/utils.py to ensure it is captured from the RunPod environment, and then updating the generation logic in src/engine.py to utilize that field when calling the Ollama/OpenAI backend. The use of Python's keyword argument unpacking (**openai_input) makes the system natively extensible for any standard OpenAI parameters.