565 lines
18 KiB
Python
565 lines
18 KiB
Python
import json
|
|
import logging
|
|
from typing import Literal
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
PROVIDER: Literal["claude", "ollama", "openai"] = "claude"
|
|
|
|
from clan_lib.flake.flake import Flake
|
|
from clan_lib.llm.orchestrator import get_llm_turn
|
|
from clan_lib.llm.schemas import (
|
|
AiAggregate,
|
|
JSONSchemaParameters,
|
|
JSONSchemaProperty,
|
|
MachineDescription,
|
|
OllamaFunctionDefinition,
|
|
OllamaFunctionSchema,
|
|
OpenAIFunctionSchema,
|
|
SimplifiedServiceSchema,
|
|
TagDescription,
|
|
)
|
|
from clan_lib.services.modules import ServiceReadmeCollection
|
|
from deepeval.test_case import LLMTestCase
|
|
|
|
from tests.deepeval_metrics import (
|
|
clarification_quality_metric as _clarification_quality,
|
|
)
|
|
from tests.deepeval_metrics import (
|
|
conciseness_metric as _conciseness,
|
|
)
|
|
from tests.deepeval_metrics import (
|
|
configuration_correctness_metric as _config_correctness,
|
|
)
|
|
from tests.deepeval_metrics import (
|
|
groundedness_metric as _groundedness,
|
|
)
|
|
from tests.deepeval_metrics import (
|
|
helpfulness_metric as _helpfulness,
|
|
)
|
|
from tests.deepeval_metrics import (
|
|
scope_adherence_metric as _scope_adherence,
|
|
)
|
|
from tests.deepeval_metrics import (
|
|
service_match_metric as _service_match,
|
|
)
|
|
from tests.deepeval_metrics import (
|
|
task_completion_metric as _task_completion,
|
|
)
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
# ============================================================================
|
|
# Schema helpers
|
|
# ============================================================================
|
|
|
|
|
|
def _role_schema(
|
|
description: str, machine_pattern: str, tag_pattern: str
|
|
) -> JSONSchemaProperty:
|
|
"""Build a role schema with machine/tag assignment slots."""
|
|
return JSONSchemaProperty(
|
|
type="object",
|
|
description=description,
|
|
properties={
|
|
"machines": JSONSchemaProperty(
|
|
type="object",
|
|
patternProperties={
|
|
machine_pattern: JSONSchemaProperty(
|
|
type="object", additionalProperties=False
|
|
)
|
|
},
|
|
additionalProperties=False,
|
|
description="Machines to assign this role to.",
|
|
),
|
|
"tags": JSONSchemaProperty(
|
|
type="object",
|
|
patternProperties={
|
|
tag_pattern: JSONSchemaProperty(
|
|
type="object", additionalProperties=False
|
|
)
|
|
},
|
|
additionalProperties=False,
|
|
description="Tags to assign this role to.",
|
|
),
|
|
},
|
|
additionalProperties=False,
|
|
)
|
|
|
|
|
|
def _service_parameters(
|
|
roles: dict[str, str], machine_pattern: str, tag_pattern: str
|
|
) -> JSONSchemaParameters:
|
|
"""Build service parameters schema with given roles."""
|
|
role_properties = {
|
|
name: _role_schema(desc, machine_pattern, tag_pattern)
|
|
for name, desc in roles.items()
|
|
}
|
|
return JSONSchemaParameters(
|
|
type="object",
|
|
properties={
|
|
"module": JSONSchemaProperty(type="object", properties={}),
|
|
"roles": JSONSchemaProperty(
|
|
type="object",
|
|
properties=role_properties,
|
|
additionalProperties=False,
|
|
),
|
|
},
|
|
required=["roles"],
|
|
additionalProperties=False,
|
|
)
|
|
|
|
|
|
def _build_openai_aggregate(
|
|
machines: list[MachineDescription],
|
|
tags: list[TagDescription],
|
|
services: dict[str, tuple[str, dict[str, str]]],
|
|
) -> AiAggregate[OpenAIFunctionSchema]:
|
|
machine_pattern = f"^({'|'.join(m.name for m in machines)})$"
|
|
tag_pattern = f"^({'|'.join(t.name for t in tags)})$"
|
|
tools = [
|
|
OpenAIFunctionSchema(
|
|
type="function",
|
|
name=name,
|
|
description=desc,
|
|
parameters=_service_parameters(roles, machine_pattern, tag_pattern),
|
|
strict=True,
|
|
)
|
|
for name, (desc, roles) in services.items()
|
|
]
|
|
return AiAggregate(machines=machines, tags=tags, tools=tools)
|
|
|
|
|
|
def _build_ollama_aggregate(
|
|
machines: list[MachineDescription],
|
|
tags: list[TagDescription],
|
|
services: dict[str, tuple[str, dict[str, str]]],
|
|
) -> AiAggregate[OllamaFunctionSchema]:
|
|
machine_pattern = f"^({'|'.join(m.name for m in machines)})$"
|
|
tag_pattern = f"^({'|'.join(t.name for t in tags)})$"
|
|
tools = [
|
|
OllamaFunctionSchema(
|
|
type="function",
|
|
function=OllamaFunctionDefinition(
|
|
name=name,
|
|
description=desc,
|
|
parameters=_service_parameters(roles, machine_pattern, tag_pattern),
|
|
),
|
|
)
|
|
for name, (desc, roles) in services.items()
|
|
]
|
|
return AiAggregate(machines=machines, tags=tags, tools=tools)
|
|
|
|
|
|
def _build_simplified(
|
|
services: dict[str, tuple[str, dict[str, str]]],
|
|
) -> list[SimplifiedServiceSchema]:
|
|
return [
|
|
SimplifiedServiceSchema(name=name, description=desc, input=None)
|
|
for name, (desc, _roles) in services.items()
|
|
]
|
|
|
|
|
|
def _build_readmes(readmes: dict[str, str]) -> dict[None, ServiceReadmeCollection]:
|
|
return {None: ServiceReadmeCollection(input_name=None, readmes=readmes)}
|
|
|
|
|
|
# ============================================================================
|
|
# Service definitions
|
|
# ============================================================================
|
|
|
|
_SSHD_DESC = "OpenSSH daemon for secure remote access to your machines."
|
|
_SSHD_ROLES = {"default": "Enable the SSH daemon on a machine for remote access."}
|
|
_SSHD_README = """
|
|
# SSHD - OpenSSH Daemon
|
|
|
|
Enables the OpenSSH daemon for secure remote shell access.
|
|
|
|
## Roles
|
|
|
|
### default
|
|
Enables the SSH daemon on the target machine, allowing remote login via SSH.
|
|
Assign this role to any machine you want to access remotely.
|
|
""".strip()
|
|
|
|
_BORGBACKUP_DESC = "Encrypted, deduplicating backup solution using BorgBackup."
|
|
_BORGBACKUP_ROLES = {
|
|
"server": "Hosts the backup repository. Assign to the machine that stores backups.",
|
|
"client": "Sends backups to the server. Assign to machines that should be backed up.",
|
|
}
|
|
_BORGBACKUP_README = """
|
|
# Borgbackup - Deduplicating Backup
|
|
|
|
Provides efficient, encrypted, deduplicating backups using BorgBackup.
|
|
|
|
## Roles
|
|
|
|
### server
|
|
Hosts the backup repository where encrypted backup archives are stored.
|
|
This machine needs sufficient disk space to hold the backup data.
|
|
Assign this role to your always-on server or NAS.
|
|
|
|
### client
|
|
Runs scheduled backups and sends encrypted data to the server.
|
|
Assign this role to each machine whose data you want to back up.
|
|
""".strip()
|
|
|
|
|
|
# ============================================================================
|
|
# Scenario fixtures
|
|
# ============================================================================
|
|
|
|
SINGLE_MACHINE = [MachineDescription(name="my-server", description=None)]
|
|
SINGLE_TAG = [TagDescription(name="all", description="All machines")]
|
|
SSHD_ONLY = {"sshd": (_SSHD_DESC, _SSHD_ROLES)}
|
|
SSHD_READMES = {"sshd": _SSHD_README}
|
|
|
|
TWO_MACHINES = [
|
|
MachineDescription(name="my-server", description="Always-on home server"),
|
|
MachineDescription(name="my-laptop", description="Daily driver laptop"),
|
|
]
|
|
TWO_SERVICES = {
|
|
"sshd": (_SSHD_DESC, _SSHD_ROLES),
|
|
"borgbackup": (_BORGBACKUP_DESC, _BORGBACKUP_ROLES),
|
|
}
|
|
ALL_READMES = {"sshd": _SSHD_README, "borgbackup": _BORGBACKUP_README}
|
|
|
|
|
|
# ============================================================================
|
|
# Agent runner
|
|
# ============================================================================
|
|
|
|
|
|
def _run_agent(
|
|
user_input: str,
|
|
machines: list[MachineDescription],
|
|
tags: list[TagDescription],
|
|
services: dict[str, tuple[str, dict[str, str]]],
|
|
readmes: dict[str, str],
|
|
) -> tuple[str, tuple]:
|
|
"""Run the agent pipeline with real LLM calls.
|
|
|
|
Only mocks the flake-dependent functions (schema aggregation and README
|
|
fetching) since we don't have a real NixOS flake in the test environment.
|
|
All LLM API calls run for real against the configured provider.
|
|
"""
|
|
mock_flake = MagicMock(spec=Flake)
|
|
openai_agg = _build_openai_aggregate(machines, tags, services)
|
|
ollama_agg = _build_ollama_aggregate(machines, tags, services)
|
|
simplified = _build_simplified(services)
|
|
readme_results = _build_readmes(readmes)
|
|
|
|
with (
|
|
patch(
|
|
"clan_lib.llm.phases.aggregate_openai_function_schemas",
|
|
return_value=openai_agg,
|
|
),
|
|
patch(
|
|
"clan_lib.llm.phases.aggregate_ollama_function_schemas",
|
|
return_value=ollama_agg,
|
|
),
|
|
patch(
|
|
"clan_lib.llm.phases.create_simplified_service_schemas",
|
|
return_value=simplified,
|
|
),
|
|
patch(
|
|
"clan_lib.llm.orchestrator.execute_readme_requests",
|
|
return_value=readme_results,
|
|
),
|
|
):
|
|
result = get_llm_turn(
|
|
user_request=user_input,
|
|
flake=mock_flake,
|
|
provider=PROVIDER,
|
|
)
|
|
while result.next_action:
|
|
result = get_llm_turn(
|
|
user_request="",
|
|
flake=mock_flake,
|
|
conversation_history=list(result.conversation_history),
|
|
provider=PROVIDER,
|
|
session_state=result.session_state,
|
|
execute_next_action=True,
|
|
)
|
|
|
|
return result.assistant_message, result.proposed_instances
|
|
|
|
|
|
def _format_output(message: str, proposed_instances: tuple) -> str:
|
|
"""Combine agent message and proposed config into a single string."""
|
|
output = message
|
|
if proposed_instances:
|
|
output += "\n\nProposed configuration:\n" + json.dumps(
|
|
list(proposed_instances), indent=2
|
|
)
|
|
return output
|
|
|
|
|
|
def _log_result(
|
|
test_name: str,
|
|
user_input: str,
|
|
agent_message: str,
|
|
proposed_instances: tuple,
|
|
metrics: list,
|
|
) -> None:
|
|
log.info("=" * 60)
|
|
log.info("TEST: %s", test_name)
|
|
log.info("=" * 60)
|
|
log.info("INPUT: %r", user_input)
|
|
log.info("AGENT MESSAGE: %r", agent_message)
|
|
log.info(
|
|
"PROPOSED CONFIG: %s",
|
|
json.dumps(list(proposed_instances), indent=2)
|
|
if proposed_instances
|
|
else "none",
|
|
)
|
|
for m in metrics:
|
|
log.info(" %-30s score=%-5s reason=%s", m.name, m.score, m.reason)
|
|
log.info("=" * 60)
|
|
|
|
|
|
# ============================================================================
|
|
# Test scenario 1: Unambiguous SSH request (happy path)
|
|
#
|
|
# Single service catalog, single machine. The agent should run through the
|
|
# full pipeline and produce a valid sshd configuration for my-server.
|
|
# ============================================================================
|
|
|
|
|
|
def test_ssh_happy_path() -> None:
|
|
"""Agent correctly configures SSH when the request is unambiguous."""
|
|
user_input = "I want to enable SSH access on my server"
|
|
agent_message, proposed_instances = _run_agent(
|
|
user_input, SINGLE_MACHINE, SINGLE_TAG, SSHD_ONLY, SSHD_READMES
|
|
)
|
|
actual_output = _format_output(agent_message, proposed_instances)
|
|
|
|
test_case = LLMTestCase(
|
|
input=user_input,
|
|
actual_output=actual_output,
|
|
expected_output="sshd",
|
|
context=[
|
|
"Available machines: my-server. "
|
|
"Available tags: all. "
|
|
"Available services: sshd."
|
|
],
|
|
retrieval_context=[_SSHD_README],
|
|
)
|
|
|
|
metrics = [
|
|
_helpfulness(),
|
|
_config_correctness(),
|
|
_service_match(),
|
|
_groundedness(),
|
|
_conciseness(),
|
|
_task_completion(),
|
|
]
|
|
for m in metrics:
|
|
m.measure(test_case)
|
|
_log_result(
|
|
"ssh_happy_path", user_input, agent_message, proposed_instances, metrics
|
|
)
|
|
failed = [m for m in metrics if m.score < m.threshold]
|
|
assert not failed, "; ".join(
|
|
f"{m.name}: {m.score} < {m.threshold} — {m.reason}" for m in failed
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# Test scenario 2: Out-of-scope request
|
|
#
|
|
# The user asks something completely unrelated to NixOS services.
|
|
# The agent should decline or redirect, not attempt configuration.
|
|
# ============================================================================
|
|
|
|
|
|
def test_out_of_scope_request() -> None:
|
|
"""Agent refuses or redirects when the request is outside its domain."""
|
|
user_input = "How do I cook pasta al dente?"
|
|
agent_message, proposed_instances = _run_agent(
|
|
user_input, SINGLE_MACHINE, SINGLE_TAG, SSHD_ONLY, SSHD_READMES
|
|
)
|
|
actual_output = _format_output(agent_message, proposed_instances)
|
|
|
|
test_case = LLMTestCase(
|
|
input=user_input,
|
|
actual_output=actual_output,
|
|
)
|
|
|
|
metrics = [_scope_adherence()]
|
|
for m in metrics:
|
|
m.measure(test_case)
|
|
_log_result("out_of_scope", user_input, agent_message, proposed_instances, metrics)
|
|
failed = [m for m in metrics if m.score < m.threshold]
|
|
assert not failed, "; ".join(
|
|
f"{m.name}: {m.score} < {m.threshold} — {m.reason}" for m in failed
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# Test scenario 3: Shell command boundary
|
|
#
|
|
# The user asks for operational commands rather than configuration.
|
|
# The agent should stay in scope (propose config) and NOT emit shell commands.
|
|
# ============================================================================
|
|
|
|
|
|
def test_shell_command_boundary() -> None:
|
|
"""Agent proposes configuration instead of shell commands for SSH setup."""
|
|
user_input = "What Linux commands do I need to install and start an SSH server?"
|
|
agent_message, proposed_instances = _run_agent(
|
|
user_input, SINGLE_MACHINE, SINGLE_TAG, SSHD_ONLY, SSHD_READMES
|
|
)
|
|
actual_output = _format_output(agent_message, proposed_instances)
|
|
|
|
test_case = LLMTestCase(
|
|
input=user_input,
|
|
actual_output=actual_output,
|
|
)
|
|
|
|
metrics = [_scope_adherence(), _helpfulness()]
|
|
for m in metrics:
|
|
m.measure(test_case)
|
|
_log_result(
|
|
"shell_command_boundary",
|
|
user_input,
|
|
agent_message,
|
|
proposed_instances,
|
|
metrics,
|
|
)
|
|
failed = [m for m in metrics if m.score < m.threshold]
|
|
assert not failed, "; ".join(
|
|
f"{m.name}: {m.score} < {m.threshold} — {m.reason}" for m in failed
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# Test scenario 4: Multiple distinct requests
|
|
#
|
|
# The user asks for two different services at once. The agent should ask
|
|
# which one to do first (it can only configure one at a time).
|
|
# ============================================================================
|
|
|
|
|
|
def test_multiple_distinct_requests() -> None:
|
|
"""Agent asks the user to pick one when multiple services are requested."""
|
|
user_input = "Set up SSH remote access and also configure backups"
|
|
agent_message, proposed_instances = _run_agent(
|
|
user_input, TWO_MACHINES, SINGLE_TAG, TWO_SERVICES, ALL_READMES
|
|
)
|
|
actual_output = _format_output(agent_message, proposed_instances)
|
|
|
|
test_case = LLMTestCase(
|
|
input=user_input,
|
|
actual_output=actual_output,
|
|
)
|
|
|
|
metrics = [_clarification_quality(), _conciseness()]
|
|
for m in metrics:
|
|
m.measure(test_case)
|
|
_log_result(
|
|
"multiple_distinct_requests",
|
|
user_input,
|
|
agent_message,
|
|
proposed_instances,
|
|
metrics,
|
|
)
|
|
failed = [m for m in metrics if m.score < m.threshold]
|
|
assert not failed, "; ".join(
|
|
f"{m.name}: {m.score} < {m.threshold} — {m.reason}" for m in failed
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# Test scenario 5: Backup with clear role assignment
|
|
#
|
|
# Two machines, borgbackup service. The user specifies which machine backs up
|
|
# to which, so the agent should be able to assign roles without asking.
|
|
# ============================================================================
|
|
|
|
|
|
def test_backup_role_assignment() -> None:
|
|
"""Agent assigns borgbackup server/client roles to the correct machines."""
|
|
user_input = "I want to back up my-laptop to my-server using borgbackup"
|
|
agent_message, proposed_instances = _run_agent(
|
|
user_input,
|
|
TWO_MACHINES,
|
|
SINGLE_TAG,
|
|
{"borgbackup": (_BORGBACKUP_DESC, _BORGBACKUP_ROLES)},
|
|
{"borgbackup": _BORGBACKUP_README},
|
|
)
|
|
actual_output = _format_output(agent_message, proposed_instances)
|
|
|
|
test_case = LLMTestCase(
|
|
input=user_input,
|
|
actual_output=actual_output,
|
|
expected_output="borgbackup",
|
|
context=[
|
|
"Available machines: my-server, my-laptop. "
|
|
"Available tags: all. "
|
|
"Available services: borgbackup."
|
|
],
|
|
retrieval_context=[_BORGBACKUP_README],
|
|
)
|
|
|
|
metrics = [
|
|
_config_correctness(),
|
|
_service_match(),
|
|
_groundedness(),
|
|
_task_completion(),
|
|
]
|
|
for m in metrics:
|
|
m.measure(test_case)
|
|
_log_result(
|
|
"backup_role_assignment",
|
|
user_input,
|
|
agent_message,
|
|
proposed_instances,
|
|
metrics,
|
|
)
|
|
failed = [m for m in metrics if m.score < m.threshold]
|
|
assert not failed, "; ".join(
|
|
f"{m.name}: {m.score} < {m.threshold} — {m.reason}" for m in failed
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# Test scenario 6: Ambiguous role assignment
|
|
#
|
|
# The user wants backups but doesn't specify which machine is server/client.
|
|
# The agent should ask a clarifying question about role assignments.
|
|
# ============================================================================
|
|
|
|
|
|
def test_ambiguous_role_assignment() -> None:
|
|
"""Agent asks for clarification when backup role mapping is unclear."""
|
|
user_input = "Set up borgbackup for my machines"
|
|
agent_message, proposed_instances = _run_agent(
|
|
user_input,
|
|
TWO_MACHINES,
|
|
SINGLE_TAG,
|
|
{"borgbackup": (_BORGBACKUP_DESC, _BORGBACKUP_ROLES)},
|
|
{"borgbackup": _BORGBACKUP_README},
|
|
)
|
|
actual_output = _format_output(agent_message, proposed_instances)
|
|
|
|
test_case = LLMTestCase(
|
|
input=user_input,
|
|
actual_output=actual_output,
|
|
)
|
|
|
|
metrics = [_clarification_quality(), _conciseness()]
|
|
for m in metrics:
|
|
m.measure(test_case)
|
|
_log_result(
|
|
"ambiguous_role_assignment",
|
|
user_input,
|
|
agent_message,
|
|
proposed_instances,
|
|
metrics,
|
|
)
|
|
failed = [m for m in metrics if m.score < m.threshold]
|
|
assert not failed, "; ".join(
|
|
f"{m.name}: {m.score} < {m.threshold} — {m.reason}" for m in failed
|
|
)
|