agentscope-runtime 0.1.5b1__py3-none-any.whl → 0.2.0b1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (106) hide show
  1. agentscope_runtime/common/__init__.py +0 -0
  2. agentscope_runtime/common/collections/in_memory_mapping.py +27 -0
  3. agentscope_runtime/common/collections/redis_mapping.py +42 -0
  4. agentscope_runtime/common/container_clients/__init__.py +0 -0
  5. agentscope_runtime/common/container_clients/agentrun_client.py +1098 -0
  6. agentscope_runtime/common/container_clients/docker_client.py +250 -0
  7. agentscope_runtime/engine/__init__.py +12 -0
  8. agentscope_runtime/engine/agents/agentscope_agent.py +488 -0
  9. agentscope_runtime/engine/agents/agno_agent.py +19 -18
  10. agentscope_runtime/engine/agents/autogen_agent.py +13 -8
  11. agentscope_runtime/engine/agents/utils.py +53 -0
  12. agentscope_runtime/engine/app/__init__.py +6 -0
  13. agentscope_runtime/engine/app/agent_app.py +239 -0
  14. agentscope_runtime/engine/app/base_app.py +181 -0
  15. agentscope_runtime/engine/app/celery_mixin.py +92 -0
  16. agentscope_runtime/engine/deployers/base.py +1 -0
  17. agentscope_runtime/engine/deployers/cli_fc_deploy.py +72 -12
  18. agentscope_runtime/engine/deployers/kubernetes_deployer.py +12 -5
  19. agentscope_runtime/engine/deployers/local_deployer.py +61 -3
  20. agentscope_runtime/engine/deployers/modelstudio_deployer.py +77 -27
  21. agentscope_runtime/engine/deployers/utils/docker_image_utils/docker_image_builder.py +3 -3
  22. agentscope_runtime/engine/deployers/utils/docker_image_utils/runner_image_factory.py +9 -0
  23. agentscope_runtime/engine/deployers/utils/package_project_utils.py +234 -3
  24. agentscope_runtime/engine/deployers/utils/service_utils/fastapi_factory.py +567 -7
  25. agentscope_runtime/engine/deployers/utils/service_utils/standalone_main.py.j2 +211 -0
  26. agentscope_runtime/engine/deployers/utils/wheel_packager.py +1 -1
  27. agentscope_runtime/engine/helpers/helper.py +60 -41
  28. agentscope_runtime/engine/runner.py +35 -24
  29. agentscope_runtime/engine/schemas/agent_schemas.py +42 -0
  30. agentscope_runtime/engine/schemas/modelstudio_llm.py +14 -14
  31. agentscope_runtime/engine/services/sandbox_service.py +62 -70
  32. agentscope_runtime/engine/services/tablestore_memory_service.py +304 -0
  33. agentscope_runtime/engine/services/tablestore_rag_service.py +143 -0
  34. agentscope_runtime/engine/services/tablestore_session_history_service.py +293 -0
  35. agentscope_runtime/engine/services/utils/__init__.py +0 -0
  36. agentscope_runtime/engine/services/utils/tablestore_service_utils.py +352 -0
  37. agentscope_runtime/engine/tracing/__init__.py +9 -3
  38. agentscope_runtime/engine/tracing/asyncio_util.py +24 -0
  39. agentscope_runtime/engine/tracing/base.py +66 -34
  40. agentscope_runtime/engine/tracing/local_logging_handler.py +45 -31
  41. agentscope_runtime/engine/tracing/message_util.py +528 -0
  42. agentscope_runtime/engine/tracing/tracing_metric.py +20 -8
  43. agentscope_runtime/engine/tracing/tracing_util.py +130 -0
  44. agentscope_runtime/engine/tracing/wrapper.py +794 -169
  45. agentscope_runtime/sandbox/__init__.py +2 -0
  46. agentscope_runtime/sandbox/box/base/__init__.py +4 -0
  47. agentscope_runtime/sandbox/box/base/base_sandbox.py +6 -4
  48. agentscope_runtime/sandbox/box/browser/__init__.py +4 -0
  49. agentscope_runtime/sandbox/box/browser/browser_sandbox.py +10 -14
  50. agentscope_runtime/sandbox/box/dummy/__init__.py +4 -0
  51. agentscope_runtime/sandbox/box/dummy/dummy_sandbox.py +2 -1
  52. agentscope_runtime/sandbox/box/filesystem/__init__.py +4 -0
  53. agentscope_runtime/sandbox/box/filesystem/filesystem_sandbox.py +10 -7
  54. agentscope_runtime/sandbox/box/gui/__init__.py +4 -0
  55. agentscope_runtime/sandbox/box/gui/box/__init__.py +0 -0
  56. agentscope_runtime/sandbox/box/gui/gui_sandbox.py +81 -0
  57. agentscope_runtime/sandbox/box/sandbox.py +5 -2
  58. agentscope_runtime/sandbox/box/shared/routers/generic.py +20 -1
  59. agentscope_runtime/sandbox/box/training_box/__init__.py +4 -0
  60. agentscope_runtime/sandbox/box/training_box/training_box.py +10 -15
  61. agentscope_runtime/sandbox/build.py +143 -58
  62. agentscope_runtime/sandbox/client/http_client.py +87 -59
  63. agentscope_runtime/sandbox/client/training_client.py +0 -1
  64. agentscope_runtime/sandbox/constant.py +27 -1
  65. agentscope_runtime/sandbox/custom/custom_sandbox.py +7 -6
  66. agentscope_runtime/sandbox/custom/example.py +4 -3
  67. agentscope_runtime/sandbox/enums.py +1 -0
  68. agentscope_runtime/sandbox/manager/sandbox_manager.py +212 -106
  69. agentscope_runtime/sandbox/manager/server/app.py +82 -14
  70. agentscope_runtime/sandbox/manager/server/config.py +50 -3
  71. agentscope_runtime/sandbox/model/container.py +12 -23
  72. agentscope_runtime/sandbox/model/manager_config.py +93 -5
  73. agentscope_runtime/sandbox/registry.py +1 -1
  74. agentscope_runtime/sandbox/tools/gui/__init__.py +7 -0
  75. agentscope_runtime/sandbox/tools/gui/tool.py +77 -0
  76. agentscope_runtime/sandbox/tools/mcp_tool.py +6 -2
  77. agentscope_runtime/sandbox/tools/tool.py +4 -0
  78. agentscope_runtime/sandbox/utils.py +124 -0
  79. agentscope_runtime/version.py +1 -1
  80. {agentscope_runtime-0.1.5b1.dist-info → agentscope_runtime-0.2.0b1.dist-info}/METADATA +209 -101
  81. {agentscope_runtime-0.1.5b1.dist-info → agentscope_runtime-0.2.0b1.dist-info}/RECORD +95 -79
  82. agentscope_runtime/engine/agents/agentscope_agent/__init__.py +0 -6
  83. agentscope_runtime/engine/agents/agentscope_agent/agent.py +0 -401
  84. agentscope_runtime/engine/agents/agentscope_agent/hooks.py +0 -169
  85. agentscope_runtime/engine/agents/llm_agent.py +0 -51
  86. agentscope_runtime/engine/llms/__init__.py +0 -3
  87. agentscope_runtime/engine/llms/base_llm.py +0 -60
  88. agentscope_runtime/engine/llms/qwen_llm.py +0 -47
  89. agentscope_runtime/sandbox/manager/collections/in_memory_mapping.py +0 -22
  90. agentscope_runtime/sandbox/manager/collections/redis_mapping.py +0 -26
  91. agentscope_runtime/sandbox/manager/container_clients/__init__.py +0 -10
  92. agentscope_runtime/sandbox/manager/container_clients/docker_client.py +0 -422
  93. /agentscope_runtime/{sandbox/manager → common}/collections/__init__.py +0 -0
  94. /agentscope_runtime/{sandbox/manager → common}/collections/base_mapping.py +0 -0
  95. /agentscope_runtime/{sandbox/manager → common}/collections/base_queue.py +0 -0
  96. /agentscope_runtime/{sandbox/manager → common}/collections/base_set.py +0 -0
  97. /agentscope_runtime/{sandbox/manager → common}/collections/in_memory_queue.py +0 -0
  98. /agentscope_runtime/{sandbox/manager → common}/collections/in_memory_set.py +0 -0
  99. /agentscope_runtime/{sandbox/manager → common}/collections/redis_queue.py +0 -0
  100. /agentscope_runtime/{sandbox/manager → common}/collections/redis_set.py +0 -0
  101. /agentscope_runtime/{sandbox/manager → common}/container_clients/base_client.py +0 -0
  102. /agentscope_runtime/{sandbox/manager → common}/container_clients/kubernetes_client.py +0 -0
  103. {agentscope_runtime-0.1.5b1.dist-info → agentscope_runtime-0.2.0b1.dist-info}/WHEEL +0 -0
  104. {agentscope_runtime-0.1.5b1.dist-info → agentscope_runtime-0.2.0b1.dist-info}/entry_points.txt +0 -0
  105. {agentscope_runtime-0.1.5b1.dist-info → agentscope_runtime-0.2.0b1.dist-info}/licenses/LICENSE +0 -0
  106. {agentscope_runtime-0.1.5b1.dist-info → agentscope_runtime-0.2.0b1.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,239 @@
1
+ # -*- coding: utf-8 -*-
2
+ import asyncio
3
+ import logging
4
+ from contextlib import asynccontextmanager
5
+ from typing import Optional, Any, Callable, List
6
+
7
+ import uvicorn
8
+ from fastapi import FastAPI
9
+ from pydantic import BaseModel
10
+
11
+ from .base_app import BaseApp
12
+ from ..agents.base_agent import Agent
13
+ from ..deployers.adapter.a2a import A2AFastAPIDefaultAdapter
14
+ from ..deployers.adapter.responses.response_api_protocol_adapter import (
15
+ ResponseAPIDefaultAdapter,
16
+ )
17
+ from ..deployers.utils.deployment_modes import DeploymentMode
18
+ from ..deployers.utils.service_utils.fastapi_factory import FastAPIAppFactory
19
+ from ..deployers.utils.service_utils.service_config import (
20
+ DEFAULT_SERVICES_CONFIG,
21
+ )
22
+ from ..runner import Runner
23
+ from ..schemas.agent_schemas import AgentRequest
24
+ from ..services.context_manager import ContextManager
25
+ from ..services.environment_manager import EnvironmentManager
26
+ from ...version import __version__
27
+
28
+ logger = logging.getLogger(__name__)
29
+
30
+
31
+ class AgentApp(BaseApp):
32
+ """
33
+ The AgentApp class represents an application that runs as an agent.
34
+ """
35
+
36
+ def __init__(
37
+ self,
38
+ *,
39
+ agent: Optional[Agent] = None,
40
+ environment_manager: Optional[EnvironmentManager] = None,
41
+ context_manager: Optional[ContextManager] = None,
42
+ endpoint_path: str = "/process",
43
+ response_type: str = "sse",
44
+ stream: bool = True,
45
+ request_model: Optional[type[BaseModel]] = AgentRequest,
46
+ before_start: Optional[Callable] = None,
47
+ after_finish: Optional[Callable] = None,
48
+ broker_url: Optional[str] = None,
49
+ backend_url: Optional[str] = None,
50
+ **kwargs,
51
+ ):
52
+ """
53
+ Initialize the AgentApp.
54
+
55
+ Args:
56
+ *args: Variable length argument list.
57
+ **kwargs: Arbitrary keyword arguments.
58
+ """
59
+
60
+ self.endpoint_path = endpoint_path
61
+ self.response_type = response_type
62
+ self.stream = stream
63
+ self.request_model = request_model
64
+ self.before_start = before_start
65
+ self.after_finish = after_finish
66
+ self.broker_url = broker_url
67
+ self.backend_url = backend_url
68
+
69
+ self._agent = agent
70
+ self._runner = None
71
+ self.custom_endpoints = [] # Store custom endpoints
72
+
73
+ a2a_protocol = A2AFastAPIDefaultAdapter(agent=self._agent)
74
+ response_protocol = ResponseAPIDefaultAdapter()
75
+ self.protocol_adapters = [a2a_protocol, response_protocol]
76
+
77
+ if self._agent:
78
+ self._runner = Runner(
79
+ agent=self._agent,
80
+ environment_manager=environment_manager,
81
+ context_manager=context_manager,
82
+ )
83
+
84
+ @asynccontextmanager
85
+ async def lifespan(app: FastAPI) -> Any:
86
+ """Manage the application lifespan."""
87
+ if hasattr(self, "before_start") and self.before_start:
88
+ if asyncio.iscoroutinefunction(self.before_start):
89
+ await self.before_start(app, **getattr(self, "kwargs", {}))
90
+ else:
91
+ self.before_start(app, **getattr(self, "kwargs", {}))
92
+ yield
93
+ if hasattr(self, "after_finish") and self.after_finish:
94
+ if asyncio.iscoroutinefunction(self.after_finish):
95
+ await self.after_finish(app, **getattr(self, "kwargs", {}))
96
+ else:
97
+ self.after_finish(app, **getattr(self, "kwargs", {}))
98
+
99
+ kwargs = {
100
+ "title": "Agent Service",
101
+ "version": __version__,
102
+ "description": "Production-ready Agent Service API",
103
+ "lifespan": lifespan,
104
+ **kwargs,
105
+ }
106
+
107
+ if self._runner:
108
+ if self.stream:
109
+ self.func = self._runner.stream_query
110
+ else:
111
+ self.func = self._runner.query
112
+
113
+ super().__init__(
114
+ broker_url=broker_url,
115
+ backend_url=backend_url,
116
+ **kwargs,
117
+ )
118
+
119
+ # Store custom endpoints and tasks for deployment
120
+ # but don't add them to FastAPI here - let FastAPIAppFactory handle it
121
+
122
+ def run(
123
+ self,
124
+ host="0.0.0.0",
125
+ port=8090,
126
+ embed_task_processor=False,
127
+ services_config=None,
128
+ **kwargs,
129
+ ):
130
+ """
131
+ Run the AgentApp using FastAPIAppFactory directly.
132
+
133
+ Args:
134
+ host: Host to bind to
135
+ port: Port to bind to
136
+ embed_task_processor: Whether to embed task processor
137
+ services_config: Optional services configuration
138
+ **kwargs: Additional keyword arguments
139
+ """
140
+
141
+ try:
142
+ logger.info(
143
+ "[AgentApp] Starting AgentApp with FastAPIAppFactory...",
144
+ )
145
+
146
+ # Use default services config if not provided
147
+ if services_config is None:
148
+ services_config = DEFAULT_SERVICES_CONFIG
149
+
150
+ # Create FastAPI application using the factory
151
+ fastapi_app = FastAPIAppFactory.create_app(
152
+ runner=self._runner,
153
+ endpoint_path=self.endpoint_path,
154
+ request_model=self.request_model,
155
+ response_type=self.response_type,
156
+ stream=self.stream,
157
+ before_start=self.before_start,
158
+ after_finish=self.after_finish,
159
+ mode=DeploymentMode.DAEMON_THREAD,
160
+ services_config=services_config,
161
+ protocol_adapters=self.protocol_adapters,
162
+ custom_endpoints=self.custom_endpoints,
163
+ broker_url=self.broker_url,
164
+ backend_url=self.backend_url,
165
+ enable_embedded_worker=embed_task_processor,
166
+ **kwargs,
167
+ )
168
+
169
+ logger.info(f"[AgentApp] Starting server on {host}:{port}")
170
+
171
+ # Start the FastAPI application with uvicorn
172
+ uvicorn.run(
173
+ fastapi_app,
174
+ host=host,
175
+ port=port,
176
+ log_level="info",
177
+ access_log=True,
178
+ )
179
+
180
+ except Exception as e:
181
+ logger.error(f"[AgentApp] Error while running: {e}")
182
+ raise
183
+
184
+ async def deploy(self, deployer, **kwargs):
185
+ """Deploy the agent app with custom endpoints support"""
186
+ # Pass custom endpoints and tasks to the deployer
187
+
188
+ deploy_kwargs = {
189
+ **kwargs,
190
+ "custom_endpoints": self.custom_endpoints,
191
+ "agent": self._agent,
192
+ "runner": self._runner,
193
+ "endpoint_path": self.endpoint_path,
194
+ "stream": self.stream,
195
+ "protocol_adapters": self.protocol_adapters,
196
+ }
197
+ return await deployer.deploy(**deploy_kwargs)
198
+
199
+ def endpoint(self, path: str, methods: Optional[List[str]] = None):
200
+ """Decorator to register custom endpoints"""
201
+
202
+ if methods is None:
203
+ methods = ["POST"]
204
+
205
+ def decorator(func: Callable):
206
+ endpoint_info = {
207
+ "path": path,
208
+ "handler": func,
209
+ "methods": methods,
210
+ "module": getattr(func, "__module__", None),
211
+ "function_name": getattr(func, "__name__", None),
212
+ }
213
+ self.custom_endpoints.append(endpoint_info)
214
+ return func
215
+
216
+ return decorator
217
+
218
+ def task(self, path: str, queue: str = "default"):
219
+ """Decorator to register custom task endpoints"""
220
+
221
+ def decorator(func: Callable):
222
+ # Store task configuration for FastAPIAppFactory to handle
223
+ task_info = {
224
+ "path": path,
225
+ "handler": func, # Store original function
226
+ "methods": ["POST"],
227
+ "module": getattr(func, "__module__", None),
228
+ "function_name": getattr(func, "__name__", None),
229
+ "queue": queue,
230
+ "task_type": True, # Mark as task endpoint
231
+ "original_func": func,
232
+ }
233
+ self.custom_endpoints.append(
234
+ task_info,
235
+ ) # Add to endpoints for deployment
236
+
237
+ return func
238
+
239
+ return decorator
@@ -0,0 +1,181 @@
1
+ # -*- coding: utf-8 -*-
2
+ import inspect
3
+ import logging
4
+ import threading
5
+ from typing import Callable, Optional
6
+
7
+ import uvicorn
8
+ from fastapi import FastAPI, Request
9
+ from fastapi.responses import StreamingResponse
10
+
11
+ from .celery_mixin import CeleryMixin
12
+
13
+ logger = logging.getLogger(__name__)
14
+
15
+
16
+ class BaseApp(FastAPI, CeleryMixin):
17
+ """
18
+ BaseApp extends FastAPI and integrates with Celery
19
+ for asynchronous background task execution.
20
+ """
21
+
22
+ def __init__(
23
+ self,
24
+ broker_url: Optional[str] = None,
25
+ backend_url: Optional[str] = None,
26
+ **kwargs,
27
+ ):
28
+ # Initialize CeleryMixin
29
+ CeleryMixin.__init__(self, broker_url, backend_url)
30
+
31
+ self.server = None
32
+
33
+ # Initialize FastAPI
34
+ FastAPI.__init__(self, **kwargs)
35
+
36
+ def task(self, path: str, queue: str = "celery"):
37
+ """
38
+ Register an asynchronous task endpoint.
39
+ POST <path> -> Create a task and return task ID
40
+ GET <path>/{task_id} -> Check the task status and result
41
+ Combines Celery and FastAPI routing functionality.
42
+ """
43
+ if self.celery_app is None:
44
+ raise RuntimeError(
45
+ f"[AgentApp] Cannot register task endpoint '{path}'.\n"
46
+ f"Reason: The @task decorator requires a background task "
47
+ f"queue to run asynchronous jobs.\n\n"
48
+ "If you want to use async task queue, you must initialize "
49
+ "AgentApp with broker_url and backend_url, e.g.: \n\n"
50
+ " app = AgentApp(\n"
51
+ " broker_url='redis://localhost:6379/0',\n"
52
+ " backend_url='redis://localhost:6379/0'\n"
53
+ " )\n",
54
+ )
55
+
56
+ def decorator(func: Callable):
57
+ # Register Celery task using CeleryMixin
58
+ celery_task = self.register_celery_task(func, queue=queue)
59
+
60
+ # Add FastAPI HTTP routes
61
+ @self.post(path)
62
+ async def create_task(request: Request):
63
+ if len(inspect.signature(func).parameters) > 0:
64
+ body = await request.json()
65
+ task = celery_task.delay(body)
66
+ else:
67
+ task = celery_task.delay()
68
+ return {"task_id": task.id}
69
+
70
+ @self.get(path + "/{task_id}")
71
+ async def get_task(task_id: str):
72
+ return self.get_task_status(task_id)
73
+
74
+ return func
75
+
76
+ return decorator
77
+
78
+ def endpoint(self, path: str):
79
+ """
80
+ Unified POST endpoint decorator.
81
+ Pure FastAPI functionality, independent of Celery.
82
+ Supports:
83
+ - Sync functions
84
+ - Async functions (coroutines)
85
+ - Sync/async generator functions (streaming responses)
86
+ """
87
+
88
+ def decorator(func: Callable):
89
+ is_async_gen = inspect.isasyncgenfunction(func)
90
+ is_sync_gen = inspect.isgeneratorfunction(func)
91
+
92
+ if is_async_gen or is_sync_gen:
93
+ # Handle streaming responses
94
+ async def _stream_generator(request: Request):
95
+ if is_async_gen:
96
+ async for chunk in func(request):
97
+ yield chunk
98
+ else:
99
+ for chunk in func(request):
100
+ yield chunk
101
+
102
+ @self.post(path)
103
+ async def _wrapped(request: Request):
104
+ return StreamingResponse(
105
+ _stream_generator(request),
106
+ media_type="text/plain",
107
+ )
108
+
109
+ else:
110
+ # Handle regular responses
111
+ @self.post(path)
112
+ async def _wrapped(request: Request):
113
+ if inspect.iscoroutinefunction(func):
114
+ return await func(request)
115
+ else:
116
+ return func(request)
117
+
118
+ return func
119
+
120
+ return decorator
121
+
122
+ def run(
123
+ self,
124
+ host="0.0.0.0",
125
+ port=8090,
126
+ embed_task_processor=False,
127
+ **kwargs,
128
+ ):
129
+ """
130
+ Run FastAPI with uvicorn.
131
+ """
132
+ if embed_task_processor:
133
+ if self.celery_app is None:
134
+ logger.warning(
135
+ "[AgentApp] Celery is not configured. "
136
+ "Cannot run embedded worker.",
137
+ )
138
+ else:
139
+ logger.warning(
140
+ "[AgentApp] embed_task_processor=True: Running "
141
+ "task_processor in embedded thread mode. This is "
142
+ "intended for development/debug purposes only. In "
143
+ "production, run Celery worker in a separate process!",
144
+ )
145
+
146
+ queues = self._registered_queues or {"celery"}
147
+ queue_list = ",".join(sorted(queues))
148
+
149
+ def start_celery_worker():
150
+ logger.info(
151
+ f"[AgentApp] Embedded worker listening "
152
+ f"queues: {queue_list}",
153
+ )
154
+ self.celery_app.worker_main(
155
+ [
156
+ "worker",
157
+ "--loglevel=INFO",
158
+ "-Q",
159
+ queue_list,
160
+ ],
161
+ )
162
+
163
+ threading.Thread(
164
+ target=start_celery_worker,
165
+ daemon=True,
166
+ ).start()
167
+ logger.info(
168
+ "[AgentApp] Embedded task processor started in background "
169
+ "thread (DEV mode).",
170
+ )
171
+
172
+ # TODO: Add CLI to main entrypoint to control run/deploy
173
+
174
+ config = uvicorn.Config(
175
+ app=self,
176
+ host=host,
177
+ port=port,
178
+ **kwargs,
179
+ )
180
+ self.server = uvicorn.Server(config)
181
+ self.server.run()
@@ -0,0 +1,92 @@
1
+ # -*- coding: utf-8 -*-
2
+ import inspect
3
+ import asyncio
4
+ import logging
5
+ from typing import Callable, Optional, List
6
+ from celery import Celery
7
+
8
+ logger = logging.getLogger(__name__)
9
+
10
+
11
+ class CeleryMixin:
12
+ """
13
+ Celery task processing mixin that provides core Celery functionality.
14
+ Can be reused by BaseApp and FastAPIAppFactory.
15
+ """
16
+
17
+ def __init__(
18
+ self,
19
+ broker_url: Optional[str] = None,
20
+ backend_url: Optional[str] = None,
21
+ ):
22
+ if broker_url and backend_url:
23
+ self.celery_app = Celery(
24
+ "agentscope_runtime",
25
+ broker=broker_url,
26
+ backend=backend_url,
27
+ )
28
+ else:
29
+ self.celery_app = None
30
+
31
+ self._registered_queues: set[str] = set()
32
+
33
+ def get_registered_queues(self) -> set[str]:
34
+ return self._registered_queues
35
+
36
+ def register_celery_task(self, func: Callable, queue: str = "celery"):
37
+ """Register a Celery task for the given function."""
38
+ if self.celery_app is None:
39
+ raise RuntimeError("Celery is not configured.")
40
+
41
+ self._registered_queues.add(queue)
42
+
43
+ @self.celery_app.task(queue=queue)
44
+ def wrapper(*args, **kwargs):
45
+ if inspect.iscoroutinefunction(func):
46
+ return asyncio.run(func(*args, **kwargs))
47
+ else:
48
+ return func(*args, **kwargs)
49
+
50
+ return wrapper
51
+
52
+ def run_task_processor(
53
+ self,
54
+ loglevel: str = "INFO",
55
+ concurrency: Optional[int] = None,
56
+ queues: Optional[List[str]] = None,
57
+ ):
58
+ """Run Celery worker in this process."""
59
+ if self.celery_app is None:
60
+ raise RuntimeError("Celery is not configured.")
61
+
62
+ cmd = ["worker", f"--loglevel={loglevel}"]
63
+ if concurrency:
64
+ cmd.append(f"--concurrency={concurrency}")
65
+ if queues:
66
+ cmd.append(f"-Q {','.join(queues)}")
67
+
68
+ self.celery_app.worker_main(cmd)
69
+
70
+ def get_task_status(self, task_id: str):
71
+ """Get task status from Celery result backend."""
72
+ if self.celery_app is None:
73
+ return {"error": "Celery not configured"}
74
+
75
+ result = self.celery_app.AsyncResult(task_id)
76
+ if result.state == "PENDING":
77
+ return {"status": "pending", "result": None}
78
+ elif result.state == "SUCCESS":
79
+ return {"status": "finished", "result": result.result}
80
+ elif result.state == "FAILURE":
81
+ return {"status": "error", "result": str(result.info)}
82
+ else:
83
+ return {"status": result.state, "result": None}
84
+
85
+ def submit_task(self, func: Callable, *args, **kwargs):
86
+ """Submit task directly to Celery queue."""
87
+ if not hasattr(func, "celery_task"):
88
+ raise RuntimeError(
89
+ f"Function {func.__name__} is not registered as a task",
90
+ )
91
+
92
+ return func.celery_task.delay(*args, **kwargs)
@@ -9,6 +9,7 @@ from typing import Dict
9
9
  class DeployManager(ABC):
10
10
  def __init__(self):
11
11
  self.deploy_id = str(uuid.uuid4())
12
+ self._app = None
12
13
 
13
14
  @abstractmethod
14
15
  async def deploy(self, *args, **kwargs) -> Dict[str, str]:
@@ -4,6 +4,10 @@ import asyncio
4
4
  from pathlib import Path
5
5
  from typing import Optional
6
6
 
7
+ from rich.console import Console
8
+ from rich.panel import Panel
9
+ from rich.table import Table
10
+
7
11
  from .modelstudio_deployer import ModelstudioDeployManager
8
12
  from .utils.wheel_packager import build_wheel
9
13
 
@@ -60,6 +64,19 @@ def _parse_args() -> argparse.Namespace:
60
64
  default=None,
61
65
  help="Custom directory for temporary build artifacts (optional)",
62
66
  )
