Merge pull request #1304 from kritinv/local-model-fix

local model/azure fix
confident-ai · Jan 23, 2025 · 69c558f · 69c558f
2 parents 714ad12 + c7ade1d
commit 69c558f
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/deepeval/models/gpt_model.py b/deepeval/models/gpt_model.py
@@ -245,6 +245,14 @@ def generate(
                     completion.usage.completion_tokens,
                 )
                 return schema.model_validate(json_output), cost
+        elif schema and not using_openai_model:
+            chat_model = self.load_model()
+            with get_openai_callback() as cb:
+                res = chat_model.invoke(prompt)
+                json_output = self.trim_and_load_json(
+                    res.content
+                )
+                return schema.model_validate(json_output), cb.total_cost
         else:
             chat_model = self.load_model()
             with get_openai_callback() as cb:
@@ -297,6 +305,14 @@ async def a_generate(
                     completion.usage.completion_tokens,
                 )
                 return schema.model_validate(json_output), cost
+        elif schema and not using_openai_model:
+            chat_model = self.load_model()
+            with get_openai_callback() as cb:
+                res = await chat_model.ainvoke(prompt)
+                json_output = self.trim_and_load_json(
+                    res.content
+                )
+                return schema.model_validate(json_output), cb.total_cost
         else:
             chat_model = self.load_model()
             with get_openai_callback() as cb: