Files
ai-ticket-platform/backend/apps/agents/tests/helpers.py
2026-07-14 13:55:18 +08:00

92 lines
2.9 KiB
Python

from apps.agents.models import AgentDefinition, AgentStepBinding, AgentVersion
from apps.modelhub.crypto import encrypt_secret
from apps.modelhub.models import ModelDeployment, ModelProvider
from apps.skills.models import SkillDefinition, SkillVersion
def create_deployment():
provider = ModelProvider.objects.create(
name="模型渠道",
base_url="https://models.example.com/v1",
api_key_encrypted=encrypt_secret("secret"),
)
return ModelDeployment.objects.create(
provider=provider,
display_name="支持模型",
model_name="support-model",
)
def create_skill(code: str, output_schema: dict, input_schema: dict | None = None):
definition = SkillDefinition.objects.create(code=code, name=code)
version = SkillVersion.objects.create(
definition=definition,
version=1,
status="published",
prompt_template="处理:{{ input.value }}",
input_schema=input_schema
or {
"type": "object",
"properties": {"value": {}},
"required": ["value"],
},
output_schema=output_schema,
)
definition.active_version = version
definition.save(update_fields=["active_version"])
return version
def create_agent(deployment=None, *, published=False, two_steps=True):
deployment = deployment or create_deployment()
first_skill = create_skill(
"first-skill",
{
"type": "object",
"properties": {"summary": {"type": "string"}},
"required": ["summary"],
},
)
second_skill = None
if two_steps:
second_skill = create_skill(
"second-skill",
{
"type": "object",
"properties": {"reply": {"type": "string"}},
"required": ["reply"],
},
)
definition = AgentDefinition.objects.create(code="support-agent", name="支持 Agent")
version = AgentVersion.objects.create(
definition=definition,
version=1,
status="published" if published else "draft",
model_deployment=deployment,
input_schema={
"type": "object",
"properties": {"problem": {"type": "string"}},
"required": ["problem"],
},
example_input={"problem": "无法登录"},
)
AgentStepBinding.objects.create(
agent_version=version,
code="structure",
order=1,
skill_version=first_skill,
input_mapping={"value": "$input.problem"},
)
if second_skill:
AgentStepBinding.objects.create(
agent_version=version,
code="reply",
order=2,
skill_version=second_skill,
input_mapping={"value": "$steps.structure.output.summary"},
)
if published:
definition.active_version = version
definition.save(update_fields=["active_version"])
return definition, version