clan_lib: Add llm unit tests
This commit is contained in:
Qubasa
289
pkgs/clan-cli/clan_lib/llm/llm_test.py
Normal file
289
pkgs/clan-cli/clan_lib/llm/llm_test.py
Normal file
@@ -0,0 +1,289 @@
|
|||||||
|
from collections.abc import Callable
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from clan_cli.tests.fixtures_flakes import nested_dict
|
||||||
|
from clan_lib.flake.flake import Flake
|
||||||
|
from clan_lib.llm.llm import (
|
||||||
|
OpenAIFunctionSchema,
|
||||||
|
aggregate_openai_function_schemas,
|
||||||
|
llm_final_decision_to_inventory_instances,
|
||||||
|
)
|
||||||
|
from clan_lib.llm.schemas import FunctionCallType, clan_module_to_openai_spec
|
||||||
|
from clan_lib.services.modules import list_service_modules
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.with_core
|
||||||
|
def test_clan_module_to_llm_func(
|
||||||
|
clan_flake: Callable[..., Flake],
|
||||||
|
) -> None:
|
||||||
|
# ATTENTION! This method lacks Typechecking
|
||||||
|
config = nested_dict()
|
||||||
|
# explicit module selection
|
||||||
|
# We use this random string in test to avoid code dependencies on the input name
|
||||||
|
config["inventory"]["instances"]["foo"]["module"]["input"] = (
|
||||||
|
"Y2xhbi1jaW9yZS1uZXZlci1kZXBlbmQtb24tbWU"
|
||||||
|
)
|
||||||
|
config["inventory"]["instances"]["foo"]["module"]["name"] = "sshd"
|
||||||
|
# input = null
|
||||||
|
config["inventory"]["instances"]["bar"]["module"]["input"] = None
|
||||||
|
config["inventory"]["instances"]["bar"]["module"]["name"] = "sshd"
|
||||||
|
|
||||||
|
config["inventory"]["machines"] = {
|
||||||
|
"machine1": {
|
||||||
|
"tags": ["production", "backup"],
|
||||||
|
},
|
||||||
|
"machine2": {
|
||||||
|
"tags": ["client"],
|
||||||
|
},
|
||||||
|
"machine3": {
|
||||||
|
"tags": ["client"],
|
||||||
|
},
|
||||||
|
}
|
||||||
|
config["inventory"]["tags"] = {
|
||||||
|
"production": [],
|
||||||
|
"backup": [],
|
||||||
|
"client": [],
|
||||||
|
}
|
||||||
|
|
||||||
|
# Omit input
|
||||||
|
config["inventory"]["instances"]["baz"]["module"]["name"] = "sshd"
|
||||||
|
# external input
|
||||||
|
flake = clan_flake(config)
|
||||||
|
|
||||||
|
service_modules = list_service_modules(flake)
|
||||||
|
|
||||||
|
# Module(usage_ref={'name': 'borgbackup', 'input': None}, info=ModuleInfo(manifest=ModuleManifest(name='borgbackup', description='Efficient, deduplicating backup program with optional compression and secure encryption.', categories=['System'], features={'API': True}), roles={'client': Role(name='client', description='A borgbackup client that backs up to all borgbackup server roles.'), 'server': Role(name='server', description='A borgbackup server that stores the backups of clients.')}), native=True, instance_refs=[]),
|
||||||
|
borgbackup_service = next(
|
||||||
|
m for m in service_modules.modules if m.usage_ref.get("name") == "borgbackup"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert borgbackup_service is not None
|
||||||
|
|
||||||
|
available_machines = ["machine1", "machine2", "server1"]
|
||||||
|
available_tags = ["production", "backup", "client"]
|
||||||
|
|
||||||
|
generated_tool_func = clan_module_to_openai_spec(
|
||||||
|
borgbackup_service, available_tags, available_machines
|
||||||
|
)
|
||||||
|
|
||||||
|
expected_tool_func: OpenAIFunctionSchema = {
|
||||||
|
"type": "function",
|
||||||
|
"name": "borgbackup",
|
||||||
|
"description": "Efficient, deduplicating backup program with optional compression and secure encryption.",
|
||||||
|
"parameters": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"module": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {},
|
||||||
|
},
|
||||||
|
"roles": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"client": {
|
||||||
|
"type": "object",
|
||||||
|
"description": "A borgbackup client that backs up to all borgbackup server roles.",
|
||||||
|
"properties": {
|
||||||
|
"machines": {
|
||||||
|
"type": "object",
|
||||||
|
"patternProperties": {
|
||||||
|
"^(machine1|machine2|server1)$": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": False,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
"description": 'Machines to assign this role to. Format: each machine name is a key with an empty object {} as value. Example: {"wintux": {}, "gchq-local": {}}',
|
||||||
|
},
|
||||||
|
"tags": {
|
||||||
|
"type": "object",
|
||||||
|
"patternProperties": {
|
||||||
|
"^(production|backup|client)$": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": False,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
"description": 'Tags to assign this role to. Format: each tag name is a key with an empty object {} as value. Example: {"all": {}, "nixos": {}}',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
},
|
||||||
|
"server": {
|
||||||
|
"type": "object",
|
||||||
|
"description": "A borgbackup server that stores the backups of clients.",
|
||||||
|
"properties": {
|
||||||
|
"machines": {
|
||||||
|
"type": "object",
|
||||||
|
"patternProperties": {
|
||||||
|
"^(machine1|machine2|server1)$": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": False,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
"description": 'Machines to assign this role to. Format: each machine name is a key with an empty object {} as value. Example: {"wintux": {}, "gchq-local": {}}',
|
||||||
|
},
|
||||||
|
"tags": {
|
||||||
|
"type": "object",
|
||||||
|
"patternProperties": {
|
||||||
|
"^(production|backup|client)$": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": False,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
"description": 'Tags to assign this role to. Format: each tag name is a key with an empty object {} as value. Example: {"all": {}, "nixos": {}}',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["roles"],
|
||||||
|
"additionalProperties": False,
|
||||||
|
},
|
||||||
|
"strict": True,
|
||||||
|
}
|
||||||
|
|
||||||
|
assert generated_tool_func == expected_tool_func
|
||||||
|
|
||||||
|
certificate_service = next(
|
||||||
|
m for m in service_modules.modules if m.usage_ref.get("name") == "certificates"
|
||||||
|
)
|
||||||
|
assert certificate_service is not None
|
||||||
|
|
||||||
|
generated_tool_func2 = clan_module_to_openai_spec(
|
||||||
|
certificate_service, available_tags, available_machines
|
||||||
|
)
|
||||||
|
|
||||||
|
expected_tool_func2: OpenAIFunctionSchema = {
|
||||||
|
"type": "function",
|
||||||
|
"name": "certificates",
|
||||||
|
"description": "Sets up a PKI certificate chain using step-ca",
|
||||||
|
"parameters": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"module": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {},
|
||||||
|
},
|
||||||
|
"roles": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"ca": {
|
||||||
|
"type": "object",
|
||||||
|
"description": "A certificate authority that issues and signs certificates for other machines.",
|
||||||
|
"properties": {
|
||||||
|
"machines": {
|
||||||
|
"type": "object",
|
||||||
|
"patternProperties": {
|
||||||
|
"^(machine1|machine2|server1)$": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": False,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
"description": 'Machines to assign this role to. Format: each machine name is a key with an empty object {} as value. Example: {"wintux": {}, "gchq-local": {}}',
|
||||||
|
},
|
||||||
|
"tags": {
|
||||||
|
"type": "object",
|
||||||
|
"patternProperties": {
|
||||||
|
"^(production|backup|client)$": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": False,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
"description": 'Tags to assign this role to. Format: each tag name is a key with an empty object {} as value. Example: {"all": {}, "nixos": {}}',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
},
|
||||||
|
"default": {
|
||||||
|
"type": "object",
|
||||||
|
"description": "A machine that trusts the CA and can get certificates issued by it.",
|
||||||
|
"properties": {
|
||||||
|
"machines": {
|
||||||
|
"type": "object",
|
||||||
|
"patternProperties": {
|
||||||
|
"^(machine1|machine2|server1)$": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": False,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
"description": 'Machines to assign this role to. Format: each machine name is a key with an empty object {} as value. Example: {"wintux": {}, "gchq-local": {}}',
|
||||||
|
},
|
||||||
|
"tags": {
|
||||||
|
"type": "object",
|
||||||
|
"patternProperties": {
|
||||||
|
"^(production|backup|client)$": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": False,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
"description": 'Tags to assign this role to. Format: each tag name is a key with an empty object {} as value. Example: {"all": {}, "nixos": {}}',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"additionalProperties": False,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["roles"],
|
||||||
|
"additionalProperties": False,
|
||||||
|
},
|
||||||
|
"strict": True,
|
||||||
|
}
|
||||||
|
|
||||||
|
assert generated_tool_func2 == expected_tool_func2
|
||||||
|
|
||||||
|
aggregate = aggregate_openai_function_schemas(flake)
|
||||||
|
|
||||||
|
assert len(aggregate.tools) >= 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_llm_final_decision_to_inventory_conversion() -> None:
|
||||||
|
"""Test conversion of LLM final decision to inventory format."""
|
||||||
|
final_decision: list[FunctionCallType] = [
|
||||||
|
{
|
||||||
|
"id": "toolu_01XHjHUMzZVTcDCqaYQJEWu5",
|
||||||
|
"call_id": "toolu_01XHjHUMzZVTcDCqaYQJEWu5",
|
||||||
|
"type": "function_call",
|
||||||
|
"name": "matrix-synapse",
|
||||||
|
"arguments": '{"roles": {"default": {"machines": {"gchq-local": {}}}}}',
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": "toolu_01TsjKZ87J3fi6RNzNzu33ff",
|
||||||
|
"call_id": "toolu_01TsjKZ87J3fi6RNzNzu33ff",
|
||||||
|
"type": "function_call",
|
||||||
|
"name": "monitoring",
|
||||||
|
"arguments": '{"module": { "input": "qubasas-clan" }, "roles": {"telegraf": {"tags": {"all": {}}}}}',
|
||||||
|
},
|
||||||
|
]
|
||||||
|
assert isinstance(final_decision, list)
|
||||||
|
|
||||||
|
expected = [
|
||||||
|
{
|
||||||
|
"module": {
|
||||||
|
"input": None,
|
||||||
|
"name": "matrix-synapse",
|
||||||
|
},
|
||||||
|
"roles": {"default": {"machines": {"gchq-local": {}}}},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"module": {
|
||||||
|
"input": "qubasas-clan",
|
||||||
|
"name": "monitoring",
|
||||||
|
},
|
||||||
|
"roles": {"telegraf": {"tags": {"all": {}}}},
|
||||||
|
},
|
||||||
|
]
|
||||||
|
|
||||||
|
result = llm_final_decision_to_inventory_instances(final_decision)
|
||||||
|
assert result == expected
|
||||||
1957
pkgs/clan-cli/clan_lib/llm/test_process_chat_turn.py
Normal file
1957
pkgs/clan-cli/clan_lib/llm/test_process_chat_turn.py
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user