Merge branch 'main' into service-contract

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "uagents"
-version = "0.4.1"
+version = "0.5.0"
 description = "Lightweight framework for rapid agent-based development"
 authors = ["Ed FitzGerald <[email protected]>", "James Riehl <[email protected]>", "Alejandro Morales <[email protected]>"]
 license = "Apache 2.0"

src/uagents/agent.py

@@ -1,6 +1,7 @@
 import asyncio
 import functools
 from typing import Dict, List, Optional, Set, Union, Type, Tuple, Any
+import uuid
 
 from cosmpy.aerial.wallet import LocalWallet, PrivateKey
 from cosmpy.crypto.address import Address
@@ -359,8 +360,10 @@ def include(self, protocol: Protocol):
         if protocol.digest is not None:
             self.protocols[protocol.digest] = protocol
 
-    async def handle_message(self, sender, schema_digest: str, message: JsonStr):
-        await self._message_queue.put((schema_digest, sender, message))
+    async def handle_message(
+        self, sender, schema_digest: str, message: JsonStr, session: uuid.UUID
+    ):
+        await self._message_queue.put((schema_digest, sender, message, session))
 
     async def _startup(self):
         for handler in self._on_startup:
@@ -425,7 +428,7 @@ def run(self):
     async def _process_message_queue(self):
         while True:
             # get an element from the queue
-            schema_digest, sender, message = await self._message_queue.get()
+            schema_digest, sender, message, session = await self._message_queue.get()
 
             # lookup the model definition
             model_class: Model = self._models.get(schema_digest)
@@ -444,6 +447,7 @@ async def _process_message_queue(self):
                 self._wallet,
                 self._ledger,
                 self._queries,
