Serializing OpenAI API Responses for Logging with Pydantic

TL;DR

from openai import OpenAI
from openai.types import CreateEmbeddingResponse

client: OpenAI = OpenAI(api_key="YOUR_OPENAI_KEY")

embeddings_response: CreateEmbeddingResponse = (
  client.embeddings.create(input=["Some embeddings to generate"])
)

# Using Pydantic's `model_dump_json` method,
# we can easily serialize any OpenAI responses
serialized_response: str = embeddings_response.model_dump_json()

# Now that our response is serialized,
# we can log and store it as needed
openai_response: OpenAiResponse = OpenAiResponse(
  embeddings_response=serialized_response,
  ...  # other model fields
)

Why you might consider serializing responses

When working with external APIs, I lean heavily on logging as much information as I can for requests and responses. Recently, I've been working on a side project that relies on OpenAI's embeddings and chat completions APIs, as well as Amazon Textract's OCR service. For these types of external APIs, logging entire responses can be incredibly useful for many reasons, whether it's trying to iron out pricing for a product or needing to track if you're receiving a non-deterministic output (despite configuring for reproducibility). Over time, the quality of outputs can drift over time, and you'll inevitably need to iterate on input prompts as your product evolves or parameters change, which is where detailed logging provides the opportunity to look at how cost and quality may be changing - almost like a passive A/B test.

For most small to medium-sized projects, I'll typically store these logs as a dedicated table or explicit metadata columns on an existing table using Postgres. You could break a response object and its items into dedicated columns, but APIs tend to change over time (sometimes without warning), and I don't want to spend time maintaining something that will be infrequently used for debugging, replaying requests, or analysis.

So, as part of this side project, I was attempting to store the response directly into an SQLAlchemy model like so:

from openai import OpenAI
from openai.types import CreateEmbeddingResponse

client: OpenAI = OpenAI(api_key="YOUR_OPENAI_KEY")

embeddings_response: CreateEmbeddingResponse = (
  client.embeddings.create(input=["Some embeddings to generate"])
)

# This will throw - we can't store the object directly!
openai_response: OpenAiResponse = OpenAiResponse(
  embeddings_response=embeddings_response,  # embeddings_response is a JSONB column
)

This, of course, would not work. Trying to insert this directly into the model would produce an error like TypeError: Object of type CreateEmbeddingResponse is not JSON serializable. While exploring solutions, I came across a gist that created an explicit serialize_completion function to accomplish this. While the approach is straightforward (and unblocked me temporarily while trying to solve an unrelated task!), it’s more work than necessary and would require more maintenance and a similar manual approach for other response types.

When in doubt, just read the source code

After creating another serializer for the embeddings response, I was frustrated. What if I miss a key or the API changes silently, meaning I may be losing information or have exceptions creeping up? There are many solutions, but I decided I'd jump into the definition for the CreateEmbeddingResponse type. Luckily, OpenAI represents all of their responses as Pydantic models. Pydantic is an extremely powerful data validation library that lets us solve this serialization issue natively, instead of having to come up with our own solution:

# Using Pydantic's `model_dump_json` method,
# we can easily serialize any OpenAI responses
serialized_response: str = embeddings_response.model_dump_json()

Reviewing the documentation for pydantic.BaseModel, we can find the model_dump_json() method that allows us to serialize it directly to a JSON-encoded string. A simple and obvious solution, but a source of frustration if you don’t know (or forget) where to look. Another personal reminder of why enforcing strong typing is so useful!

A closing note: I’m currently building an ETL product based around documents and document-heavy workflows. If you or your organization might be interested, please let me know!

This post was originally published at romansorin.com