wehub-resource-sync
214 lines
7.7 KiB
Python
214 lines
7.7 KiB
Python
import json
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, ToolMessage
|
|
|
|
import mlflow
|
|
from mlflow.langchain.chat_agent_langgraph import parse_message
|
|
from mlflow.types.agent import ChatAgentMessage
|
|
|
|
LC_TOOL_CALL_MSG = AIMessage(**{
|
|
"content": "",
|
|
"additional_kwargs": {
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_a9b9afd5-d23a-4973-8417-ac283b1413d5",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "system__ai__python_exec",
|
|
"arguments": '{ "code": "print(5+5)" }',
|
|
},
|
|
}
|
|
]
|
|
},
|
|
"response_metadata": {"prompt_tokens": 2658, "completion_tokens": 24, "total_tokens": 2682},
|
|
"type": "ai",
|
|
"name": None,
|
|
"id": "run-3a2ad83b-a5cf-4d51-97c8-9f68205df787-0",
|
|
"example": False,
|
|
"tool_calls": [
|
|
{
|
|
"name": "system__ai__python_exec",
|
|
"args": {"code": "print(5+5)"},
|
|
"id": "call_a9b9afd5-d23a-4973-8417-ac283b1413d5",
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
"invalid_tool_calls": [],
|
|
"usage_metadata": None,
|
|
})
|
|
CHAT_AGENT_TOOL_CALL_MSG = ChatAgentMessage(**{
|
|
"role": "assistant",
|
|
"content": "",
|
|
"name": "llm",
|
|
"id": "run-3a2ad83b-a5cf-4d51-97c8-9f68205df787-0",
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_a9b9afd5-d23a-4973-8417-ac283b1413d5",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "system__ai__python_exec",
|
|
"arguments": '{"code": "print(5+5)"}',
|
|
},
|
|
}
|
|
],
|
|
}).model_dump(exclude_none=True)
|
|
LC_TOOL_MSG = ToolMessage(**{
|
|
"content": '{"content": "Successfully generated array of 5 random ints in [1, 100].", "attachments": {"key1": "attach1", "key2": "attach2"}, "custom_outputs": {"random_nums": [1, 82, 9, 12, 22]}}', # noqa: E501
|
|
"additional_kwargs": {},
|
|
"response_metadata": {},
|
|
"type": "tool",
|
|
"name": "generate_random_ints",
|
|
"id": None,
|
|
"tool_call_id": "call_ee823299-62d7-4407-95e8-168412904471",
|
|
"artifact": None,
|
|
"status": "success",
|
|
})
|
|
CHAT_AGENT_TOOL_MSG = ChatAgentMessage(
|
|
role="tool",
|
|
content='{"content": "Successfully generated array of 5 random ints in [1, 100].", "attachments": {"key1": "attach1", "key2": "attach2"}, "custom_outputs": {"random_nums": [1, 82, 9, 12, 22]}}', # noqa: E501
|
|
name="generate_random_ints",
|
|
tool_calls=None,
|
|
tool_call_id="call_ee823299-62d7-4407-95e8-168412904471",
|
|
attachments={"key1": "attach1", "key2": "attach2"},
|
|
finish_reason=None,
|
|
).model_dump(exclude_none=True) # id will be a generated UUID
|
|
TOOL_MSG_ATTACHMENTS = {"key1": "attach1", "key2": "attach2"}
|
|
LC_ASSISTANT_MSG = AIMessage(**{
|
|
"content": "The generated random numbers are 1, 82, 9, 12, and 22.",
|
|
"additional_kwargs": {},
|
|
"response_metadata": {"prompt_tokens": 2763, "completion_tokens": 22, "total_tokens": 2785},
|
|
"type": "ai",
|
|
"name": None,
|
|
"id": "run-4972ab0f-8b90-4650-8a84-a689fbd912f1-0",
|
|
"example": False,
|
|
"tool_calls": [],
|
|
"invalid_tool_calls": [],
|
|
"usage_metadata": None,
|
|
})
|
|
CHAT_AGENT_ASSISTANT_MSG = ChatAgentMessage(
|
|
role="assistant",
|
|
content="The generated random numbers are 1, 82, 9, 12, and 22.",
|
|
name="llm",
|
|
id="run-4972ab0f-8b90-4650-8a84-a689fbd912f1-0",
|
|
tool_calls=None,
|
|
tool_call_id=None,
|
|
attachments=None,
|
|
finish_reason=None,
|
|
).model_dump(exclude_none=True)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("lc_msg", "chat_agent_msg", "name", "attachments"),
|
|
[
|
|
(LC_TOOL_CALL_MSG, CHAT_AGENT_TOOL_CALL_MSG, "llm", None),
|
|
(LC_TOOL_MSG, CHAT_AGENT_TOOL_MSG, None, TOOL_MSG_ATTACHMENTS),
|
|
(LC_ASSISTANT_MSG, CHAT_AGENT_ASSISTANT_MSG, "llm", None),
|
|
],
|
|
)
|
|
def test_parse_message(lc_msg, chat_agent_msg, name, attachments):
|
|
# id is autogenerated
|
|
if lc_msg.id is None:
|
|
lc_msg.id = chat_agent_msg.get("id")
|
|
assert parse_message(lc_msg, name, attachments) == chat_agent_msg
|
|
|
|
|
|
def test_langgraph_chat_agent_save_as_code():
|
|
# (role, content)
|
|
expected_messages = [
|
|
("assistant", ""), # tool message does not have content
|
|
(
|
|
"tool",
|
|
json.dumps({
|
|
"format": "SCALAR",
|
|
"value": '{"content":"hi","attachments":{"a":"b"},"custom_outputs":{"c":"d"}}',
|
|
"truncated": False,
|
|
}),
|
|
),
|
|
("assistant", ""),
|
|
(
|
|
"tool",
|
|
json.dumps({
|
|
"content": f"Successfully generated array of 2 random ints: {[1, 2]}.",
|
|
"attachments": {"key1": "attach1", "key2": "attach2"},
|
|
"custom_outputs": {"random_nums": [1, 2]},
|
|
}),
|
|
),
|
|
("assistant", "Successfully generated"),
|
|
]
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="agent",
|
|
python_model="tests/langgraph/sample_code/langgraph_chat_agent.py",
|
|
)
|
|
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
response = loaded_model.predict({"messages": [{"role": "user", "content": "hi"}]})
|
|
messages = response["messages"]
|
|
assert len(messages) == len(expected_messages)
|
|
for msg, (role, expected_content) in zip(messages, expected_messages):
|
|
assert msg["role"] == role
|
|
assert msg["content"] == expected_content
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
response = loaded_model.predict_stream({"messages": [{"role": "user", "content": "hi"}]})
|
|
for event, (role, expected_content) in zip(response, expected_messages):
|
|
assert event["delta"]["content"] == expected_content
|
|
assert event["delta"]["role"] == role
|
|
|
|
|
|
def test_langgraph_chat_agent_custom_inputs():
|
|
# (role, content)
|
|
expected_messages = [
|
|
("assistant", ""), # tool message does not have content
|
|
(
|
|
"tool",
|
|
json.dumps({
|
|
"format": "SCALAR",
|
|
"value": '{"content":"hi","attachments":{"a":"b"},"custom_outputs":{"c":"d"}}',
|
|
"truncated": False,
|
|
}),
|
|
),
|
|
("assistant", ""),
|
|
(
|
|
"tool",
|
|
json.dumps({
|
|
"content": f"Successfully generated array of 2 random ints: {[1, 2]}.",
|
|
"attachments": {"key1": "attach1", "key2": "attach2"},
|
|
"custom_outputs": {"random_nums": [1, 2]},
|
|
}),
|
|
),
|
|
("assistant", "Successfully generated"),
|
|
("assistant", "adding custom outputs"),
|
|
]
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="agent",
|
|
python_model="tests/langgraph/sample_code/langgraph_chat_agent_custom_inputs.py",
|
|
)
|
|
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
response = loaded_model.predict({
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"custom_inputs": {"asdf": "jkl;"},
|
|
})
|
|
assert response["custom_outputs"]["asdf"] == "jkl;"
|
|
messages = response["messages"]
|
|
assert len(messages) == len(expected_messages)
|
|
for msg, (role, expected_content) in zip(messages, expected_messages):
|
|
assert msg["role"] == role
|
|
assert msg["content"] == expected_content
|
|
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
response = loaded_model.predict_stream({
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"custom_inputs": {"asdf": "jkl;"},
|
|
})
|
|
counter = 0
|
|
for chunk, (role, expected_content) in zip(response, expected_messages):
|
|
assert chunk["delta"]["content"] == expected_content
|
|
assert chunk["delta"]["role"] == role
|
|
if "custom_outputs" in chunk:
|
|
assert chunk["custom_outputs"]["asdf"] == "jkl;"
|
|
counter += 1
|
|
assert counter == 1
|