wehub-resource-sync
b7f52be4c9
CI / Run CI (push) Has been cancelled
CI / check-backend (push) Has been cancelled
CI / check-frontend (push) Has been cancelled
CI / tests (push) Has been cancelled
CI / e2e-tests (push) Has been cancelled
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
from unittest.mock import patch
|
|
|
|
from llama_index.core.callbacks.schema import CBEventType, EventPayload
|
|
from llama_index.core.tools.types import ToolMetadata
|
|
|
|
from chainlit.llama_index.callbacks import LlamaIndexCallbackHandler
|
|
from chainlit.step import Step
|
|
|
|
|
|
async def test_on_event_start_for_function_calls(mock_chainlit_context):
|
|
TEST_EVENT_ID = "test_event_id"
|
|
async with mock_chainlit_context:
|
|
handler = LlamaIndexCallbackHandler()
|
|
|
|
with patch.object(Step, "send") as mock_send:
|
|
result = handler.on_event_start(
|
|
CBEventType.FUNCTION_CALL,
|
|
{
|
|
EventPayload.TOOL: ToolMetadata(
|
|
name="test_tool", description="test_description"
|
|
),
|
|
EventPayload.FUNCTION_CALL: {"arg1": "value1"},
|
|
},
|
|
TEST_EVENT_ID,
|
|
)
|
|
|
|
assert result == TEST_EVENT_ID
|
|
assert TEST_EVENT_ID in handler.steps
|
|
step = handler.steps[TEST_EVENT_ID]
|
|
assert isinstance(step, Step)
|
|
assert step.name == "test_tool"
|
|
assert step.type == "tool"
|
|
assert step.id == TEST_EVENT_ID
|
|
assert step.input == '{\n "arg1": "value1"\n}'
|
|
mock_send.assert_called_once()
|
|
|
|
|
|
async def test_on_event_start_for_function_calls_missing_payload(mock_chainlit_context):
|
|
TEST_EVENT_ID = "test_event_id"
|
|
async with mock_chainlit_context:
|
|
handler = LlamaIndexCallbackHandler()
|
|
|
|
with patch.object(Step, "send") as mock_send:
|
|
result = handler.on_event_start(
|
|
CBEventType.FUNCTION_CALL,
|
|
None,
|
|
TEST_EVENT_ID,
|
|
)
|
|
|
|
assert result == TEST_EVENT_ID
|
|
assert TEST_EVENT_ID in handler.steps
|
|
step = handler.steps[TEST_EVENT_ID]
|
|
assert isinstance(step, Step)
|
|
assert step.name == "function_call"
|
|
assert step.type == "tool"
|
|
assert step.id == TEST_EVENT_ID
|
|
assert step.input == "{}"
|
|
mock_send.assert_called_once()
|
|
|
|
|
|
async def test_on_event_end_for_function_calls(mock_chainlit_context):
|
|
TEST_EVENT_ID = "test_event_id"
|
|
async with mock_chainlit_context:
|
|
handler = LlamaIndexCallbackHandler()
|
|
# Pretend that we have started a step before.
|
|
step = Step(name="test_tool", type="tool", id=TEST_EVENT_ID)
|
|
handler.steps[TEST_EVENT_ID] = step
|
|
|
|
with patch.object(step, "update") as mock_send:
|
|
handler.on_event_end(
|
|
CBEventType.FUNCTION_CALL,
|
|
payload={EventPayload.FUNCTION_OUTPUT: "test_output"},
|
|
event_id=TEST_EVENT_ID,
|
|
)
|
|
|
|
assert step.output == "test_output"
|
|
assert TEST_EVENT_ID not in handler.steps
|
|
mock_send.assert_called_once()
|
|
|
|
|
|
async def test_on_event_end_for_function_calls_missing_payload(mock_chainlit_context):
|
|
TEST_EVENT_ID = "test_event_id"
|
|
async with mock_chainlit_context:
|
|
handler = LlamaIndexCallbackHandler()
|
|
# Pretend that we have started a step before.
|
|
step = Step(name="test_tool", type="tool", id=TEST_EVENT_ID)
|
|
handler.steps[TEST_EVENT_ID] = step
|
|
|
|
with patch.object(step, "update") as mock_send:
|
|
handler.on_event_end(
|
|
CBEventType.FUNCTION_CALL,
|
|
payload=None,
|
|
event_id=TEST_EVENT_ID,
|
|
)
|
|
# TODO: Is this the desired behavior? Shouldn't we still remove the step as long as we've been told it has ended, even if the payload is missing?
|
|
assert TEST_EVENT_ID in handler.steps
|
|
mock_send.assert_not_called()
|