31 lines
1.2 KiB
Python
31 lines
1.2 KiB
Python
import logging
|
|
|
|
from app.modules.agent.llm.service import AgentLlmService
|
|
|
|
|
|
class _FakeClient:
|
|
def complete(self, *, system_prompt: str, user_prompt: str) -> str:
|
|
assert system_prompt == "System prompt"
|
|
assert user_prompt == "User input"
|
|
return "LLM output"
|
|
|
|
|
|
class _FakePrompts:
|
|
def load(self, prompt_name: str) -> str:
|
|
assert prompt_name == "general_answer"
|
|
return "System prompt"
|
|
|
|
|
|
def test_llm_service_logs_input_and_output_for_graph_context(caplog) -> None:
|
|
service = AgentLlmService(_FakeClient(), _FakePrompts())
|
|
|
|
with caplog.at_level(logging.WARNING, logger="app.modules.agent.llm.service"):
|
|
result = service.generate("general_answer", "User input", log_context="graph.default.answer")
|
|
|
|
assert result == "LLM output"
|
|
messages = [record.getMessage() for record in caplog.records]
|
|
assert any("graph llm input: context=graph.default.answer" in message for message in messages)
|
|
assert any("graph llm output: context=graph.default.answer" in message for message in messages)
|
|
assert any("User input" in message for message in messages)
|
|
assert any("LLM output" in message for message in messages)
|