forked from HAL9000/cleveragents-core
140 lines
4.1 KiB
Python
140 lines
4.1 KiB
Python
from typing import Any
|
|
from typing import Dict
|
|
from typing import List
|
|
from typing import Optional
|
|
|
|
from behave import given
|
|
from behave import then
|
|
from behave import when
|
|
from hamcrest import assert_that
|
|
from hamcrest import equal_to
|
|
from hamcrest import has_key
|
|
from hamcrest import is_
|
|
from hamcrest import not_
|
|
|
|
from cleveragents.agents.base import Agent
|
|
from cleveragents.agents.base import AgentWithMemory
|
|
from cleveragents.core.exceptions import AgentCreationError
|
|
from cleveragents.templates.renderer import TemplateEngine
|
|
from cleveragents.templates.renderer import TemplateRenderer
|
|
|
|
|
|
# A simple test agent for instantiation
|
|
class CoverageTestAgent(Agent):
|
|
async def process(
|
|
self, message: str, context: Optional[Dict[str, Any]] = None
|
|
) -> str:
|
|
return "processed"
|
|
|
|
def get_capabilities(self) -> List[str]:
|
|
return ["coverage-test"]
|
|
|
|
|
|
class CoverageTestAgentWithMemory(AgentWithMemory):
|
|
async def process(
|
|
self, message: str, context: Optional[Dict[str, Any]] = None
|
|
) -> str:
|
|
self.memory["last_message"] = message
|
|
return "processed with memory"
|
|
|
|
def get_capabilities(self) -> List[str]:
|
|
return ["coverage-test-memory"]
|
|
|
|
|
|
def _setup_context(context):
|
|
if not hasattr(context, "template_renderer"):
|
|
context.template_renderer = TemplateRenderer(TemplateEngine.SIMPLE)
|
|
|
|
|
|
@given("an Agent is initialized with model and provider in config")
|
|
def step_impl(context):
|
|
_setup_context(context)
|
|
config = {"model": "test-model", "provider": "test-provider"}
|
|
context.agent = CoverageTestAgent("test-agent", config, context.template_renderer)
|
|
|
|
|
|
@given("an Agent is initialized with only a model in config")
|
|
def step_impl(context):
|
|
_setup_context(context)
|
|
config = {"model": "test-model"}
|
|
context.agent = CoverageTestAgent("test-agent", config, context.template_renderer)
|
|
|
|
|
|
@given("an Agent is initialized with only a provider in config")
|
|
def step_impl(context):
|
|
_setup_context(context)
|
|
config = {"provider": "test-provider"}
|
|
context.agent = CoverageTestAgent("test-agent", config, context.template_renderer)
|
|
|
|
|
|
@when("I get the agent's metadata")
|
|
def step_impl(context):
|
|
context.metadata = context.agent.get_metadata()
|
|
|
|
|
|
@then('the metadata should contain the model "{model}"')
|
|
def step_impl(context, model):
|
|
assert_that(context.metadata, has_key("model"))
|
|
assert_that(context.metadata["model"], equal_to(model))
|
|
|
|
|
|
@then('the metadata should contain the provider "{provider}"')
|
|
def step_impl(context, provider):
|
|
assert_that(context.metadata, has_key("provider"))
|
|
assert_that(context.metadata["provider"], equal_to(provider))
|
|
|
|
|
|
@then("the metadata should not contain a provider")
|
|
def step_impl(context):
|
|
assert_that(context.metadata, not_(has_key("provider")))
|
|
|
|
|
|
@then("the metadata should not contain a model")
|
|
def step_impl(context):
|
|
assert_that(context.metadata, not_(has_key("model")))
|
|
|
|
|
|
@given("an AgentWithMemory instance")
|
|
def step_impl(context):
|
|
_setup_context(context)
|
|
context.agent = CoverageTestAgentWithMemory(
|
|
"memory-agent", {}, context.template_renderer
|
|
)
|
|
|
|
|
|
@when("I try to load memory with a non-dictionary value")
|
|
def step_impl(context):
|
|
context.error = None
|
|
try:
|
|
context.agent.load_memory("not a dict")
|
|
except AgentCreationError as e:
|
|
context.error = e
|
|
|
|
|
|
@then("an AgentCreationError should be raised")
|
|
def step_impl(context):
|
|
assert_that(context.error, is_(AgentCreationError))
|
|
|
|
|
|
@given("I save the initial memory state")
|
|
def step_impl(context):
|
|
context.agent.memory = {"initial": "state"}
|
|
context.saved_memory = context.agent.save_memory()
|
|
# Also check that it's a deep copy
|
|
assert_that(context.saved_memory is not context.agent.memory)
|
|
|
|
|
|
@when("I modify the agent's memory")
|
|
def step_impl(context):
|
|
context.agent.memory["modified"] = True
|
|
|
|
|
|
@when("I load the saved memory state")
|
|
def step_impl(context):
|
|
context.agent.load_memory(context.saved_memory)
|
|
|
|
|
|
@then("the agent's memory should be restored to the initial state")
|
|
def step_impl(context):
|
|
assert_that(context.agent.memory, equal_to({"initial": "state"}))
|