In this post, we explore how to use OpenAI and Nebius APIs to build an interactive response model with structured outputs. The whole code can be found in this Google Colab Notebook. You'll learn the following:

Configure API keys for OpenAI and Nebius using .env files. In Colab, you can set this up directly in the terminal by running:

cat <<EOL > .env
OPENAI_API_KEY="your_API_key"
NEBIUS_API_KEY="your_API_key"
EOL

Next, we create a custom print_wrapper function that formats text, enabling markdown-style formatting for better readability and wrapping long lines. This function provides a cleaner output structure, which is useful when working with detailed responses and reasoning.

from textwrap import TextWrapper
import re

chars = 70

def print_wrapper(print_func):
    """Wraps text to specified width and handles markdown-style and escaped newlines."""
    def function_wrapper(text):
        if not isinstance(text, str):
            text = str(text)

        # Handle **bold** formatting
        text = re.sub(r'\*\*(.*?)\*\*', r'\033[1m\1\033[0m', text)  # ANSI for bold

        # Handle escaped newlines
        text = text.replace("\\n", "\n")
        text = text.replace("\\", "\n")

        # Wrap text
        wrapper = TextWrapper(width=chars)
        wrapped_text = "\n".join([wrapper.fill(line) for line in text.split("\n")])

        return print_func(wrapped_text)

    return function_wrapper

# Override the built-in print function with the wrapped version
print = print_wrapper(print)

# Pretty print function using the wrapped print
def pretty_print_response(response_obj):
    # Title for response
    print("Response:\n" + "-" * chars)
    print(response_obj.response)
    print("\n" + "=" * chars)

    # Title for thought
    print("Thought Process:\n" + "-" * chars)
    print(response_obj.thought)
    print("\n" + "-" * chars)

Using pydantic, we define a structured Response model, containing both an answer and a "Thought Process" section, providing transparent, step-by-step reasoning behind each answer.

class Response(BaseModel):
    response: str
    thought: str = Field(
        description="Provide and answer in a structered step-by-step reasoning."
    )

The run_model function allows for flexible interactions with either OpenAI or Nebius models by specifying the provider and model name. This makes it simple to test multiple models and configurations for your questions.

system_prompt = """You are an assistant that answers questions with clear reasoning. 
For each question, provide both a response and a detailed reasoning process that 
explains how you reached the answer."""

question = "What is the impact of reinforcement learning in robotics?"


def run_model(system_prompt: str, question:str, response_model: BaseModel, model_name:str, provider:str):
    if provider == "openai":
        llm = instructor.from_openai(OpenAI())

    elif provider == "nebius":
        llm = instructor.from_openai(OpenAI(
            base_url="https://api.studio.nebius.ai/v1/",
            api_key=os.environ.get("NEBIUS_API_KEY"),
        ))
    response = llm.chat.completions.create(
    model=model_name,
    response_model=Response,
    messages=[
        {"role": "system", "content": system_prompt},
        {
            "role": "user",
            "content": question,
        },
    ],
    temperature=0.1,
    top_p=0.8,
    frequency_penalty=0.6,
    )
    return response

The pretty_print_response function displays responses with distinct sections for the main answer and the thought process, ensuring that complex answers remain easy to follow.

%%time
response = run_model(system_prompt, question, Response, "meta-llama/Meta-Llama-3.1-405B-Instruct", "nebius")
pretty_print_response(response)

Test different models by specifying different model names in run_model, like "meta-llama/Meta-Llama-3.1-405B-Instruct" or "gpt-4o" to see how they differ in handling the prompt.

%%time
response = run_model(system_prompt, question, Response, "gpt-4o", "openai")
pretty_print_response(response)