refacto: Use native OpenAI API JSON validation

microsoft · Jun 15, 2024 · 762cc9d · 762cc9d
1 parent 5da1201
commit 762cc9d
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/function_app.py b/function_app.py
@@ -405,16 +405,14 @@ async def page_to_fact(input: BlobClientTrigger) -> None:
     facts: list[FactModel] = []
     for _ in range(CONFIG.features.fact_iterations):  # We will generate facts 10 times
         def _validate(req: Optional[str]) -> tuple[bool, Optional[str], Optional[FactedLlmModel]]:
-            if not req:
-                return False, "Empty response", None
-            req = req.strip().strip("```json\n").strip("\n```").strip()
             try:
                 return True, None, FactedLlmModel.model_validate_json(req)
             except ValidationError as e:
                 return False, str(e), None
         facted_llm_model = await llm_client.generate(
             res_object=FactedLlmModel,
             temperature=1,  # We want creative answers
+            validate_json=True,
             validation_callback=_validate,
             prompt=f"""
             Assistant is an expert data analyst with 20 years of experience.

diff --git a/persistence/illm.py b/persistence/illm.py
@@ -12,8 +12,9 @@ async def generate(
         prompt: str,
         res_object: type[T],
         validation_callback: Callable[[Optional[str]], tuple[bool, Optional[str], Optional[T]]],
-        temperature: float = 0,
         max_tokens: Optional[int] = None,
+        temperature: float = 0,
+        validate_json: bool = False,
         _previous_result: Optional[str] = None,
         _retries_remaining: Optional[int] = None,
         _validation_error: Optional[str] = None,

diff --git a/persistence/openai.py b/persistence/openai.py
@@ -24,8 +24,9 @@ async def generate(
         prompt: str,
         res_object: type[T],
         validation_callback: Callable[[Optional[str]], tuple[bool, Optional[str], Optional[T]]],
-        temperature: float = 0,
         max_tokens: Optional[int] = None,
+        temperature: float = 0,
+        validate_json: bool = False,
         _previous_result: Optional[str] = None,
         _retries_remaining: Optional[int] = None,
         _validation_error: Optional[str] = None,
@@ -59,13 +60,18 @@ async def generate(
                 )
             )
 
+        extra = {}
+        if validate_json:
+            extra["response_format"] = {"type": "json_object"}
+
         # Generate
         client = self._use_client()
         res = await client.chat.completions.create(
             max_tokens=max_tokens,
             messages=messages,
             model=self._config.model,
             temperature=temperature,
+            **extra,
         )
         res_content = res.choices[0].message.content  # type: ignore
 
@@ -81,6 +87,7 @@ async def generate(
                 prompt=prompt,
                 res_object=res_object,
                 temperature=temperature,
+                validate_json=validate_json,
                 validation_callback=validation_callback,
                 _previous_result=res_content,
                 _retries_remaining=_retries_remaining - 1,