67
+ parser.add_argument(
68
+ "--update",
69
+ dest="agent_id",
70
+ default=None,
71
+ help="Update an existing agent. "
72
+ "Specify agent_id to update a particular agent (optional)",
73
+ )
74
+ parser.add_argument(
75
+ "--desc",
76
+ dest="agent_desc",
77
+ default=None,
78
+ help="Add description to current agent(optional)",
79
+ )
63
80
  return parser.parse_args()
64
81
 
65
82
 
@@ -72,6 +89,8 @@ async def _run(
72
89
  build_root: Optional[str],
73
90
  mode: str,
74
91
  whl_path: Optional[str],
92
+ agent_id: Optional[str],
93
+ agent_desc: Optional[str],
75
94
  ):
76
95
  deployer = ModelstudioDeployManager(build_root=build_root)
77
96
  # If a wheel path is provided, skip local build entirely
@@ -83,6 +102,8 @@ async def _run(
83
102
  skip_upload=skip_upload,
84
103
  telemetry_enabled=telemetry_enabled,
85
104
  external_whl_path=whl_path,
105
+ agent_id=agent_id,
106
+ agent_desc=agent_desc,
86
107
  )
87
108
 
88
109
  if mode == "native":
@@ -96,20 +117,25 @@ async def _run(
96
117
  skip_upload=skip_upload,
97
118
  telemetry_enabled=telemetry_enabled,
98
119
  external_whl_path=str(built_whl),
120
+ agent_id=agent_id,
121
+ agent_desc=agent_desc,
99
122
  )
100
123
 
101
124
  # wrapper mode (default): require dir and cmd
102
- if not dir_path or not cmd:
103
- raise SystemExit(
104
- "In wrapper mode, --dir and --cmd are required. Alternatively "
105
- "use --mode native or --whl-path.",
106
- )
125
+ if not agent_id:
126
+ if not dir_path or not cmd:
127
+ raise SystemExit(
128
+ "In wrapper mode, --dir and --cmd are required. Alternatively "
129
+ "use --mode native or --whl-path.",
130
+ )
107
131
  return await deployer.deploy(
108
132
  project_dir=dir_path,
109
133
  cmd=cmd,
110
134
  deploy_name=deploy_name,
111
135
  skip_upload=skip_upload,
112
136
  telemetry_enabled=telemetry_enabled,
137
+ agent_id=agent_id,
138
+ agent_desc=agent_desc,
113
139
  )
114
140
 
115
141
 
@@ -126,17 +152,51 @@ def main() -> None:
126
152
  build_root=args.build_root,
127
153
  mode=args.mode,
128
154
  whl_path=args.whl_path,
155
+ agent_id=args.agent_id,
156
+ agent_desc=args.agent_desc,
129
157
  ),
130
158
  )
131
- print("Built wheel at:", result.get("wheel_path", ""))
159
+
160
+ console = Console()
161
+
162
+ # Create a table for basic information
163
+ info_table = Table(show_header=False, box=None, padding=(0, 2))
164
+ info_table.add_column("Key", style="bold cyan")
165
+ info_table.add_column("Value", style="white")
166
+
167
+ if result.get("wheel_path"):
168
+ info_table.add_row("Built wheel at", result.get("wheel_path", ""))
169
+
132
170
  if result.get("artifact_url"):
