From dcdda9131d46b904a794b5a588b372e8f61c5513 Mon Sep 17 00:00:00 2001
From: ysolanky <yash@phidata.com>
Date: Thu, 19 Sep 2024 18:17:44 -0400
Subject: [PATCH] fix-agent-output-model-phi-1236

---
 cookbook/providers/__init__.py                |  0
 .../providers/openai/structured_output.py     | 26 +++++++++++++++++++
 phi/agent/agent.py                            |  4 +--
 phi/playground/playground.py                  |  4 ++-
 4 files changed, 31 insertions(+), 3 deletions(-)
 create mode 100644 cookbook/providers/__init__.py
 create mode 100644 cookbook/providers/openai/structured_output.py

diff --git a/cookbook/providers/__init__.py b/cookbook/providers/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/cookbook/providers/openai/structured_output.py b/cookbook/providers/openai/structured_output.py
new file mode 100644
index 000000000..52d17fded
--- /dev/null
+++ b/cookbook/providers/openai/structured_output.py
@@ -0,0 +1,26 @@
+from typing import List
+from pydantic import BaseModel, Field
+from rich.pretty import pprint
+from phi.agent import Agent
+from phi.model.openai import OpenAIChat
+
+
+class MovieScript(BaseModel):
+    setting: str = Field(..., description="Provide a nice setting for a blockbuster movie.")
+    ending: str = Field(..., description="Ending of the movie. If not available, provide a happy ending.")
+    genre: str = Field(
+        ..., description="Genre of the movie. If not available, select action, thriller or romantic comedy."
+    )
+    name: str = Field(..., description="Give a name to this movie")
+    characters: List[str] = Field(..., description="Name of characters for this movie.")
+    storyline: str = Field(..., description="3 sentence storyline for the movie. Make it exciting!")
+
+
+movie_agent = Agent(
+    model=OpenAIChat(model="gpt-4o"),
+    description="You help people write movie scripts.",
+    output_model=MovieScript,
+    # debug_mode=True,
+)
+
+pprint(movie_agent.run("New York"))
diff --git a/phi/agent/agent.py b/phi/agent/agent.py
index 57c3a8f5f..f08aae60b 100644
--- a/phi/agent/agent.py
+++ b/phi/agent/agent.py
@@ -1079,7 +1079,7 @@ def run(
                 try:
                     structured_output = None
                     try:
-                        structured_output = self.output_model.model_validate(run_response.content)
+                        structured_output = self.output_model.model_validate(json.loads(run_response.content))
                     except ValidationError:
                         # Check if response starts with ```json
                         if run_response.content.startswith("```json"):
@@ -1362,7 +1362,7 @@ async def arun(
                 try:
                     structured_output = None
                     try:
-                        structured_output = self.output_model.model_validate_json(run_response.content)
+                        structured_output = self.output_model.model_validate(json.loads(run_response.content))
                     except ValidationError:
                         # Check if response starts with ```json
                         if run_response.content.startswith("```json"):
diff --git a/phi/playground/playground.py b/phi/playground/playground.py
index 6403f2555..e750f42fb 100644
--- a/phi/playground/playground.py
+++ b/phi/playground/playground.py
@@ -96,7 +96,9 @@ def agent_get():
                         ),
                         enable_rag=agent.enable_rag,
                         tools=formatted_tools,
-                        memory={"name": agent.memory.db.__class__.__name__} if agent.memory and agent.memory.db else None,
+                        memory={"name": agent.memory.db.__class__.__name__}
+                        if agent.memory and agent.memory.db
+                        else None,
                         storage={"name": agent.storage.__class__.__name__} if agent.storage else None,
                         knowledge={"name": agent.knowledge.__class__.__name__} if agent.knowledge else None,
                         description=agent.description,