From 48299913c598c1feb2ea7cd8986492515f532718 Mon Sep 17 00:00:00 2001
From: Xuchen Pan <32844285+pan-x-c@users.noreply.github.com>
Date: Mon, 22 Jan 2024 16:07:37 +0800
Subject: [PATCH] fix rpc import error (#14)

Fixed the bug that RPC related modules throw ModuleNotFoundError under default installation conditions
---
 src/agentscope/agents/rpc_agent.py     |  6 +++---
 src/agentscope/message.py              |  2 +-
 src/agentscope/rpc/__init__.py         | 26 ++++++++++++++++++++++++++
 src/agentscope/rpc/rpc_agent_client.py | 10 ++++++++--
 4 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/src/agentscope/agents/rpc_agent.py b/src/agentscope/agents/rpc_agent.py
index b441facea..a2aeff030 100644
--- a/src/agentscope/agents/rpc_agent.py
+++ b/src/agentscope/agents/rpc_agent.py
@@ -32,15 +32,15 @@
 from agentscope._init import _INIT_SETTINGS
 from agentscope._init import init
 from agentscope.agents.agent import AgentBase
-from agentscope.rpc.rpc_agent_pb2 import RpcMsg  # pylint: disable=E0611
 from agentscope.message import MessageBase
 from agentscope.message import Msg
 from agentscope.message import PlaceholderMessage
 from agentscope.message import deserialize
 from agentscope.message import serialize
 from agentscope.utils.logging_utils import logger
-from agentscope.rpc.rpc_agent_client import RpcAgentClient
-from agentscope.rpc.rpc_agent_pb2_grpc import (
+from agentscope.rpc import (
+    RpcAgentClient,
+    RpcMsg,
     RpcAgentServicer,
     add_RpcAgentServicer_to_server,
 )
diff --git a/src/agentscope/message.py b/src/agentscope/message.py
index 2057a3308..8aa675f68 100644
--- a/src/agentscope/message.py
+++ b/src/agentscope/message.py
@@ -7,7 +7,7 @@
 
 from loguru import logger
 
-from .rpc.rpc_agent_client import RpcAgentClient
+from .rpc import RpcAgentClient
 from .utils.tools import _get_timestamp
 
 
diff --git a/src/agentscope/rpc/__init__.py b/src/agentscope/rpc/__init__.py
index e69de29bb..0a5ffe826 100644
--- a/src/agentscope/rpc/__init__.py
+++ b/src/agentscope/rpc/__init__.py
@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+"""Import all rpc related modules in the package."""
+from typing import Any
+from .rpc_agent_client import RpcAgentClient
+
+try:
+    from .rpc_agent_pb2 import RpcMsg  # pylint: disable=E0611
+except ModuleNotFoundError:
+    RpcMsg = Any
+try:
+    from .rpc_agent_pb2_grpc import RpcAgentServicer
+    from .rpc_agent_pb2_grpc import RpcAgentStub
+    from .rpc_agent_pb2_grpc import add_RpcAgentServicer_to_server
+except ImportError:
+    RpcAgentServicer = object
+    RpcAgentStub = Any
+    add_RpcAgentServicer_to_server = Any
+
+
+__all__ = [
+    "RpcAgentClient",
+    "RpcMsg",
+    "RpcAgentServicer",
+    "RpcAgentStub",
+    "add_RpcAgentServicer_to_server",
+]
diff --git a/src/agentscope/rpc/rpc_agent_client.py b/src/agentscope/rpc/rpc_agent_client.py
index ab0d63a0c..e1f5f3937 100644
--- a/src/agentscope/rpc/rpc_agent_client.py
+++ b/src/agentscope/rpc/rpc_agent_client.py
@@ -1,13 +1,19 @@
 # -*- coding: utf-8 -*-
 """ Client of rpc agent server """
 
+from typing import Any
+
 try:
     import grpc
 except ImportError:
     grpc = None
 
-from agentscope.rpc.rpc_agent_pb2 import RpcMsg  # pylint: disable=E0611
-from agentscope.rpc.rpc_agent_pb2_grpc import RpcAgentStub
+try:
+    from agentscope.rpc.rpc_agent_pb2 import RpcMsg  # pylint: disable=E0611
+    from agentscope.rpc.rpc_agent_pb2_grpc import RpcAgentStub
+except ModuleNotFoundError:
+    RpcMsg = Any
+    RpcAgentStub = Any
 
 
 class RpcAgentClient: