forked from xiaohei/taiji-AI-PAD
201 lines
6.5 KiB
Python
201 lines
6.5 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
简化的 Agent Manager API 格式验证脚本
|
||
不依赖外部模块，仅验证数据结构
|
||
"""
|
||
|
||
|
||
def test_payload_format():
|
||
"""验证新的 API Payload 格式"""
|
||
print("=" * 60)
|
||
print("Agent Manager API 新格式验证")
|
||
print("=" * 60)
|
||
|
||
# 模拟 AgentConfig
|
||
class AgentConfig:
|
||
def __init__(self, user_id=None, cpu_request="100m", cpu_limit="500m",
|
||
memory_request="128Mi", memory_limit="512Mi", replicas=1):
|
||
self.user_id = user_id
|
||
self.cpu_request = cpu_request
|
||
self.cpu_limit = cpu_limit
|
||
self.memory_request = memory_request
|
||
self.memory_limit = memory_limit
|
||
self.replicas = replicas
|
||
|
||
def to_dict(self):
|
||
result = {}
|
||
if self.user_id:
|
||
result["user_id"] = self.user_id
|
||
if self.cpu_request:
|
||
result["cpu_request"] = self.cpu_request
|
||
if self.cpu_limit:
|
||
result["cpu_limit"] = self.cpu_limit
|
||
if self.memory_request:
|
||
result["memory_request"] = self.memory_request
|
||
if self.memory_limit:
|
||
result["memory_limit"] = self.memory_limit
|
||
if self.replicas is not None:
|
||
result["replicas"] = self.replicas
|
||
return result
|
||
|
||
# 测试 1: 平台 Agent（单副本）
|
||
print("\n[测试 1] 平台 Agent 部署格式")
|
||
print("-" * 60)
|
||
|
||
config1 = AgentConfig(
|
||
user_id="user-123",
|
||
cpu_request="100m",
|
||
cpu_limit="500m",
|
||
memory_request="128Mi",
|
||
memory_limit="512Mi",
|
||
replicas=1
|
||
)
|
||
|
||
env1 = {}
|
||
|
||
payload1 = {
|
||
"name": "echo-agent-user123",
|
||
"template": "echo_agent",
|
||
"config": config1.to_dict()
|
||
}
|
||
|
||
if env1:
|
||
payload1["env"] = env1
|
||
|
||
if config1.replicas:
|
||
payload1["replicas"] = config1.replicas
|
||
|
||
print(f"✓ Agent 名称: {payload1['name']}")
|
||
print(f"✓ 模板类型: {payload1['template']}")
|
||
print(f"✓ 副本数（顶层）: {payload1.get('replicas', 'N/A')}")
|
||
print(f"✓ 配置内副本数: {payload1['config'].get('replicas', 'N/A')}")
|
||
print(f"✓ 环境变量字段: {'env' if 'env' in payload1 else '无'}")
|
||
print(f"✓ 用户 ID: {payload1['config'].get('user_id')}")
|
||
|
||
# 测试 2: Azure Blob Agent（多副本 + 环境变量）
|
||
print("\n[测试 2] Azure Blob Agent 部署格式")
|
||
print("-" * 60)
|
||
|
||
config2 = AgentConfig(
|
||
user_id="user-456",
|
||
cpu_request="100m",
|
||
cpu_limit="500m",
|
||
memory_request="256Mi",
|
||
memory_limit="1Gi",
|
||
replicas=2
|
||
)
|
||
|
||
env2 = {
|
||
"LITELLM_API_BASE": "http://litellm-service:4000",
|
||
"LITELLM_MODEL": "gpt-4",
|
||
"LITELLM_API_KEY": "sk-test-key",
|
||
"AZURE_STORAGE_CONNECTION_STRING": "DefaultEndpointsProtocol=https;...",
|
||
"SERVICE_PORT": "8080"
|
||
}
|
||
|
||
payload2 = {
|
||
"name": "azure-blob-agent-user456",
|
||
"template": "azure_blob_agent",
|
||
"config": config2.to_dict()
|
||
}
|
||
|
||
if env2:
|
||
payload2["env"] = env2
|
||
|
||
if config2.replicas:
|
||
payload2["replicas"] = config2.replicas
|
||
|
||
print(f"✓ Agent 名称: {payload2['name']}")
|
||
print(f"✓ 模板类型: {payload2['template']}")
|
||
print(f"✓ 副本数（顶层）: {payload2.get('replicas', 'N/A')}")
|
||
print(f"✓ 配置内副本数: {payload2['config'].get('replicas', 'N/A')}")
|
||
print(f"✓ 环境变量字段: {'env' if 'env' in payload2 else '无'}")
|
||
print(f"✓ 环境变量数量: {len(payload2.get('env', {}))}")
|
||
print(f"✓ LiteLLM 配置:")
|
||
if 'env' in payload2:
|
||
print(f" - LITELLM_API_BASE: {payload2['env'].get('LITELLM_API_BASE')}")
|
||
print(f" - LITELLM_MODEL: {payload2['env'].get('LITELLM_MODEL')}")
|
||
print(f" - SERVICE_PORT: {payload2['env'].get('SERVICE_PORT')}")
|
||
|
||
# 测试 3: 关键差异对比
|
||
print("\n[测试 3] 新旧格式关键差异")
|
||
print("-" * 60)
|
||
|
||
old_format = {
|
||
"name": "my-agent",
|
||
"template": "my_template",
|
||
"config": {"user_id": "user", "cpu_request": "100m"},
|
||
"env": {"KEY": "value"} # 旧字段名
|
||
}
|
||
|
||
new_format = {
|
||
"name": "my-agent",
|
||
"template": "my_template",
|
||
"replicas": 2, # 新增顶层字段
|
||
"config": {
|
||
"user_id": "user",
|
||
"cpu_request": "100m",
|
||
"replicas": 2 # 新增配置字段
|
||
},
|
||
"env": {"KEY": "value"} # 统一后字段名
|
||
}
|
||
|
||
print("旧格式:")
|
||
print(f" - 环境变量字段: 'env_variables' (不统一)")
|
||
print(f" - 副本数字段: ❌ 不支持")
|
||
print(f" - config.replicas: ❌ 不支持")
|
||
|
||
print("\n新格式 (统一后):")
|
||
print(f" - 环境变量字段: 'env' ✓")
|
||
print(f" - 副本数字段: 'replicas' ✓")
|
||
print(f" - config.replicas: {new_format['config']['replicas']} ✓")
|
||
|
||
# 测试 4: 字段映射验证
|
||
print("\n[测试 4] 字段映射验证")
|
||
print("-" * 60)
|
||
|
||
checks = [
|
||
("AgentConfig 包含 replicas 字段", True),
|
||
("AgentConfig.to_dict() 返回 replicas", 'replicas' in config2.to_dict()),
|
||
("payload 使用 env 而非 env_variables", 'env' in payload2 and 'env_variables' not in payload2),
|
||
("payload 包含顶层 replicas 字段", 'replicas' in payload2),
|
||
("config 和顶层 replicas 值一致", payload2.get('replicas') == payload2['config'].get('replicas')),
|
||
]
|
||
|
||
all_passed = True
|
||
for check_name, passed in checks:
|
||
status = "✓ 通过" if passed else "✗ 失败"
|
||
print(f"{status}: {check_name}")
|
||
if not passed:
|
||
all_passed = False
|
||
|
||
# 总结
|
||
print("\n" + "=" * 60)
|
||
print("验证结果总结")
|
||
print("=" * 60)
|
||
|
||
if all_passed:
|
||
print("✓ 所有验证项通过")
|
||
print("✓ 代码已成功适配新的 Agent Manager API 格式")
|
||
print("\n主要变更:")
|
||
print(" 1. AgentConfig 新增 replicas 字段（默认为 1）")
|
||
print(" 2. API payload 统一使用 env")
|
||
print(" 3. 支持多副本部署")
|
||
print(" 4. 兼容 LiteLLM 和 Azure 配置")
|
||
return True
|
||
else:
|
||
print("✗ 部分验证项失败，请检查代码")
|
||
return False
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import sys
|
||
try:
|
||
success = test_payload_format()
|
||
sys.exit(0 if success else 1)
|
||
except Exception as e:
|
||
print(f"\n✗ 验证过程出错: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
sys.exit(1)
|