Repository: run-llama/llama_deploy

Issue:
Make register-consumer-remote idempotent
Currently registering a service via SimpleRemoteMessageQueueClient multiple times leads to a service being registered as a consumer more than once. However, this should only happen once.


Discussion:

Patch:
diff --git a/llama_agents/message_queues/simple.py b/llama_agents/message_queues/simple.py
index ca7ef53d..3efa678e 100644
--- a/llama_agents/message_queues/simple.py
+++ b/llama_agents/message_queues/simple.py
@@ -7,7 +7,7 @@
 
 from collections import deque
 from contextlib import asynccontextmanager
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, status
 from logging import getLogger
 from pydantic import Field, PrivateAttr
 from typing import Any, AsyncGenerator, Dict, List, Optional
@@ -218,8 +218,10 @@ async def register_consumer(
             )
         else:
             if consumer.id_ in self.consumers[message_type_str]:
-                raise ValueError("Consumer has already been added.")
-
+                raise HTTPException(
+                    status_code=status.HTTP_409_CONFLICT,
+                    detail="Consumer with the same id_ has already been previously added.",
+                )
             self.consumers[message_type_str][consumer.id_] = consumer
             logger.info(
                 f"Consumer {consumer.id_}: {message_type_str} has been registered."
@@ -232,14 +234,43 @@ async def register_remote_consumer(
         self, consumer_def: RemoteMessageConsumerDef
     ) -> Dict[str, str]:
         consumer = RemoteMessageConsumer(**consumer_def.model_dump())
+        message_type = consumer.message_type
+
+        # check if consumer with same url already exists
+        if message_type in self.consumers:
+
+            def consumer_with_same_url(
+                c: RemoteMessageConsumer,
+            ) -> bool:
+                return c.url == consumer.url
+
+            try:
+                next(
+                    filter(
+                        consumer_with_same_url,
+                        [
+                            c
+                            for c in self.consumers[message_type].values()
+                            if isinstance(c, RemoteMessageConsumer)
+                        ],
+                    )
+                )
+                raise HTTPException(
+                    status_code=status.HTTP_409_CONFLICT,
+                    detail="A consumer with the same url has previously been registered.",
+                )
+            except StopIteration:
+                pass
+
         await self.register_consumer(consumer)
         return {"consumer": consumer.id_}
 
     async def deregister_consumer(self, consumer: BaseMessageQueueConsumer) -> None:
         message_type_str = consumer.message_type
         if consumer.id_ not in self.consumers.get(message_type_str, {}):
-            raise ValueError(
-                f"No consumer found for associated message type. {consumer.id_}: {message_type_str}"
+            raise HTTPException(
+                detail=f"No consumer found for associated message type. {consumer.id_}: {message_type_str}",
+                status_code=status.HTTP_404_NOT_FOUND,
             )
 
         del self.consumers[message_type_str][consumer.id_]
@@ -310,3 +341,8 @@ async def home(self) -> Dict[str, str]:
             "service_name": "message_queue",
             "description": "Message queue for multi-agent system",
         }
+
+
+if __name__ == "__main__":
+    mq = SimpleMessageQueue()
+    asyncio.run(mq.launch_server())


Test patch:
diff --git a/tests/message_queues/test_simple.py b/tests/message_queues/test_simple.py
index 2fa7a0e3..97fda2cb 100644
--- a/tests/message_queues/test_simple.py
+++ b/tests/message_queues/test_simple.py
@@ -1,5 +1,6 @@
 import asyncio
 import pytest
+from fastapi import HTTPException
 from pydantic import PrivateAttr
 from typing import Any, List
 
@@ -27,7 +28,7 @@ async def test_simple_register_consumer() -> None:
     # Act
     await mq.register_consumer(consumer_one)
     await mq.register_consumer(consumer_two)
-    with pytest.raises(ValueError):
+    with pytest.raises(HTTPException):
         await mq.register_consumer(consumer_two)
 
     # Assert
@@ -54,7 +55,7 @@ async def test_simple_deregister_consumer() -> None:
     # Act
     await mq.deregister_consumer(consumer_one)
     await mq.deregister_consumer(consumer_three)
-    with pytest.raises(ValueError):
+    with pytest.raises(HTTPException):
         await mq.deregister_consumer(consumer_three)
 
     # Assert