133
- print("Artifact URL:", result.get("artifact_url"))
134
- print("Deploy ID:", result.get("deploy_id"))
135
- print("Resource Name:", result.get("resource_name"))
171
+ info_table.add_row("Artifact URL", result.get("artifact_url"))
172
+
173
+ if result.get("resource_name"):
174
+ info_table.add_row("Resource Name", result.get("resource_name"))
175
+
136
176
  if result.get("workspace_id"):
137
- print("Workspace:", result.get("workspace_id"))
138
- if result.get("url"):
139
- print("Console Url:", result.get("url"))
177
+ info_table.add_row("Workspace", result.get("workspace_id"))
178
+
179
+ console.print(info_table)
180
+
181
+ # Display deploy result in a panel
182
+ console_url = result.get("url")
183
+ deploy_id = result.get("deploy_id")
184
+ if console_url and deploy_id:
185
+ console.print() # Add spacing
186
+ panel_content = (
187
+ f"[bold cyan]Console URL:[/bold cyan] {console_url}\n"
188
+ f"[bold cyan]Deploy ID:[/bold cyan] {deploy_id}"
189
+ )
190
+ console.print(
191
+ Panel(
192
+ panel_content,
193
+ title="[bold green]Deploy Result[/bold green]",
194
+ title_align="center",
195
+ expand=False,
196
+ border_style="green",
197
+ ),
198
+ )
199
+ console.print() # Add spacing
140
200
 
141
201
 
142
202
  if __name__ == "__main__": # pragma: no cover