Update agents.py

agents.py

@@ -1,26 +1,38 @@
-from phi.assistant import Assistant
-from phi.llm.google import Gemini
-from phi.tools.wikipedia import WikipediaTools
-
-def get_content_enhancer_agent(knowledge_base=None):
-    """Factory function to create the content enhancement agent."""
-    return Assistant(
-        llm=Gemini(model="gemini-1.5-flash-latest"),
-        tools=[WikipediaTools()], 
-        knowledge_base=knowledge_base,
-        search_knowledge=True,
-        description="You are an expert multimedia content creator and writer.",
-        instructions=[
-            "You will receive a summary from a video analysis.",
-            "Your primary goal is to enrich this summary by:",
-            "  1. Using your knowledge base and Wikipedia to find more details, facts, and context about the key topics mentioned.",
-            "  2. Weaving all information (video summary, knowledge base, Wikipedia) into a single, cohesive, and engaging narrative.",
-            "  3. Identifying the most visually compelling moments or concepts in the narrative.",
-            "At these points, you MUST insert a placeholder in the exact format: [IMAGE: caption=\"A detailed, vivid prompt for an AI image generator\"]",
-            "The caption prompt should be descriptive: think about subject, action, style, lighting, and composition.",
-            "Prioritize the information from the video summary, then augment it with external knowledge.",
-            "Return only the final, complete text with the placeholders."
-        ],
-        # shows the LLM thinking process
-        show_tool_calls=True,
+from phi.assistant import Assistant
+from phi.tools.wikipedia import WikipediaTools
+from google import genai
+import os
+
+client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))
+
+class GeminiWrapper:
+    """Custom wrapper to use new Gemini SDK inside phi"""
+
+    def __init__(self, model="gemini-2.0-flash"):
+        self.model = model
+
+    def __call__(self, prompt, **kwargs):
+        response = client.models.generate_content(
+            model=self.model,
+            contents=prompt
+        )
+        return response.text
+
+
+def get_content_enhancer_agent(knowledge_base=None):
+    return Assistant(
+        llm=GeminiWrapper(),
+        tools=[WikipediaTools()],
+        knowledge_base=knowledge_base,
+        search_knowledge=True,
+        description="You are an expert multimedia content creator and writer.",
+        instructions=[
+            "You will receive a summary from a video analysis.",
+            "Enhance it using knowledge + Wikipedia.",
+            "Create a cohesive narrative.",
+            "Insert [IMAGE: caption=\"...\"] placeholders.",
+            "Make prompts vivid and visual.",
+            "Return only final text."
+        ],
+        show_tool_calls=True,
     )