feat: add chat proto image support (#76)

* feat: add chat proto image support

* chore: update deployed agent addresses

Author: jrriehl

6-deployed-agents/knowledge-base/asi1-agent/README.md

@@ -0,0 +1,8 @@
+# ASI1-Mini Agent
+
+![domain:knowledge-base](https://img.shields.io/badge/knowledge--base-3D8BD3?style=flat&logo=data%3Aimage%2Fsvg%2Bxml%3Bbase64%2CPHN2ZyB3aWR0aD0iMTAiIGhlaWdodD0iOCIgdmlld0JveD0iMCAwIDEwIDgiIGZpbGw9Im5vbmUiIHhtbG5zPSJodHRwOi8vd3d3LnczLm9yZy8yMDAwL3N2ZyI%2BCjxwYXRoIGQ9Ik00Ljc1IDAuNTAwMDA4VjYuMzc1MDFMMS41IDUuNzUwMDFWMC41NDY4ODNDMS41IDAuMjM0MzgzIDEuNzY1NjIgNy44MzMzNWUtMDYgMi4wNzgxMiAwLjA2MjUwNzhMNC43NSAwLjUwMDAwOFpNMS4zOTA2MiA2LjM0Mzc2TDUgNy4wNjI1MUw4LjU5Mzc1IDYuMzQzNzZDOC44MjgxMiA2LjI5Njg4IDkgNi4wNzgxMyA5IDUuODQzNzZWMC40NTMxMzNMOS4zOTA2MiAwLjM3NTAwOEM5LjcwMzEyIDAuMzEyNTA4IDEwIDAuNTQ2ODgzIDEwIDAuODU5MzgzVjYuNTkzNzZDMTAgNi44NDM3NiA5LjgyODEyIDcuMDQ2ODggOS41OTM3NSA3LjA5Mzc2TDUgOC4wMDAwMUwwLjM5MDYyNSA3LjA5Mzc2QzAuMTU2MjUgNy4wNDY4OCAwIDYuODI4MTMgMCA2LjU5Mzc2VjAuODU5MzgzQzAgMC41NDY4ODMgMC4yODEyNSAwLjMxMjUwOCAwLjU5Mzc1IDAuMzc1MDA4TDEgMC40NTMxMzNWNS44NDM3NkMxIDYuMDkzNzYgMS4xNTYyNSA2LjI5Njg4IDEuMzkwNjIgNi4zNDM3NlpNNS4yNSA2LjM3NTAxVjAuNTAwMDA4TDcuOTA2MjUgMC4wNjI1MDc4QzguMjE4NzUgNy44MzMzNWUtMDYgOC41IDAuMjM0MzgzIDguNSAwLjU0Njg4M1Y1Ljc1MDAxTDUuMjUgNi4zNzUwMVoiIGZpbGw9IndoaXRlIi8%2BCjwvc3ZnPgo%3D)
+![tech:llm](https://img.shields.io/badge/llm-E85D2E?style=flat&logo=data%3Aimage%2Fsvg%2Bxml%3Bbase64%2CPHN2ZyB3aWR0aD0iMTAiIGhlaWdodD0iOCIgdmlld0JveD0iMCAwIDEwIDgiIGZpbGw9Im5vbmUiIHhtbG5zPSJodHRwOi8vd3d3LnczLm9yZy8yMDAwL3N2ZyI%2BCjxwYXRoIGQ9Ik00LjUgMUM0LjUgMS4yMTg3NSA0LjQyMTg4IDEuNDIxODggNC4zMTI1IDEuNTc4MTJMNC43NjU2MiAyLjU2MjVDNC45MjE4OCAyLjUzMTI1IDUuMDc4MTIgMi41IDUuMjUgMi41QzUuODEyNSAyLjUgNi4zMjgxMiAyLjcxODc1IDYuNzE4NzUgMy4wNjI1TDggMi4xMDkzOEM4IDIuMDc4MTIgOCAyLjA0Njg4IDggMkM4IDEuNDUzMTIgOC40Mzc1IDEgOSAxQzkuNTQ2ODggMSAxMCAxLjQ1MzEyIDEwIDJDMTAgMi41NjI1IDkuNTQ2ODggMyA5IDNDOC44NDM3NSAzIDguNzE4NzUgMi45ODQzOCA4LjU5Mzc1IDIuOTIxODhMNy4zMTI1IDMuODU5MzhDNy40MjE4OCA0LjE0MDYyIDcuNSA0LjQzNzUgNy41IDQuNzVDNy41IDUgNy40NTMxMiA1LjIzNDM4IDcuMzc1IDUuNDUzMTJMOC41IDYuMTI1QzguNjU2MjUgNi4wNDY4OCA4LjgxMjUgNiA5IDZDOS41NDY4OCA2IDEwIDYuNDUzMTIgMTAgN0MxMCA3LjU2MjUgOS41NDY4OCA4IDkgOEM4LjQzNzUgOCA4IDcuNTYyNSA4IDdWNi45ODQzOEw2Ljg1OTM4IDYuMzEyNUM2LjQ1MzEyIDYuNzM0MzggNS44NzUgNyA1LjI1IDdDNC4xNzE4OCA3IDMuMjgxMjUgNi4yNjU2MiAzLjA0Njg4IDUuMjVIMS44NTkzOEMxLjY4NzUgNS41NjI1IDEuMzU5MzggNS43NSAxIDUuNzVDMC40Mzc1IDUuNzUgMCA1LjMxMjUgMCA0Ljc1QzAgNC4yMDMxMiAwLjQzNzUgMy43NSAxIDMuNzVDMS4zNTkzOCAzLjc1IDEuNjg3NSAzLjk1MzEyIDEuODU5MzggNC4yNUgzLjA0Njg4QzMuMTcxODggMy43MzQzOCAzLjQ1MzEyIDMuMjk2ODggMy44NTkzOCAyLjk4NDM4TDMuNDA2MjUgMkMyLjg5MDYyIDEuOTUzMTIgMi41IDEuNTMxMjUgMi41IDFDMi41IDAuNDUzMTI1IDIuOTM3NSAwIDMuNSAwQzQuMDQ2ODggMCA0LjUgMC40NTMxMjUgNC41IDFaTTUuMjUgNS41QzUuNTE1NjIgNS41IDUuNzUgNS4zNTkzOCA1Ljg5MDYyIDUuMTI1QzYuMDMxMjUgNC45MDYyNSA2LjAzMTI1IDQuNjA5MzggNS44OTA2MiA0LjM3NUM1Ljc1IDQuMTU2MjUgNS41MTU2MiA0IDUuMjUgNEM0Ljk2ODc1IDQgNC43MzQzOCA0LjE1NjI1IDQuNTkzNzUgNC4zNzVDNC40NTMxMiA0LjYwOTM4IDQuNDUzMTIgNC45MDYyNSA0LjU5Mzc1IDUuMTI1QzQuNzM0MzggNS4zNTkzOCA0Ljk2ODc1IDUuNSA1LjI1IDUuNVoiIGZpbGw9IndoaXRlIi8%2BCjwvc3ZnPgo%3D)
+[![link to source code](https://img.shields.io/badge/Source%20Code-E8ECF1?style=flat&logo=data%3Aimage%2Fsvg%2Bxml%3Bbase64%2CPHN2ZyB3aWR0aD0iOCIgaGVpZ2h0PSI4IiB2aWV3Qm94PSIwIDAgOCA4IiBmaWxsPSJub25lIiB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciPgo8cGF0aCBkPSJNNCAwLjA5ODk5OUMxLjc5IDAuMDk4OTk5IDAgMS44OSAwIDQuMDk5QzAgNS44NjY2NyAxLjE0NiA3LjM2NTY2IDIuNzM1IDcuODk0QzIuOTM1IDcuOTMxNjYgMy4wMDgzMyA3LjgwOCAzLjAwODMzIDcuNzAxNjZDMy4wMDgzMyA3LjYwNjY2IDMuMDA1IDcuMzU1IDMuMDAzMzMgNy4wMjE2N0MxLjg5MDY3IDcuMjYzIDEuNjU2IDYuNDg1IDEuNjU2IDYuNDg1QzEuNDc0IDYuMDIzMzMgMS4yMTEgNS45IDEuMjExIDUuOUMwLjg0ODY2NyA1LjY1MiAxLjIzOSA1LjY1NyAxLjIzOSA1LjY1N0MxLjY0MDY3IDUuNjg1IDEuODUxNjcgNi4wNjkgMS44NTE2NyA2LjA2OUMyLjIwODMzIDYuNjgwNjcgMi43ODggNi41MDQgMy4wMTY2NyA2LjQwMTY2QzMuMDUyNjcgNi4xNDMgMy4xNTU2NyA1Ljk2NjY3IDMuMjcgNS44NjY2N0MyLjM4MTY3IDUuNzY2NjcgMS40NDggNS40MjI2NyAxLjQ0OCAzLjg5QzEuNDQ4IDMuNDUzMzMgMS42MDMgMy4wOTY2NyAxLjg1OTY3IDIuODE2NjdDMS44MTQ2NyAyLjcxNTY3IDEuNjc5NjcgMi4zMDkgMS44OTQ2NyAxLjc1OEMxLjg5NDY3IDEuNzU4IDIuMjI5NjcgMS42NTA2NyAyLjk5NDY3IDIuMTY4QzMuMzE0NjcgMi4wNzkgMy42NTQ2NyAyLjAzNSAzLjk5NDY3IDIuMDMzQzQuMzM0NjcgMi4wMzUgNC42NzQ2NyAyLjA3OSA0Ljk5NDY3IDIuMTY4QzUuNzU0NjcgMS42NTA2NyA2LjA4OTY3IDEuNzU4IDYuMDg5NjcgMS43NThDNi4zMDQ2NyAyLjMwOSA2LjE2OTY3IDIuNzE1NjcgNi4xMjk2NyAyLjgxNjY3QzYuMzg0NjcgMy4wOTY2NyA2LjUzOTY3IDMuNDUzMzMgNi41Mzk2NyAzLjg5QzYuNTM5NjcgNS40MjY2NyA1LjYwNDY3IDUuNzY1IDQuNzE0NjcgNS44NjMzM0M0Ljg1NDY3IDUuOTgzMzMgNC45ODQ2NyA2LjIyODY2IDQuOTg0NjcgNi42MDMzM0M0Ljk4NDY3IDcuMTM4NjYgNC45Nzk2NyA3LjU2ODY3IDQuOTc5NjcgNy42OTg2N0M0Ljk3OTY3IDcuODAzNjcgNS4wNDk2NyA3LjkyODY3IDUuMjU0NjcgNy44ODg2N0M2Ljg1NSA3LjM2NCA4IDUuODY0IDggNC4wOTlDOCAxLjg5IDYuMjA5IDAuMDk4OTk5IDQgMC4wOTg5OTlaIiBmaWxsPSIjNTU2NTc4Ii8%2BCjwvc3ZnPgo%3D)](https://github.com/fetchai/uAgent-Examples/tree/main/6-deployed-agents/knowledge-base/asi1-agent)
+[![live](https://img.shields.io/badge/Live-8A2BE2?style=flat&logo=data%3Aimage%2Fsvg%2Bxml%3Bbase64%2CPHN2ZyB3aWR0aD0iMTAiIGhlaWdodD0iOCIgdmlld0JveD0iMCAwIDEwIDgiIGZpbGw9Im5vbmUiIHhtbG5zPSJodHRwOi8vd3d3LnczLm9yZy8yMDAwL3N2ZyI%2BCjxwYXRoIGQ9Ik0yLjI1IDcuNUMxIDcuNSAwIDYuNSAwIDUuMjVDMCA0LjI4MTI1IDAuNjI1IDMuNDM3NSAxLjUgMy4xNDA2MkMxLjUgMy4wOTM3NSAxLjUgMy4wNDY4OCAxLjUgM0MxLjUgMS42MjUgMi42MDkzOCAwLjUgNCAwLjVDNC45MjE4OCAwLjUgNS43MzQzOCAxLjAxNTYyIDYuMTU2MjUgMS43NjU2MkM2LjM5MDYyIDEuNTkzNzUgNi42ODc1IDEuNSA3IDEuNUM3LjgyODEyIDEuNSA4LjUgMi4xNzE4OCA4LjUgM0M4LjUgMy4yMDMxMiA4LjQ1MzEyIDMuMzc1IDguMzkwNjIgMy41NDY4OEM5LjMxMjUgMy43MzQzOCAxMCA0LjU0Njg4IDEwIDUuNUMxMCA2LjYwOTM4IDkuMDkzNzUgNy41IDggNy41SDIuMjVaTTYuNzY1NjIgMy43NjU2MkM2LjkwNjI1IDMuNjI1IDYuOTA2MjUgMy4zOTA2MiA2Ljc2NTYyIDMuMjVDNi42MDkzOCAzLjA5Mzc1IDYuMzc1IDMuMDkzNzUgNi4yMzQzOCAzLjI1TDQuNSA0Ljk4NDM4TDMuNzY1NjIgNC4yNUMzLjYwOTM4IDQuMDkzNzUgMy4zNzUgNC4wOTM3NSAzLjIzNDM4IDQuMjVDMy4wNzgxMiA0LjM5MDYyIDMuMDc4MTIgNC42MjUgMy4yMzQzOCA0Ljc2NTYyTDQuMjM0MzggNS43NjU2MkM0LjM3NSA1LjkyMTg4IDQuNjA5MzggNS45MjE4OCA0Ljc2NTYyIDUuNzY1NjJMNi43NjU2MiAzLjc2NTYyWiIgZmlsbD0id2hpdGUiLz4KPC9zdmc%2BCg%3D%3D)](https://agentverse.ai/agents/details/agent1q0h70caed8ax769shpemapzkyk65uscw4xwk6dc4t3emvp5jdcvqs9xs32y/profile)
+
+This agent is a simple wrapper around `asi1-mini` Large Language Model. For more details go to https://docs.asi1.ai/docs

6-deployed-agents/knowledge-base/asi1-agent/agent.py

@@ -0,0 +1,70 @@
+import os
+from enum import Enum
+
+from protocols import chat_proto
+from uagents import Agent, Context, Model
+from uagents.experimental.quota import QuotaProtocol
+
+AGENT_SEED = os.getenv("AGENT_SEED", "asi1-test-agent")
+AGENT_NAME = os.getenv("AGENT_NAME", "ASI1-Mini Agent")
+
+
+PORT = 8000
+agent = Agent(
+    name=AGENT_NAME,
+    seed=AGENT_SEED,
+    port=PORT,
+    endpoint=f"http://localhost:{PORT}/submit",
+)
+
+agent.include(chat_proto, publish_manifest=True)
+
+
+### Health check related code
+def agent_is_healthy() -> bool:
+    """
+    Implement the actual health check logic here.
+
+    For example, check if the agent can connect to a third party API,
+    check if the agent has enough resources, etc.
+    """
+    condition = True  # TODO: logic here
+    return bool(condition)
+
+
+class HealthCheck(Model):
+    pass
+
+
+class HealthStatus(str, Enum):
+    HEALTHY = "healthy"
+    UNHEALTHY = "unhealthy"
+
+
+class AgentHealth(Model):
+    agent_name: str
+    status: HealthStatus
+
+
+health_protocol = QuotaProtocol(
+    storage_reference=agent.storage, name="HealthProtocol", version="0.1.0"
+)
+
+
+@health_protocol.on_message(HealthCheck, replies={AgentHealth})
+async def handle_health_check(ctx: Context, sender: str, msg: HealthCheck):
+    status = HealthStatus.UNHEALTHY
+    try:
+        if agent_is_healthy():
+            status = HealthStatus.HEALTHY
+    except Exception as err:
+        ctx.logger.error(err)
+    finally:
+        await ctx.send(sender, AgentHealth(agent_name=AGENT_NAME, status=status))
+
+
+agent.include(health_protocol, publish_manifest=True)
+
+
+if __name__ == "__main__":
+    agent.run()

6-deployed-agents/knowledge-base/asi1-agent/ai.py

@@ -0,0 +1,66 @@
+import os
+from typing import Any, List
+from openai import OpenAI, OpenAIError
+from openai.types.chat import ChatCompletionMessageParam
+
+
+ASI1_URL = "https://api.asi1.ai/v1"
+MODEL_NAME = os.getenv("MODEL_NAME", "asi1-mini")
+ASI1_API_KEY = os.getenv("ASI1_API_KEY", "")
+MAX_TOKENS = int(os.getenv("MAX_TOKENS", "4096"))
+
+
+if not ASI1_API_KEY:
+    raise ValueError(
+        "You need to provide an API key:https://docs.asi1.ai/docs/getting-api-key"
+    )
+
+
+client = OpenAI(api_key=ASI1_API_KEY, base_url=ASI1_URL)
+
+
+# Send a prompt and context to the AI model and return the content of the completion
+def get_completion(
+    content: List[dict[str, Any]],
+    max_tokens: int = MAX_TOKENS,
+) -> str:
+    message_parts = []
+
+    for item in content:
+        if item.get("type") == "text":
+            message_parts.append({"type": "text", "text": item["text"]})
+        elif item.get("type") == "resource":
+            mime_type = item["mime_type"]
+            if mime_type.startswith("image/"):
+                message_parts.append({
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:{mime_type};base64,{item['contents']}",
+                    },
+                })
+            else:
+                return f"Unsupported mime type: {mime_type}"
+
+
+    if len(message_parts) == 1 and message_parts[0].get("type") == "text":
+        user_content = message_parts[0]["text"]
+    else:
+        user_content = message_parts
+
+    try:
+        messages: List[ChatCompletionMessageParam] = [
+            {"role": "system", "content": "You are Fetch.ai agent, running on AgentVerse.ai platform."},
+            {"role": "user", "content": user_content},
+        ]
+        response = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=messages,
+            max_tokens=max_tokens,
+        )
+
+        content = response.choices[0].message.content
+
+        return str(content)
+
+    except OpenAIError as e:
+        return f"An error occurred: {e}"

6-deployed-agents/knowledge-base/asi1-agent/protocols.py

@@ -0,0 +1,82 @@
+import os
+from datetime import datetime
+from uuid import uuid4
+
+from ai import get_completion
+from uagents import Context, Protocol
+from uagents_core.contrib.protocols.chat import (
+    ChatAcknowledgement,
+    ChatMessage,
+    MetadataContent,
+    ResourceContent,
+    StartSessionContent,
+    TextContent,
+    chat_protocol_spec,
+)
+from uagents_core.storage import ExternalStorage
+
+STORAGE_URL = os.getenv("AGENTVERSE_URL", "https://agentverse.ai") + "/v1/storage"
+
+
+def create_text_chat(text: str) -> ChatMessage:
+    return ChatMessage(
+        timestamp=datetime.utcnow(),
+        msg_id=uuid4(),
+        content=[TextContent(type="text", text=text)],
+    )
+
+
+def create_metadata(metadata: dict[str, str]) -> ChatMessage:
+    return ChatMessage(
+        timestamp=datetime.utcnow(),
+        msg_id=uuid4(),
+        content=[MetadataContent(type="metadata", metadata=metadata)],
+    )
+
+chat_proto = Protocol(spec=chat_protocol_spec)
+
+@chat_proto.on_message(ChatMessage)
+async def handle_message(ctx: Context, sender: str, msg: ChatMessage):
+    ctx.logger.info(f"Got a message from {sender}")
+    await ctx.send(sender, ChatAcknowledgement(
+        timestamp=datetime.utcnow(), acknowledged_msg_id=msg.msg_id)
+    )
+
+    prompt_content = []
+
+    for item in msg.content:
+        if isinstance(item, StartSessionContent):
+            await ctx.send(sender, create_metadata({"attachments": "true"}))
+
+        elif isinstance(item, TextContent):
+            prompt_content.append({"type": "text", "text": item.text})
+
+        elif isinstance(item, ResourceContent):
+            try:
+                external_storage = ExternalStorage(
+                    identity=ctx.agent.identity,
+                    storage_url=STORAGE_URL,
+                )
+                data = external_storage.download(str(item.resource_id))
+                prompt_content.append({
+                    "type": "resource",
+                    "mime_type": data["mime_type"],
+                    "contents": data["contents"],
+                })
+
+            except Exception as e:
+                ctx.logger.error(f"Failed to download resource: {e}")
+                await ctx.send(sender, create_text_chat("Failed to download resource."))
+                return
+
+        else:
+            ctx.logger.warning(f"Unexpected content type from {sender}")
+
+    if prompt_content:
+        response = get_completion(prompt_content)
+        await ctx.send(sender, create_text_chat(response or "No response generated."))
+
+
+@chat_proto.on_message(ChatAcknowledgement)
+async def handle_ack(ctx: Context, sender: str, msg: ChatAcknowledgement):
+    ctx.logger.info(f"Got an acknowledgement from {sender} for {msg.acknowledged_msg_id}")

6-deployed-agents/knowledge-base/asi1-agent/pyproject.toml

@@ -0,0 +1,14 @@
+[tool.poetry]
+name = "asi1-agent"
+version = "0.1.0"
+description = ""
+authors = ["Attila Bagoly <[email protected]>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.10,<3.13"
+uagents = "^0.20.1"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

6-deployed-agents/knowledge-base/claude-ai-agent/README.md

@@ -80,7 +80,7 @@ class StructuredOutputResponse(Model):
 agent = Agent()
 
 
-AI_AGENT_ADDRESS = "<deployed_agent_address>"
+AI_AGENT_ADDRESS = "{{ .Agent.Address }}"
 
 
 class Location(Model):

6-deployed-agents/knowledge-base/claude-ai-agent/ai.py

@@ -49,13 +49,33 @@ def get_text_completion(prompt: str, tool: dict[str, Any] | None = None) -> str
 def get_completion(
     content: list[dict[str, Any]], tool: dict[str, Any] | None = None
 ) -> str | None:
+
+    processed_content = []
+
+    for item in content:
+        if item.get("type") == "text":
+            processed_content.append({"type": "text", "text": item["text"]})
+        elif item.get("type") == "resource":
+            mime_type = item["mime_type"]
+            if mime_type.startswith("image/"):
+                processed_content.append({
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": mime_type,
+                        "data": item["contents"],
+                    }
+                })
+            else:
+                return f"Unsupported mime type: {mime_type}"
+
     data = {
         "model": MODEL_ENGINE,
         "max_tokens": MAX_TOKENS,
         "messages": [
             {
                 "role": "user",
-                "content": content,
+                "content": processed_content,
             }
         ],
     }

6-deployed-agents/knowledge-base/claude-ai-agent/chat_proto.py

@@ -50,7 +50,7 @@ async def handle_message(ctx: Context, sender: str, msg: ChatMessage):
             timestamp=datetime.utcnow(), acknowledged_msg_id=msg.msg_id
         ),
     )
-    ctx.storage.set(str(ctx.session), sender)
+
     prompt_content = []
     for item in msg.content:
         if isinstance(item, StartSessionContent):
@@ -64,21 +64,12 @@ async def handle_message(ctx: Context, sender: str, msg: ChatMessage):
                     storage_url=STORAGE_URL,
                 )
                 data = external_storage.download(str(item.resource_id))
-                if "image" in data["mime_type"]:
-                    prompt_content.append(
-                        {
-                            "type": "image",
-                            "source": {
-                                "type": "base64",
-                                "media_type": data["mime_type"],
-                                "data": data["contents"],
-                            },
-                        }
-                    )
-                else:
-                    ctx.logger.warning(
-                        f"Got unexpected resource type: {data['mime_type']}"
-                    )
+                prompt_content.append({
+                    "type": "resource",
+                    "mime_type": data["mime_type"],
+                    "contents": data["contents"],
+                })
+
             except Exception as ex:
                 ctx.logger.error(f"Failed to download resource: {ex}")
                 await ctx.send(sender, create_text_chat("Failed to download resource."))

6-deployed-agents/knowledge-base/google-gemini-agent/README.md

@@ -103,7 +103,7 @@ class StructuredOutputResponse(Model):
 agent = Agent()
 
 
-AI_AGENT_ADDRESS = "<deployed_agent_address>"
+AI_AGENT_ADDRESS = "{{ .Agent.Address }}"
 
 
 class Location(Model):

6-deployed-agents/knowledge-base/google-gemini-agent/agent.py

@@ -3,7 +3,7 @@
 from enum import Enum
 from typing import Any
 
-from ai import get_completion
+from ai import get_completion, get_text_completion
 from chat_proto import chat_proto
 from uagents import Agent, Context, Model
 from uagents.experimental.quota import QuotaProtocol, RateLimit, AccessControlList
@@ -56,27 +56,27 @@ class StructuredOutputResponse(Model):
     storage_reference=agent.storage,
     name="LLM-Text-Response",
     version="0.1.0",
-    default_rate_limit=RateLimit(window_size_minutes=60, max_requests=6), acl=acl,
+    default_rate_limit=RateLimit(window_size_minutes=60, max_requests=6, acl=acl),
 )
 
 code_proto = QuotaProtocol(
     storage_reference=agent.storage,
     name="LLM-Code-Generator",
     version="0.1.0",
-    default_rate_limit=RateLimit(window_size_minutes=60, max_requests=6), acl=acl,
+    default_rate_limit=RateLimit(window_size_minutes=60, max_requests=6, acl=acl),
 )
 
 struct_proto = QuotaProtocol(
     storage_reference=agent.storage,
     name="LLM-Structured-Response",
     version="0.1.0",
-    default_rate_limit=RateLimit(window_size_minutes=60, max_requests=6), acl=acl,
+    default_rate_limit=RateLimit(window_size_minutes=60, max_requests=6, acl=acl),
 )
 
 
 @text_proto.on_message(TextPrompt, replies={TextResponse, ErrorMessage})
 async def handle_request(ctx: Context, sender: str, msg: TextPrompt):
-    response = get_completion(msg.text, False)
+    response = get_text_completion(msg.text)
     if response is None:
         await ctx.send(
             sender,
@@ -89,7 +89,7 @@ async def handle_request(ctx: Context, sender: str, msg: TextPrompt):
 
 @code_proto.on_message(CodePrompt, replies={CodeResponse, ErrorMessage})
 async def handle_codegen_request(ctx: Context, sender: str, msg: CodePrompt):
-    response = get_completion(msg.text, True)
+    response = get_text_completion(msg.text, True)
     if response is None:
         await ctx.send(
             sender,
@@ -106,7 +106,7 @@ async def handle_codegen_request(ctx: Context, sender: str, msg: CodePrompt):
 async def handle_structured_request(
     ctx: Context, sender: str, msg: StructuredOutputPrompt
 ):
-    response = get_completion(msg.prompt, False, msg.output_schema)
+    response = get_completion([{"type": "text", "text": msg.prompt}], response_schema=msg.output_schema)
     if response is None:
         await ctx.send(
             sender,