@@ -30,6 +30,8 @@ class ModelMetadata(NamedTuple):
30
30
31
31
class LlmModel (str , Enum ):
32
32
# OpenAI models
33
+ O1_PREVIEW = "o1-preview"
34
+ O1_MINI = "o1-mini"
33
35
GPT4O_MINI = "gpt-4o-mini"
34
36
GPT4O = "gpt-4o"
35
37
GPT4_TURBO = "gpt-4-turbo"
@@ -57,6 +59,8 @@ def metadata(self) -> ModelMetadata:
57
59
58
60
59
61
MODEL_METADATA = {
62
+ LlmModel .O1_PREVIEW : ModelMetadata ("openai" , 32000 , cost_factor = 60 ),
63
+ LlmModel .O1_MINI : ModelMetadata ("openai" , 62000 , cost_factor = 30 ),
60
64
LlmModel .GPT4O_MINI : ModelMetadata ("openai" , 128000 , cost_factor = 10 ),
61
65
LlmModel .GPT4O : ModelMetadata ("openai" , 128000 , cost_factor = 12 ),
62
66
LlmModel .GPT4_TURBO : ModelMetadata ("openai" , 128000 , cost_factor = 11 ),
@@ -84,7 +88,10 @@ def metadata(self) -> ModelMetadata:
84
88
class AIStructuredResponseGeneratorBlock (Block ):
85
89
class Input (BlockSchema ):
86
90
prompt : str
87
- expected_format : dict [str , str ]
91
+ expected_format : dict [str , str ] = SchemaField (
92
+ description = "Expected format of the response. If provided, the response will be validated against this format. "
93
+ "The keys should be the expected fields in the response, and the values should be the description of the field." ,
94
+ )
88
95
model : LlmModel = LlmModel .GPT4_TURBO
89
96
api_key : BlockSecret = SecretField (value = "" )
90
97
sys_prompt : str = ""
@@ -132,7 +139,18 @@ def llm_call(
132
139
133
140
if provider == "openai" :
134
141
openai .api_key = api_key
135
- response_format = {"type" : "json_object" } if json_format else None
142
+ response_format = None
143
+
144
+ if model in [LlmModel .O1_MINI , LlmModel .O1_PREVIEW ]:
145
+ sys_messages = [p ["content" ] for p in prompt if p ["role" ] == "system" ]
146
+ usr_messages = [p ["content" ] for p in prompt if p ["role" ] != "system" ]
147
+ prompt = [
148
+ {"role" : "user" , "content" : "\n " .join (sys_messages )},
149
+ {"role" : "user" , "content" : "\n " .join (usr_messages )},
150
+ ]
151
+ elif json_format :
152
+ response_format = {"type" : "json_object" }
153
+
136
154
response = openai .chat .completions .create (
137
155
model = model .value ,
138
156
messages = prompt , # type: ignore
@@ -207,11 +225,11 @@ def trim_prompt(s: str) -> str:
207
225
format_prompt = ",\n " .join (expected_format )
208
226
sys_prompt = trim_prompt (
209
227
f"""
210
- |Reply in json format:
211
- |{{
212
- | { format_prompt }
213
- |}}
214
- """
228
+ |Reply strictly only in the following JSON format:
229
+ |{{
230
+ | { format_prompt }
231
+ |}}
232
+ """
215
233
)
216
234
prompt .append ({"role" : "system" , "content" : sys_prompt })
217
235
0 commit comments