File size: 2,941 Bytes
1d4d315 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import json
import pytest
from ollie_assistants.assistant.agent_loop import AgentLoop
from ollie_assistants.assistant.tools import default_tool_registry
from ollie_assistants.llm.types import ChatConfig, ChatMessage, ChatResult, Role
from ollie_assistants.memory.store import InMemoryStore
from ollie_assistants.memory.tools import MemoryTools
from ollie_assistants.safety.guardrails import GuardrailService
class CapturingProvider:
provider_id = "test"
def __init__(self) -> None:
self.messages = []
self.calls = []
self.main_messages = []
async def chat(self, messages, config: ChatConfig) -> ChatResult:
self.calls.append(list(messages))
self.messages = list(messages)
system_text = messages[0].content if messages and messages[0].role == Role.SYSTEM else ""
if "should_search" in system_text:
return ChatResult(
text=json.dumps(
{
"should_search": True,
"query": "answer style preference",
"reason": "user asks remembered preference",
}
)
)
if "should_write" in system_text:
return ChatResult(
text=json.dumps(
{
"should_write": True,
"content": "I prefer concise answers",
"kind": "preference",
"reason": "stable preference",
}
)
)
self.main_messages = list(messages)
return ChatResult(text="ok")
@pytest.mark.asyncio
async def test_agent_loop_executes_memory_tool_before_llm() -> None:
guardrails = GuardrailService()
provider = CapturingProvider()
memory_tools = MemoryTools(InMemoryStore())
memory_tools.remember_user_fact("s1", "I prefer concise answers", "setup", "preference")
loop = AgentLoop(
provider,
default_tool_registry(memory_tools),
guardrails,
)
safety = guardrails.check_input("what is my answer style preference?")
result = await loop.run(
run_id="run-1",
session_id="s1",
model="test",
max_new_tokens=10,
messages=[ChatMessage(Role.USER, "what is my answer style preference?")],
user_message="what is my answer style preference?",
safety_decision=safety,
)
tool_names = [tool.name for tool in result.tool_results]
assert "search_memory" in tool_names
assert "remember_user_fact" in tool_names
assert all(message.role != Role.TOOL for message in provider.main_messages)
assert any("Relevant memory" in message.content for message in provider.main_messages)
assert any(event.name == "memory_search_decision" for event in result.events)
assert any(event.name == "memory_write_decision" for event in result.events)
|