Marketing sample migration to AGNext (#234)

python/samples/marketing-agents/app.py

@@ -29,3 +29,6 @@ async def build_app(runtime: AgentRuntime) -> None:
 
     runtime.register("GraphicDesigner", lambda: GraphicDesignerAgent(client=image_client))
     runtime.register("Auditor", lambda: AuditAgent(model_client=chat_client))
+
+    runtime.get("GraphicDesigner")
+    runtime.get("Auditor")

python/samples/marketing-agents/auditor.py

@@ -30,4 +30,4 @@ class AuditAgent(TypeRoutedAgent):
         assert isinstance(completion.content, str)
         if "NOTFORME" in completion.content:
             return
-        await self.publish_message(AuditorAlert(user_id=message.user_id, auditor_alert_message=completion.content))
+        await self.publish_message(AuditorAlert(UserId=message.UserId, auditorAlertMessage=completion.content))

python/samples/marketing-agents/graphic_designer.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Literal
 
 import openai
@@ -21,11 +22,17 @@ class GraphicDesignerAgent(TypeRoutedAgent):
 
     @message_handler
     async def handle_user_chat_input(self, message: ArticleCreated, cancellation_token: CancellationToken) -> None:
-        response = await self._client.images.generate(
-            model=self._model, prompt=message.article, response_format="b64_json"
-        )
-        assert len(response.data) > 0 and response.data[0].b64_json is not None
-        image_base64 = response.data[0].b64_json
-        image_uri = f"data:image/png;base64,{image_base64}"
+        logger = logging.getLogger("graphic_designer")
+        try:
+            logger.info(f"Asking model to generate an image for the article '{message.article}'.")
+            response = await self._client.images.generate(
+                model=self._model, prompt=message.article, response_format="url"
+            )
+            logger.info(f"Image response: '{response.data[0]}'")
+            assert len(response.data) > 0 and response.data[0].url is not None
+            image_uri = response.data[0].url
+            logger.info(f"Generated image for article. Got response: '{image_uri}'")
 
-        await self.publish_message(GraphicDesignCreated(user_id=message.user_id, image_uri=image_uri))
+            await self.publish_message(GraphicDesignCreated(UserId=message.UserId, imageUri=image_uri))
+        except Exception as e:
+            logger.error(f"Failed to generate image for article. Error: {e}")

python/samples/marketing-agents/messages.py

@@ -2,20 +2,20 @@ from pydantic import BaseModel
 
 
 class ArticleCreated(BaseModel):
-    user_id: str
+    UserId: str
     article: str
 
 
 class GraphicDesignCreated(BaseModel):
-    user_id: str
-    image_uri: str
+    UserId: str
+    imageUri: str
 
 
 class AuditText(BaseModel):
-    user_id: str
+    UserId: str
     text: str
 
 
 class AuditorAlert(BaseModel):
-    user_id: str
-    auditor_alert_message: str
+    UserId: str
+    auditorAlertMessage: str

python/samples/marketing-agents/test_usage.py

@@ -17,14 +17,14 @@ class Printer(TypeRoutedAgent):
 
     @message_handler
     async def handle_graphic_design(self, message: GraphicDesignCreated, cancellation_token: CancellationToken) -> None:
-        image = Image.from_uri(message.image_uri)
+        image = Image.from_uri(message.imageUri)
         # Save image to random name in current directory
-        image.image.save(os.path.join(os.getcwd(), f"{message.user_id}.png"))
-        print(f"Received GraphicDesignCreated: user {message.user_id}, saved to {message.user_id}.png")
+        image.image.save(os.path.join(os.getcwd(), f"{message.UserId}.png"))
+        print(f"Received GraphicDesignCreated: user {message.UserId}, saved to {message.UserId}.png")
 
     @message_handler
     async def handle_auditor_alert(self, message: AuditorAlert, cancellation_token: CancellationToken) -> None:
-        print(f"Received AuditorAlert: {message.auditor_alert_message} for user {message.user_id}")
+        print(f"Received AuditorAlert: {message.auditorAlertMessage} for user {message.UserId}")
 
 
 async def main() -> None:
@@ -35,11 +35,11 @@ async def main() -> None:
     ctx = runtime.start()
 
     await runtime.publish_message(
-        AuditText(text="Buy my product for a MASSIVE 50% discount.", user_id="user-1"), namespace="default"
+        AuditText(text="Buy my product for a MASSIVE 50% discount.", UserId="user-1"), namespace="default"
     )
 
     await runtime.publish_message(
-        ArticleCreated(article="The best article ever written about trees and rocks", user_id="user-2"),
+        ArticleCreated(article="The best article ever written about trees and rocks", UserId="user-2"),
         namespace="default",
     )
 

python/samples/marketing-agents/worker.py

@@ -1,24 +1,40 @@
 import asyncio
 import logging
-import os
 
+from agnext.core._serialization import MESSAGE_TYPE_REGISTRY
 from agnext.worker.worker_runtime import WorkerAgentRuntime
 from app import build_app
+from dotenv import load_dotenv
+from messages import ArticleCreated, AuditorAlert, AuditText, GraphicDesignCreated
+
+agnext_logger = logging.getLogger("agnext")
 
 
 async def main() -> None:
+    load_dotenv()
     runtime = WorkerAgentRuntime()
-    await runtime.setup_channel(os.environ["AGENT_HOST"])
+    MESSAGE_TYPE_REGISTRY.add_type(ArticleCreated)
+    MESSAGE_TYPE_REGISTRY.add_type(GraphicDesignCreated)
+    MESSAGE_TYPE_REGISTRY.add_type(AuditText)
+    MESSAGE_TYPE_REGISTRY.add_type(AuditorAlert)
+    agnext_logger.info("1")
+    await runtime.setup_channel("localhost:5145")
+
+    agnext_logger.info("2")
 
     await build_app(runtime)
+    agnext_logger.info("3")
 
     # just to keep the runtime running
     try:
         await asyncio.sleep(1000000)
     except KeyboardInterrupt:
         pass
+    await runtime.close_channel()
 
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)
+    agnext_logger.setLevel(logging.DEBUG)
+    agnext_logger.log(logging.DEBUG, "Starting worker")
     asyncio.run(main())