feat: add storage support to claude and image generator agents (#74)

Author: jrriehl

6-deployed-agents/knowledge-base/claude.ai-agent/.avctl/config.toml renamed to 6-deployed-agents/knowledge-base/claude-ai-agent/.avctl/config.toml

@@ -2,12 +2,13 @@
 from enum import Enum
 from typing import Any
 
-from ai import get_completion, get_structured_response
-from chat_proto import chat_proto
 from uagents import Agent, Context, Model
 from uagents.experimental.quota import AccessControlList, QuotaProtocol, RateLimit
 from uagents_core.models import ErrorMessage
 
+from ai import get_structured_response, get_text_completion
+from chat_proto import chat_proto
+
 AGENT_SEED = os.getenv("AGENT_SEED", "claude-test-agent")
 AGENT_NAME = os.getenv("AGENT_NAME", "Claude.ai Agent")
 BYPASS_RATE_LIMIT = set(
@@ -62,7 +63,7 @@ class StructuredOutputResponse(Model):
 
 @proto.on_message(TextPrompt, replies={TextResponse, ErrorMessage})
 async def handle_request(ctx: Context, sender: str, msg: TextPrompt):
-    response = get_completion(msg.text)
+    response = get_text_completion(msg.text)
     if response is None:
         await ctx.send(
             sender,

6-deployed-agents/knowledge-base/claude.ai-agent/.secrets.template renamed to 6-deployed-agents/knowledge-base/claude-ai-agent/.secrets.template

@@ -11,7 +11,7 @@
     raise ValueError(
         "You need to provide an API key: https://platform.openai.com/api-keys"
     )
-MODEL_ENGINE = os.getenv("MODEL_ENGINE", "claude-3-haiku-20240307")
+MODEL_ENGINE = os.getenv("MODEL_ENGINE", "claude-3-5-haiku-latest")
 HEADERS = {
     "x-api-key": ANTHROPIC_API_KEY,
     "anthropic-version": "2023-06-01",
@@ -25,7 +25,7 @@
 
 
 def create_structured_response_tool(
-    response_model_schema: dict[str, Any]
+    response_model_schema: dict[str, Any],
 ) -> dict[str, Any]:
 
     # Exclude the title from the schema (not allowed in the API)
@@ -40,13 +40,26 @@ def create_structured_response_tool(
     }
 
 
+def get_text_completion(prompt: str, tool: dict[str, Any] | None = None) -> str | None:
+    content = [{"type": "text", "text": prompt}]
+    return get_completion(content, tool)
+
+
 # Send a prompt to the AI model and return the content of the completion
-def get_completion(prompt: str, tool: dict[str, Any] | None = None) -> str | None:
+def get_completion(
+    content: list[dict[str, Any]], tool: dict[str, Any] | None = None
+) -> str | None:
     data = {
         "model": MODEL_ENGINE,
         "max_tokens": MAX_TOKENS,
-        "messages": [{"role": "user", "content": prompt}],
+        "messages": [
+            {
+                "role": "user",
+                "content": content,
+            }
+        ],
     }
+
     if tool:
         data["tools"] = [tool]
         data["tool_choice"] = {"type": "tool", "name": tool["name"]}
@@ -55,23 +68,34 @@ def get_completion(prompt: str, tool: dict[str, Any] | None = None) -> str | Non
         response = requests.post(
             CLAUDE_URL, headers=HEADERS, data=json.dumps(data), timeout=120
         )
+        response.raise_for_status()
     except requests.exceptions.Timeout:
         return "The request timed out. Please try again."
     except requests.exceptions.RequestException as e:
         return f"An error occurred: {e}"
 
+    # Check if the response was successful
+    response_data = response.json()
+
+    # Handle error responses
+    if "error" in response_data:
+        return f"API Error: {response_data['error'].get('message', 'Unknown error')}"
+
     if tool:
-        for item in response.json()["content"]:
+        for item in response_data["content"]:
             if item["type"] == "tool_use":
                 return item["input"]
+            
+    messages = response_data["content"]
 
-    messages = response.json()["content"]
-    message = messages[0]["text"]
-    return message
+    if messages:
+        return messages[0]["text"]
+    else:
+        return None
 
 
 def get_structured_response(
     prompt: str, response_model_schema: dict[str, Any]
 ) -> dict[str, Any] | None:
     tool = create_structured_response_tool(response_model_schema)
-    return get_completion(prompt, tool)
+    return get_text_completion(prompt, tool)

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

@@ -0,0 +1,97 @@
+import os
+from datetime import datetime
+from uuid import uuid4
+
+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
+
+from ai import get_completion
+
+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
+        ),
+    )
+    ctx.storage.set(str(ctx.session), sender)
+    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({"text": item.text, "type": "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))
+                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']}"
+                    )
+            except Exception as ex:
+                ctx.logger.error(f"Failed to download resource: {ex}")
+                await ctx.send(sender, create_text_chat("Failed to download resource."))
+        else:
+            ctx.logger.warning(f"Got unexpected content from {sender}")
+
+    if prompt_content:
+        response = get_completion(prompt_content)
+        await ctx.send(sender, create_text_chat(response))
+
+
+@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/claude.ai-agent/agent.py renamed to 6-deployed-agents/knowledge-base/claude-ai-agent/agent.py

@@ -8,7 +8,8 @@ packages = [{ include = "claude" }]
 
 [tool.poetry.dependencies]
 python = "^3.10,<3.13"
-uagents = "^0.22.0"
+uagents = "^0.22.4"
+uagents-core = "^0.3.1"
 requests = "^2.32.3"
 
 

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

@@ -0,0 +1,98 @@
+import os
+from enum import Enum
+
+from uagents import Agent, Context, Model
+from uagents.experimental.quota import QuotaProtocol, RateLimit
+from uagents_core.models import ErrorMessage
+
+from chat_proto import chat_proto
+from models import ImageRequest, ImageResponse, generate_image
+
+AGENT_SEED = os.getenv("AGENT_SEED", "image-generator-agent")
+AGENT_NAME = os.getenv("AGENT_NAME", "Image Generator Agent")
+
+PORT = 8000
+agent = Agent(
+    name=AGENT_NAME,
+    seed=AGENT_SEED,
+    port=PORT,
+    endpoint=f"http://localhost:{PORT}/submit",
+)
+
+proto = QuotaProtocol(
+    storage_reference=agent.storage,
+    name="Image-Generator",
+    version="0.1.0",
+    default_rate_limit=RateLimit(window_size_minutes=60, max_requests=6),
+)
+
+@proto.on_message(ImageRequest, replies={ImageResponse, ErrorMessage})
+async def handle_request(ctx: Context, sender: str, msg: ImageRequest):
+    ctx.logger.info(f"Received Image request")
+    try:
+        image_url = generate_image(msg.image_description)
+
+    except Exception as err:
+        ctx.logger.error(err)
+        await ctx.send(
+            sender,
+            ErrorMessage(
+                error="An error occurred while processing the request. Please try again later."
+            ),
+        )
+        return
+    await ctx.send(sender, ImageResponse(image_url=image_url))
+
+
+agent.include(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)
+agent.include(chat_proto, publish_manifest=True)
+
+
+if __name__ == "__main__":
+    agent.run()