[Test] Basic AOAI JSON test (#1253)

A very simple test of JSON functionality in AOAI. This does rely on the model having some 'baseline' knowledge and reasoning power.... should be OK for GPT-4 class models.

Author: riedgar-ms

tests/model_specific/common_chat_testing.py

@@ -1,4 +1,7 @@
+import pydantic
+
 from guidance import assistant, gen, models, system, user
+from guidance import json as gen_json
 
 
 def smoke_chat(lm: models.Model, has_system_role: bool = True):
@@ -74,3 +77,23 @@ def longer_chat_2(lm: models.Model, has_system_role: bool = True):
 
     print(str(lm))
     assert len(lm["number"]) > 0
+
+
+def json_output_smoke(lm: models.Model):
+    class NameHolder(pydantic.BaseModel):
+        my_name: str
+        my_age: int
+        model_config = pydantic.ConfigDict(extra="forbid")
+
+    with user():
+        lm += "Hello, my name is Tweedledum and I am 10 years old. What is my twin brother's name and age?"
+
+    EXTRACT_KEY = "my_output"
+    with assistant():
+        lm += gen_json(name=EXTRACT_KEY, schema=NameHolder)
+
+    output_json = lm[EXTRACT_KEY]
+    name_data = NameHolder.model_validate_json(output_json)
+    print(name_data.model_dump_json(indent=4))
+    assert name_data.my_name == "Tweedledee"
+    assert name_data.my_age == 10

tests/need_credentials/test_azureai_openai.py

@@ -2,12 +2,15 @@
 
 from urllib.parse import parse_qs, urlparse
 
+import json
+import pydantic
 import pytest
 import requests
 
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 
 from guidance import assistant, gen, models, system, user, gen_audio, image
+from guidance import json as gen_json
 from guidance.models._azureai import create_azure_openai_model
 
 from ..model_specific import common_chat_testing
@@ -110,6 +113,10 @@ def test_azureai_openai_chat_smoke(azureai_chat_model):
     common_chat_testing.smoke_chat(azureai_chat_model)
 
 
+def test_azureai_openai_chat_json(azureai_chat_model: models.Model):
+    common_chat_testing.json_output_smoke(azureai_chat_model)
+
+
 def test_azureai_openai_audio_smoke(azureai_audio_model: models.Model):
     lm = azureai_audio_model
     with system():
@@ -189,7 +196,9 @@ def test_azureai_openai_completion_smoke():
     )
 
     lm = models.AzureOpenAI(
-        model=model, azure_endpoint=azureai_endpoint, azure_ad_token_provider=token_provider
+        model=model,
+        azure_endpoint=azureai_endpoint,
+        azure_ad_token_provider=token_provider,
     )
     assert isinstance(lm, models.AzureOpenAI)
     assert isinstance(lm.engine, models._openai.OpenAIEngine)
@@ -240,7 +249,9 @@ def test_azureai_openai_chat_loop(azureai_chat_model):
     for i in range(2):
         print(f"Iteration: {i}")
         with system():
-            generation = azureai_chat_model + "You will just return whatever number I give you"
+            generation = (
+                azureai_chat_model + "You will just return whatever number I give you"
+            )
 
         with user():
             generation += f"The number is: {i}"

tests/need_credentials/test_openai.py

@@ -78,3 +78,7 @@ def test_openai_image_smoke(openai_image_model: Model):
         lm += gen(name="img_describe")
     assert "img_describe" in lm
     assert len(lm["img_describe"]) > 0
+
+
+def test_openai_chat_json(openai_model: Model):
+    common_chat_testing.json_output_smoke(openai_model)