+                session=session,
                 replies=self._replies,
                 interval_messages=self._interval_messages,
                 message_received=MsgDigest(

src/uagents/asgi.py

@@ -11,8 +11,8 @@
 from uagents.crypto import is_user_address
 from uagents.dispatch import dispatcher
 from uagents.envelope import Envelope
-from uagents.models import Model, ErrorMessage
-from uagents.query import enclose_response
+from uagents.models import ErrorMessage
+from uagents.query import enclose_response_raw
 
 
 HOST = "0.0.0.0"
@@ -158,17 +158,19 @@ async def __call__(self, scope, receive, send):
             return
 
         await dispatcher.dispatch(
-            env.sender, env.target, env.schema_digest, env.decode_payload()
+            env.sender, env.target, env.schema_digest, env.decode_payload(), env.session
         )
 
         # wait for any queries to be resolved
         if expects_response:
-            response_msg: Model = await self._queries[env.sender]
+            response_msg, schema_digest = await self._queries[env.sender]
             if env.expires is not None:
                 if datetime.now() > datetime.fromtimestamp(env.expires):
                     response_msg = ErrorMessage(error="Query envelope expired")
             sender = env.target
-            response = enclose_response(response_msg, sender, str(env.session))
+            response = enclose_response_raw(
+                response_msg, schema_digest, sender, str(env.session)
+            )
         else:
             response = "{}"
 

src/uagents/context.py

@@ -12,7 +12,7 @@
 
 from uagents.config import DEFAULT_ENVELOPE_TIMEOUT_SECONDS
 from uagents.crypto import Identity
-from uagents.dispatch import dispatcher
+from uagents.dispatch import JsonStr, dispatcher
 from uagents.envelope import Envelope
 from uagents.models import Model, ErrorMessage
 from uagents.resolver import Resolver
@@ -46,6 +46,7 @@ def __init__(
         wallet: LocalWallet,
         ledger: LedgerClient,
         queries: Dict[str, asyncio.Future],
+        session: Optional[uuid.UUID] = None,
         replies: Optional[Dict[str, Set[Type[Model]]]] = None,
         interval_messages: Optional[Set[str]] = None,
         message_received: Optional[MsgDigest] = None,
@@ -60,6 +61,7 @@ def __init__(
         self._resolver = resolve
         self._identity = identity
         self._queries = queries
+        self._session = session or uuid.uuid4()
         self._replies = replies
         self._interval_messages = interval_messages
         self._message_received = message_received
@@ -84,6 +86,10 @@ def logger(self) -> logging.Logger:
     def protocols(self) -> Optional[Dict[str, Protocol]]:
         return self._protocols
 
+    @property
+    def session(self) -> uuid.UUID:
+        return self._session
+
     def get_message_protocol(self, message_schema_digest) -> Optional[str]:
         for protocol_digest, protocol in self._protocols.items():
             for reply_models in protocol.replies.values():
@@ -97,10 +103,23 @@ async def send(
         message: Model,
         timeout: Optional[int] = DEFAULT_ENVELOPE_TIMEOUT_SECONDS,
     ):
-        # convert the message into object form
-        json_message = message.json()
         schema_digest = Model.build_schema_digest(message)
+        await self.send_raw(
+            destination,
+            message.json(),
+            schema_digest,
+            message_type=type(message),
+            timeout=timeout,
+        )
 
+    async def send_raw(
+        self,
+        destination: str,
+        json_message: JsonStr,
+        schema_digest: str,
+        message_type: Optional[Type[Model]] = None,
+        timeout: Optional[int] = DEFAULT_ENVELOPE_TIMEOUT_SECONDS,
+    ):
         # check if this message is a reply
         if (
             self._message_received is not None
@@ -112,7 +131,7 @@ async def send(
                 # ensure the reply is valid
                 if schema_digest not in self._replies[received.schema_digest]:
                     self._logger.exception(
-                        f"Outgoing message {type(message)} "
+                        f"Outgoing message {message_type or ''} "
                         f"is not a valid reply to {received.message}"
                     )
                     return
@@ -121,20 +140,20 @@ async def send(
         if self._message_received is None and self._interval_messages:
             if schema_digest not in self._interval_messages:
                 self._logger.exception(
-                    f"Outgoing message {type(message)} is not a valid interval message"
+                    f"Outgoing message {message_type} is not a valid interval message"
                 )
                 return
 
         # handle local dispatch of messages
         if dispatcher.contains(destination):
             await dispatcher.dispatch(
-                self.address, destination, schema_digest, json_message
+                self.address, destination, schema_digest, json_message, self._session
             )
             return
 
         # handle queries waiting for a response
         if destination in self._queries:
-            self._queries[destination].set_result(message)
+            self._queries[destination].set_result((json_message, schema_digest))
             del self._queries[destination]
             return
 
@@ -154,7 +173,8 @@ async def send(
             version=1,
             sender=self.address,
             target=destination_address,
-            session=uuid.uuid4(),
+            session=self._session,
+
             schema_digest=schema_digest,
             protocol_digest=self.get_message_protocol(schema_digest),
             expires=expires,

src/uagents/dispatch.py

@@ -1,12 +1,15 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Set
+import uuid
 
 JsonStr = str
 
 
 class Sink(ABC):
     @abstractmethod
-    async def handle_message(self, sender: str, schema_digest: str, message: JsonStr):
+    async def handle_message(
+        self, sender: str, schema_digest: str, message: JsonStr, session: uuid.UUID
+    ):
         pass
 
 
@@ -28,10 +31,15 @@ def contains(self, address: str) -> bool:
         return address in self._sinks
 
     async def dispatch(
-        self, sender: str, destination: str, schema_digest: str, message: JsonStr
+        self,
+        sender: str,
+        destination: str,
+        schema_digest: str,
+        message: JsonStr,
+        session: uuid.UUID,
     ):
         for handler in self._sinks.get(destination, set()):
-            await handler.handle_message(sender, schema_digest, message)
+            await handler.handle_message(sender, schema_digest, message, session)
 
 
 dispatcher = Dispatcher()

src/uagents/mailbox.py

@@ -75,6 +75,7 @@ async def _handle_envelope(self, payload: dict):
             env.target,
             env.schema_digest,
             env.decode_payload(),
+            env.session,
         )
 
         # queue envelope for deletion from server

src/uagents/query.py

@@ -6,6 +6,7 @@
 
 from uagents.config import get_logger
 from uagents.crypto import generate_user_address
+from uagents.dispatch import JsonStr
 from uagents.envelope import Envelope
 from uagents.models import Model
 from uagents.resolver import Resolver, AlmanacResolver
@@ -68,12 +69,19 @@ async def query(
 
 
 def enclose_response(message: Model, sender: str, session: str) -> str:
+    schema_digest = Model.build_schema_digest(message)
+    return enclose_response_raw(message.json(), schema_digest, sender, session)
+
+
+def enclose_response_raw(
+    json_message: JsonStr, schema_digest: str, sender: str, session: str
+) -> str:
     response_env = Envelope(
         version=1,
         sender=sender,
         target="",
         session=session,
-        schema_digest=Model.build_schema_digest(message),
+        schema_digest=schema_digest,
     )
-    response_env.encode_payload(message.json())
+    response_env.encode_payload(json_message)
     return response_env.json()

tests/test_server.py

@@ -23,7 +23,9 @@ def setUp(self) -> None:
     async def mock_process_sync_message(self, sender: str, msg: Model):
         while True:
             if sender in self.agent._server._queries:
-                self.agent._server._queries[sender].set_result(msg)
+                self.agent._server._queries[sender].set_result(
+                    (msg.json(), Model.build_schema_digest(msg))
+                )
                 return
 
     async def test_message_success(self):