232 lines
7.6 KiB
Python
232 lines
7.6 KiB
Python
from dataclasses import asdict
|
|
|
|
from pam_deploy_graph.agent import PamDeployAgent
|
|
from pam_deploy_graph.checkpoint_store import redact_mapping
|
|
from pam_deploy_graph.llm.factory import build_llm_client
|
|
from pam_deploy_graph.llm.openai_compatible import OpenAICompatibleLlmClient
|
|
from pam_deploy_graph.llm.rule_based import RuleBasedLlmClient
|
|
from pam_deploy_graph.llm.validators import validate_deploy_plan
|
|
from pam_deploy_graph.models import LlmDeployPlan
|
|
from pam_deploy_graph.models import ActionResult
|
|
|
|
|
|
def test_understand_request_prefers_hybrid_for_mcp():
|
|
result = RuleBasedLlmClient().understand_request("请用 MCP 部署 HET")
|
|
assert result.intent == "deploy"
|
|
assert result.mode_preference == "MCP"
|
|
assert result.strategy_preference == "hybrid_node_mcp"
|
|
|
|
|
|
def test_extract_params_from_key_value_text():
|
|
result = RuleBasedLlmClient().extract_params(
|
|
"HOME_BASE_URL=https://x CLIENT_ID=id CLIENT_SECRET=s AIRPORT_CODE=HET "
|
|
"APP_NAME=PAM MODULE_NAME=Node VERSION_NUMBER=2.0.5 ZIP_FILE_PATH=C:/pkg.zip"
|
|
)
|
|
assert result.extracted_params["AIRPORT_CODE"] == "HET"
|
|
assert result.missing_required_params == []
|
|
assert "CLIENT_SECRET" in result.sensitive_fields_present
|
|
|
|
|
|
def test_analyze_request_returns_structured_objects():
|
|
agent = PamDeployAgent()
|
|
result = agent.analyze_request(
|
|
"不要动环境,预演部署",
|
|
{
|
|
"HOME_BASE_URL": "https://x",
|
|
"CLIENT_ID": "id",
|
|
"CLIENT_SECRET": "s",
|
|
"AIRPORT_CODE": "HET",
|
|
"APP_NAME": "PAM",
|
|
"MODULE_NAME": "Node",
|
|
"VERSION_NUMBER": "2.0.5",
|
|
"ZIP_FILE_PATH": "C:/pkg.zip",
|
|
},
|
|
)
|
|
payload = {key: asdict(value) for key, value in result.items()}
|
|
assert payload["intent"]["intent"] == "preview"
|
|
assert payload["plan"]["execution_strategy"] == "hybrid_node_mcp"
|
|
|
|
|
|
def test_analyze_payload_can_be_redacted():
|
|
agent = PamDeployAgent()
|
|
result = agent.analyze_request(
|
|
"帮我部署",
|
|
{
|
|
"HOME_BASE_URL": "https://x",
|
|
"CLIENT_ID": "id",
|
|
"CLIENT_SECRET": "super-secret",
|
|
"AIRPORT_CODE": "HET",
|
|
"APP_NAME": "PAM",
|
|
"MODULE_NAME": "Node",
|
|
"VERSION_NUMBER": "2.0.5",
|
|
"ZIP_FILE_PATH": "C:/pkg.zip",
|
|
},
|
|
)
|
|
payload = redact_mapping({key: asdict(value) for key, value in result.items()})
|
|
assert payload["params"]["extracted_params"]["CLIENT_SECRET"] == "***"
|
|
|
|
|
|
def test_plan_guardrails_reject_executable_text():
|
|
plan = LlmDeployPlan(summary="run bash ./deploy.sh", planned_actions=["get-token"])
|
|
try:
|
|
validate_deploy_plan(plan)
|
|
except ValueError as exc:
|
|
assert "禁止" in str(exc)
|
|
else:
|
|
raise AssertionError("expected guardrail failure")
|
|
|
|
|
|
def test_openai_compatible_client_uses_base_url_api_key_and_prompt():
|
|
calls = []
|
|
|
|
def transport(url, headers, payload, timeout_sec):
|
|
calls.append((url, headers, payload, timeout_sec))
|
|
return {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": (
|
|
'{"intent":"deploy","mode_preference":"MCP",'
|
|
'"strategy_preference":"hybrid_node_mcp","confidence":0.9,'
|
|
'"reasons":["ok"]}'
|
|
)
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
client = OpenAICompatibleLlmClient(
|
|
base_url="https://llm.example/v1",
|
|
api_key="secret-key",
|
|
model="model-a",
|
|
transport=transport,
|
|
)
|
|
|
|
result = client.understand_request("请用 MCP 部署")
|
|
|
|
assert result.intent == "deploy"
|
|
assert calls[0][0] == "https://llm.example/v1/chat/completions"
|
|
assert calls[0][1]["Authorization"] == "Bearer secret-key"
|
|
assert calls[0][2]["model"] == "model-a"
|
|
assert "只输出一个 JSON 对象" in calls[0][2]["messages"][0]["content"]
|
|
|
|
|
|
def test_openai_compatible_client_allows_empty_api_key():
|
|
calls = []
|
|
|
|
def transport(url, headers, payload, timeout_sec):
|
|
calls.append((url, headers, payload, timeout_sec))
|
|
return {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": (
|
|
'{"intent":"deploy","mode_preference":"未指定",'
|
|
'"strategy_preference":"fake","confidence":0.8}'
|
|
)
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
client = OpenAICompatibleLlmClient(
|
|
base_url="https://llm.example/v1",
|
|
api_key="",
|
|
model="model-a",
|
|
transport=transport,
|
|
)
|
|
|
|
result = client.understand_request("部署")
|
|
|
|
assert result.intent == "deploy"
|
|
assert "Authorization" not in calls[0][1]
|
|
assert calls[0][1]["Content-Type"] == "application/json"
|
|
|
|
|
|
def test_llm_factory_allows_empty_api_key_with_base_url_and_model():
|
|
client = build_llm_client(base_url="https://llm.example/v1", model="model-a")
|
|
|
|
assert isinstance(client, OpenAICompatibleLlmClient)
|
|
assert client.api_key == ""
|
|
|
|
|
|
def test_openai_compatible_client_does_not_send_base_secret():
|
|
calls = []
|
|
|
|
def transport(url, headers, payload, timeout_sec):
|
|
calls.append(payload)
|
|
return {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": (
|
|
'{"extracted_params":{"AIRPORT_CODE":"HET"},'
|
|
'"extracted_control":{},'
|
|
'"missing_required_params":[],'
|
|
'"ambiguous_fields":[]}'
|
|
)
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
client = OpenAICompatibleLlmClient(
|
|
base_url="https://llm.example/v1",
|
|
api_key="secret-key",
|
|
model="model-a",
|
|
transport=transport,
|
|
)
|
|
|
|
result = client.extract_params("机场 HET", {"CLIENT_SECRET": "real-secret", "CLIENT_ID": "id"})
|
|
|
|
serialized_prompt = str(calls[0])
|
|
assert "real-secret" not in serialized_prompt
|
|
assert result.extracted_params["CLIENT_SECRET"] == "real-secret"
|
|
|
|
|
|
def test_openai_compatible_client_analyzes_action_result_with_redaction():
|
|
calls = []
|
|
|
|
def transport(url, headers, payload, timeout_sec):
|
|
calls.append(payload)
|
|
return {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": (
|
|
'{"action":"verify-ip","has_anomaly":true,"severity":"high",'
|
|
'"possible_reason":"health check failed",'
|
|
'"suggested_action":"download logs","requires_confirmation":true,'
|
|
'"notes":["verify failed"]}'
|
|
)
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
client = OpenAICompatibleLlmClient(
|
|
base_url="https://llm.example/v1",
|
|
api_key="secret-key",
|
|
model="model-a",
|
|
transport=transport,
|
|
)
|
|
|
|
analysis = client.analyze_action_result(
|
|
action="verify-ip",
|
|
result=ActionResult(
|
|
action="verify-ip",
|
|
backend="fake",
|
|
ok=False,
|
|
values={"CLIENT_SECRET": "real-secret", "SUCCESS": "false"},
|
|
stderr="x" * 1200,
|
|
error_summary="failed",
|
|
),
|
|
state_summary={"params": {"CLIENT_SECRET": "real-secret"}},
|
|
)
|
|
|
|
serialized_prompt = str(calls[0])
|
|
assert analysis.has_anomaly is True
|
|
assert analysis.severity == "high"
|
|
assert "real-secret" not in serialized_prompt
|
|
assert "[已截断]" in serialized_prompt
|