Integrate OpenAI API Structured Generation

docs/overrides/home.html

@@ -117,4 +117,4 @@ <h2 class="subtitle" style="font-weight: 400; padding-top: 1rem;">
 </section>
 {% endblock %}
 {% block content %}{% endblock %}
-{% block footer %}{% endblock %}
+{% block footer %}{% endblock %}

docs/overrides/main.html

@@ -1,2 +1 @@
 {% extends "base.html" %}
-

docs/reference/models/openai.md

@@ -6,26 +6,25 @@
 
 ## OpenAI models
 
-Outlines supports models available via the OpenAI Chat API, e.g. ChatGPT and GPT-4. You can initialize the model by passing the model name to `outlines.models.openai`:
+Outlines supports models available via the OpenAI Chat API, e.g. GPT-4o, ChatGPT and GPT-4. You can initialize the model by passing the model name to `outlines.models.openai`:
 
 ```python
 from outlines import models
 
 
-model = models.openai("gpt-3.5-turbo")
-model = models.openai("gpt-4-turbo")
+model = models.openai("gpt-4o-mini")
 model = models.openai("gpt-4o")
 ```
 
-Check the [OpenAI documentation](https://platform.openai.com/docs/models/gpt-4-turbo-and-gpt-4) for an up-to-date list of available models. You can pass any parameter you would pass to `openai.AsyncOpenAI` as keyword arguments:
+Check the [OpenAI documentation](https://platform.openai.com/docs/models/gpt-4o) for an up-to-date list of available models. You can pass any parameter you would pass to `openai.AsyncOpenAI` as keyword arguments:
 
 ```python
 import os
 from outlines import models
 
 
 model = models.openai(
-    "gpt-3.5-turbo",
+    "gpt-4o-mini",
     api_key=os.environ["OPENAI_API_KEY"]
 )
 ```
@@ -56,8 +55,8 @@ from outlines import models
 
 model = models.azure_openai(
     "azure-deployment-name",
-    "gpt-3.5-turbo",
-    api_version="2023-07-01-preview",
+    "gpt-4o-mini",
+    api_version="2024-07-18",
     azure_endpoint="https://example-endpoint.openai.azure.com",
 )
 ```
@@ -111,6 +110,37 @@ model = models.openai(client, config)
 
     You need to pass the async client to be able to do batch inference.
 
+## Structured Generation Support
+
+Outlines provides support for [OpenAI Structured Outputs](https://platform.openai.com/docs/guides/structured-outputs/json-mode) via `outlines.generate.json`, `outlines.generate.choice`
+
+```python
+from pydantic import BaseModel, ConfigDict
+import outlines.models as models
+from outlines import generate
+
+model = models.openai("gpt-4o-mini")
+
+class Person(BaseModel):
+    model_config = ConfigDict(extra='forbid')  # required for openai
+    first_name: str
+    last_name: str
+    age: int
+
+generate.json(model, Person)
+generator("current indian prime minister on january 1st 2023")
+# Person(first_name='Narendra', last_name='Modi', age=72)
+
+generator = generate.choice(model, ["Chicken", "Egg"])
+print(generator("Which came first?"))
+# Chicken
+```
+
+!!! Warning
+
+    Structured generation support only provided to OpenAI-compatible endpoints which conform to OpenAI's standard. Additionally, `generate.regex` and `generate.cfg` are not supported.
+
+
 ## Advanced configuration
 
 For more advanced configuration option, such as support proxy, please consult the [OpenAI SDK's documentation](https://github.com/openai/openai-python):
@@ -146,7 +176,7 @@ config = OpenAIConfig(
     top_p=.95,
     seed=0,
 )
-model = models.openai("gpt-3.5-turbo", config)
+model = models.openai("gpt-4o-mini", config)
 ```
 
 ## Monitoring API use
@@ -158,7 +188,7 @@ from openai import AsyncOpenAI
 import outlines.models
 
 
-model = models.openai("gpt-4")
+model = models.openai("gpt-4o")
 
 print(model.prompt_tokens)
 # 0

docs/reference/text.md

@@ -9,7 +9,7 @@ Outlines provides a unified interface to generate text with many language models
 ```python
 from outlines import models, generate
 
-model = models.openai("gpt-4")
+model = models.openai("gpt-4o-mini")
 generator = generate.text(model)
 answer = generator("What is 2+2?")
 

outlines/generate/choice.py

@@ -1,10 +1,12 @@
+import json as pyjson
 from functools import singledispatch
 from typing import Callable, List
 
 from outlines.generate.api import SequenceGeneratorAdapter
 from outlines.models import OpenAI
 from outlines.samplers import Sampler, multinomial
 
+from .json import json
 from .regex import regex
 
 
@@ -24,13 +26,22 @@ def choice(
 def choice_openai(
     model: OpenAI, choices: List[str], sampler: Sampler = multinomial()
 ) -> Callable:
-    if not isinstance(sampler, multinomial):
-        raise NotImplementedError(
-            r"The OpenAI API does not support any other sampling algorithm "
-            + "that the multinomial sampler."
-        )
-
-    def generate_choice(prompt: str, max_tokens: int = 1):
-        return model.generate_choice(prompt, choices, max_tokens)
+    """
+    Call OpenAI API with response_format of a dict:
+    {"result": <one of choices>}
+    """
+
+    choices_schema = pyjson.dumps(
+        {
+            "type": "object",
+            "properties": {"result": {"type": "string", "enum": choices}},
+            "additionalProperties": False,
+            "required": ["result"],
+        }
+    )
+    generator = json(model, choices_schema, sampler)
+
+    def generate_choice(*args, **kwargs):
+        return generator(*args, **kwargs)["result"]
 
     return generate_choice

outlines/generate/json.py

@@ -72,7 +72,49 @@ def json(
 def json_openai(
     model, schema_object: Union[str, object, Callable], sampler: Sampler = multinomial()
 ):
-    raise NotImplementedError(
-        "Cannot use JSON Schema-structure generation with an OpenAI model "
-        + "due to the limitations of the OpenAI API"
+    if not isinstance(sampler, multinomial):
+        raise NotImplementedError(
+            r"The OpenAI API does not support any other sampling algorithm "
+            + "than the multinomial sampler."
+        )
+
+    if isinstance(schema_object, type(BaseModel)):
+        schema = pyjson.dumps(schema_object.model_json_schema())
+    elif callable(schema_object):
+        schema = pyjson.dumps(get_schema_from_signature(schema_object))
+    elif isinstance(schema_object, str):
+        schema = schema_object
+    else:
+        raise ValueError(
+            f"Cannot parse schema {schema_object}. The schema must be either "
+            + "a Pydantic object, a function or a string that contains the JSON "
+            + "Schema specification"
+        )
+
+    # create copied, patched model with normalized json schema set
+    generator = model.new_with_replacements(
+        response_format={
+            "type": "json_schema",
+            "json_schema": {
+                "name": "default",
+                "strict": True,
+                "schema": pyjson.loads(schema),
+            },
+        }
     )
+
+    # set generators sequence post-processor
+    if isinstance(schema_object, type(BaseModel)):
+        generator.format_sequence = lambda x: schema_object.parse_raw(x)
+    elif callable(schema_object):
+        generator.format_sequence = lambda x: pyjson.loads(x)
+    elif isinstance(schema_object, str) or isinstance(schema_object, dict):
+        generator.format_sequence = lambda x: pyjson.loads(x)
+    else:
+        raise ValueError(
+            f"Cannot parse schema {schema_object}. The schema must be either "
+            + "a Pydantic object, a function or a string that contains the JSON "
+            + "Schema specification"
+        )
+
+    return generator

outlines/models/openai.py

@@ -1,4 +1,5 @@
 """Integration with OpenAI's API."""
+import copy
 import functools
 import warnings
 from dataclasses import asdict, dataclass, field, replace
@@ -91,7 +92,6 @@ def __init__(
             parameters that cannot be set by calling this class' methods.
         tokenizer
             The tokenizer associated with the model the client connects to.
-
         """
 
         self.client = client
@@ -104,6 +104,8 @@ def __init__(
         self.prompt_tokens = 0
         self.completion_tokens = 0
 
+        self.format_sequence = lambda seq: seq
+
     def __call__(
         self,
         prompt: Union[str, List[str]],
@@ -152,7 +154,7 @@ def __call__(
         self.prompt_tokens += prompt_tokens
         self.completion_tokens += completion_tokens
 
-        return response
+        return self.format_sequence(response)
 
     def stream(self, *args, **kwargs):
         raise NotImplementedError(
@@ -250,9 +252,10 @@ def generate_choice(
 
         return choice
 
-    def generate_json(self):
-        """Call the OpenAI API to generate a JSON object."""
-        raise NotImplementedError
+    def new_with_replacements(self, **kwargs):
+        new_instance = copy.copy(self)
+        new_instance.config = replace(new_instance.config, **kwargs)
+        return new_instance
 
     def __str__(self):
         return self.__class__.__name__ + " API"