Proposal: Messages Middleware

[jxnl]

Design Proposal

An impl exists at #551

Overview

This proposal outlines the design of a middleware system for enhancing and customizing interactions with language models. The middleware system allows developers to modify the messages sent to the language model before they are processed, enabling custom preprocessing, context addition, and even simple retrieval-augmented generation (RAG) techniques.

Objectives

• Provide a flexible and extensible middleware framework for modifying messages (in the future.,response, models, kwargs, etc) sent to language models
• Enable developers to perform custom preprocessing, add context, and implement RAG techniques while preserving the create() api
• Support both simple function-based middleware and more complex stateful middleware
• Facilitate easy registration and chaining of multiple middleware components
• Allow middleware to be defined using Python classes with Pydantic for validation and serialization

System Architecture

The middleware system will consist of the following components:

1. Middleware Interface: An abstract base class MessageMiddleware that defines the structure and methods required for middleware components. This will include the __call__ method for processing messages.

2. Function-based Middleware: A decorator @messages_middleware that allows developers to define simple middleware as functions. These functions will take the list of messages as input, make any desired changes, and return the modified list of messages.

3. Stateful Middleware: Middleware defined as Python classes that inherit from MessageMiddleware and BaseModel from Pydantic. This allows for stateful middleware with input validation and serialization.

4. Middleware Registration: The with_middleware() method on the language model client for registering middleware components. This allows chaining multiple middleware together.

5. Message Processing Pipeline: The core component that manages the flow of messages through the registered middleware components. It ensures that each middleware is called in the correct order and the modified messages are passed to the language model for processing.

Middleware Examples

Function-based Middleware

Function-based middleware can be defined using the @messages_middleware decorator:

@instructor.messages_middleware
def dumb_rag(messages):
    # TODO: use RAG to generate a response
    # TODO: add the response to the messages
    return messages + [
        {
            "role": "user",
            "content": "Search retrieved: 'Jason is 20 years old'",
        }
    ]

Stateful Middleware

Stateful middleware can be defined as classes that inherit from MessageMiddleware and BaseModel:

class PrintMessages(instructor.MessageMiddleware, BaseModel):
    log: bool = False

    def __call__(
        self, messages: List[ChatCompletionMessageParam]
    ) -> List[ChatCompletionMessageParam]:
        if self.log:
            import pprint
            pprint.pprint({"messages": messages})
        return messages

Middleware Registration and Usage

Middleware components are registered with the language model client using the with_middleware() method, allowing chaining of multiple middleware:

client = (
    instructor.from_openai(openai.OpenAI())
    .with_middleware(dumb_rag)
    .with_middleware(PrintMessages(log=True))
)

Example Usage

Here's an example of using the middleware system to perform a simple RAG technique and extract structured data from the response:

class User(BaseModel):
    age: int
    name: str

client = (
    instructor.from_openai(openai.OpenAI())
    .with_middleware(dumb_rag)
    .with_middleware(PrintLastUserMessage(log=True))
)

user = client.chat.completions.create(
    model="gpt-3.5-turbo",
    messages=[
        {
            "role": "user",
            "content": "How old is jason?",
        }
    ],
    response_model=User,
)

print(user.model_json_dump())
# {'messages': [{'content': 'How old is jason?', 'role': 'user'},
#               {'content': "Search retrieved: 'Jason is 20 years old'",
#                'role': 'user'}]}
# {'age': 20, 'name': 'Jason'}