Allow literal values in pydantic model

gpt_json/prompts.py

@@ -1,5 +1,5 @@
 from types import UnionType
-from typing import List, Type, get_args, get_origin
+from typing import List, Literal, Type, get_args, get_origin
 
 from pydantic import BaseModel
 
@@ -17,6 +17,7 @@ def generate_payload(model: Type[BaseModel]):
             field_annotation = value.annotation
             annotation_origin = get_origin(field_annotation)
             annotation_arguments = get_args(field_annotation)
+            print(annotation_origin)
 
             if field_annotation is None:
                 continue
@@ -31,6 +32,9 @@ def generate_payload(model: Type[BaseModel]):
                 payload.append(
                     f'"{key}": {" | ".join([arg.__name__.lower() for arg in annotation_arguments])}'
                 )
+            elif annotation_origin == Literal:
+                allowed_values = [f'"{arg}"' for arg in annotation_arguments]
+                payload.append(f'"{key}": {" | ".join(allowed_values)}')
             elif issubclass(field_annotation, BaseModel):
                 payload.append(f'"{key}": {generate_payload(field_annotation)}')
             else:

gpt_json/tests/shared.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from typing import Literal
 
 from pydantic import BaseModel, Field
 
@@ -15,6 +16,12 @@ class MySchema(BaseModel):
     reason: bool = Field(description="Explanation")
 
 
+class LiteralSchema(BaseModel):
+    work_format: Literal["REMOTE", "OFFICE", "ANY"] = Field(
+        default="ANY", description="One of the given values"
+    )
+
+
 class UnitType(Enum):
     CELSIUS = "celsius"
     FAHRENHEIT = "fahrenheit"

gpt_json/tests/test_prompts.py

@@ -5,7 +5,7 @@
 from gpt_json.generics import resolve_generic_model
 from gpt_json.gpt import ListResponse
 from gpt_json.prompts import generate_schema_prompt
-from gpt_json.tests.shared import MySchema
+from gpt_json.tests.shared import LiteralSchema, MySchema
 
 
 def strip_whitespace(input_string: str):
@@ -45,6 +45,14 @@ def strip_whitespace(input_string: str):
             }}
             """,
         ),
+        (
+            LiteralSchema,
+            """
+            {{
+            "work_format": "REMOTE" | "OFFICE" | "ANY" // One of the given values
+            }}
+            """,
+        ),
     ],
 )
 def test_generate_schema_prompt(schema_definition, expected: str):