wehub-resource-sync
289 lines
11 KiB
Python
289 lines
11 KiB
Python
import json
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.entities.span import SpanType
|
|
from mlflow.entities.span_status import SpanStatusCode
|
|
from mlflow.tracing.constant import TokenUsageKey, TraceMetadataKey
|
|
from mlflow.version import IS_TRACING_SDK_ONLY
|
|
|
|
from tests.tracing.helper import get_traces, skip_when_testing_trace_sdk
|
|
|
|
|
|
@skip_when_testing_trace_sdk
|
|
def test_langgraph_save_as_code():
|
|
input_example = {"messages": [{"role": "user", "content": "what is the weather in sf?"}]}
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.langchain.log_model(
|
|
"tests/langgraph/sample_code/langgraph_prebuilt.py",
|
|
name="langgraph",
|
|
input_example=input_example,
|
|
)
|
|
|
|
# (role, content)
|
|
expected_messages = [
|
|
("human", "what is the weather in sf?"),
|
|
("agent", ""), # tool message does not have content
|
|
("tools", "It's always sunny in sf"),
|
|
("agent", "The weather in San Francisco is always sunny!"),
|
|
]
|
|
|
|
loaded_graph = mlflow.langchain.load_model(model_info.model_uri)
|
|
response = loaded_graph.invoke(input_example)
|
|
messages = response["messages"]
|
|
assert len(messages) == 4
|
|
for msg, (role, expected_content) in zip(messages, expected_messages):
|
|
assert msg.content == expected_content
|
|
|
|
# Need to reload to reset the iterator in FakeOpenAI
|
|
loaded_graph = mlflow.langchain.load_model(model_info.model_uri)
|
|
response = loaded_graph.stream(input_example)
|
|
# .stream() response does not includes the first Human message
|
|
for chunk, (role, expected_content) in zip(response, expected_messages[1:]):
|
|
assert chunk[role]["messages"][0].content == expected_content
|
|
|
|
loaded_pyfunc = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
response = loaded_pyfunc.predict(input_example)[0]
|
|
messages = response["messages"]
|
|
assert len(messages) == 4
|
|
for msg, (role, expected_content) in zip(messages, expected_messages):
|
|
assert msg["content"] == expected_content
|
|
# response should be json serializable
|
|
assert json.dumps(response) is not None
|
|
|
|
loaded_pyfunc = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
response = loaded_pyfunc.predict_stream(input_example)
|
|
for chunk, (role, expected_content) in zip(response, expected_messages[1:]):
|
|
assert chunk[role]["messages"][0]["content"] == expected_content
|
|
|
|
|
|
@skip_when_testing_trace_sdk
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
|
|
async def test_langgraph_tracing_prebuilt(is_async, mock_litellm_cost):
|
|
from tests.langgraph.sample_code.langgraph_prebuilt import graph
|
|
|
|
mlflow.langchain.autolog()
|
|
|
|
input_example = {"messages": [{"role": "user", "content": "what is the weather in sf?"}]}
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
if is_async:
|
|
await graph.ainvoke(input_example, config)
|
|
else:
|
|
graph.invoke(input_example, config)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert traces[0].info.status == "OK"
|
|
assert traces[0].data.spans[0].name == "LangGraph"
|
|
assert traces[0].data.spans[0].inputs == input_example
|
|
|
|
# (type, content)
|
|
expected_messages = [
|
|
("human", "what is the weather in sf?"),
|
|
("ai", ""), # tool message does not have content
|
|
("tool", "It's always sunny in sf"),
|
|
("ai", "The weather in San Francisco is always sunny!"),
|
|
]
|
|
|
|
messages = traces[0].data.spans[0].outputs["messages"]
|
|
assert len(messages) == 4
|
|
for msg, (type, expected_content) in zip(messages, expected_messages):
|
|
assert msg["type"] == type
|
|
assert msg["content"] == expected_content
|
|
|
|
# Validate tool span
|
|
tool_span = next(span for span in traces[0].data.spans if span.span_type == SpanType.TOOL)
|
|
assert tool_span.name == "get_weather"
|
|
assert tool_span.inputs == {"city": "sf"}
|
|
assert tool_span.outputs["content"] == "It's always sunny in sf"
|
|
assert tool_span.outputs["status"] == "success"
|
|
assert tool_span.status.status_code == SpanStatusCode.OK
|
|
|
|
# Validate token usage
|
|
token_usage = json.loads(traces[0].info.trace_metadata[TraceMetadataKey.TOKEN_USAGE])
|
|
assert token_usage == {
|
|
TokenUsageKey.INPUT_TOKENS: 15,
|
|
TokenUsageKey.OUTPUT_TOKENS: 30,
|
|
TokenUsageKey.TOTAL_TOKENS: 45,
|
|
}
|
|
|
|
# Thread ID should be recoded in the trace metadata
|
|
assert traces[0].info.trace_metadata[TraceMetadataKey.TRACE_SESSION] == "1"
|
|
|
|
# Verify chat model spans have model name extracted
|
|
chat_spans = [s for s in traces[0].data.spans if s.span_type == SpanType.CHAT_MODEL]
|
|
for chat_span in chat_spans:
|
|
assert chat_span.model_name == "gpt-3.5-turbo"
|
|
if not IS_TRACING_SDK_ONLY:
|
|
usage = chat_span.get_attribute("mlflow.chat.tokenUsage")
|
|
assert chat_span.llm_cost == {
|
|
"input_cost": usage["input_tokens"] * 1.0,
|
|
"output_cost": usage["output_tokens"] * 2.0,
|
|
"total_cost": usage["input_tokens"] * 1.0 + usage["output_tokens"] * 2.0,
|
|
}
|
|
|
|
|
|
@skip_when_testing_trace_sdk
|
|
def test_langgraph_tracing_diy_graph():
|
|
mlflow.langchain.autolog()
|
|
|
|
input_example = {"messages": [{"role": "user", "content": "hi"}]}
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.langchain.log_model(
|
|
"tests/langgraph/sample_code/langgraph_diy.py",
|
|
name="langgraph",
|
|
)
|
|
|
|
loaded_graph = mlflow.langchain.load_model(model_info.model_uri)
|
|
loaded_graph.invoke(input_example)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert traces[0].info.status == "OK"
|
|
assert traces[0].data.spans[0].name == "LangGraph"
|
|
assert traces[0].data.spans[0].inputs == input_example
|
|
|
|
chat_spans = [span for span in traces[0].data.spans if span.name.startswith("ChatOpenAI")]
|
|
assert len(chat_spans) == 3
|
|
# Verify all chat model spans have model name extracted
|
|
for chat_span in chat_spans:
|
|
assert chat_span.model_name == "gpt-3.5-turbo"
|
|
|
|
|
|
@skip_when_testing_trace_sdk
|
|
def test_langgraph_tracing_with_custom_span():
|
|
mlflow.langchain.autolog()
|
|
|
|
input_example = {"messages": [{"role": "user", "content": "what is the weather in sf?"}]}
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.langchain.log_model(
|
|
"tests/langgraph/sample_code/langgraph_with_custom_span.py",
|
|
name="langgraph",
|
|
input_example=input_example,
|
|
)
|
|
|
|
loaded_graph = mlflow.langchain.load_model(model_info.model_uri)
|
|
|
|
# No trace should be created for the first call
|
|
assert mlflow.get_trace(mlflow.get_last_active_trace_id()) is None
|
|
|
|
loaded_graph.invoke(input_example)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert traces[0].info.status == "OK"
|
|
assert traces[0].data.spans[0].name == "LangGraph"
|
|
assert traces[0].data.spans[0].inputs == input_example
|
|
|
|
spans = traces[0].data.spans
|
|
|
|
# Validate chat model spans
|
|
chat_spans = [s for s in spans if s.span_type == SpanType.CHAT_MODEL]
|
|
assert len(chat_spans) == 3
|
|
# Verify all chat model spans have model name extracted
|
|
for chat_span in chat_spans:
|
|
assert chat_span.model_name == "gpt-3.5-turbo"
|
|
|
|
# Validate tool span
|
|
tool_span = next(s for s in spans if s.span_type == SpanType.TOOL)
|
|
assert tool_span.name == "get_weather"
|
|
assert tool_span.inputs == {"city": "sf"}
|
|
assert tool_span.outputs["content"] == "It's always sunny in sf"
|
|
assert tool_span.outputs["status"] == "success"
|
|
assert tool_span.status.status_code == SpanStatusCode.OK
|
|
|
|
# Validate inner span
|
|
inner_span = next(s for s in spans if s.name == "get_weather_inner")
|
|
assert inner_span.parent_id == tool_span.span_id
|
|
assert inner_span.inputs == "sf"
|
|
assert inner_span.outputs == "It's always sunny in sf"
|
|
|
|
inner_runnable_span = next(s for s in spans if s.parent_id == inner_span.span_id)
|
|
assert inner_runnable_span.name == "RunnableSequence"
|
|
|
|
|
|
@skip_when_testing_trace_sdk
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
|
|
async def test_langgraph_tracing_with_parent_span(is_async):
|
|
from tests.langgraph.sample_code.langgraph_prebuilt import graph
|
|
|
|
mlflow.langchain.autolog()
|
|
|
|
input_example = {"messages": [{"role": "user", "content": "what is the weather in sf?"}]}
|
|
|
|
with mlflow.start_span("parent"):
|
|
if is_async:
|
|
await graph.ainvoke(input_example)
|
|
else:
|
|
graph.invoke(input_example)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert traces[0].info.status == "OK"
|
|
|
|
# Validate structure
|
|
span_id_to_span = {span.span_id: span for span in traces[0].data.spans}
|
|
tool_span = next(span for span in traces[0].data.spans if span.span_type == SpanType.TOOL)
|
|
assert tool_span.name == "get_weather"
|
|
|
|
tool_parent_span = span_id_to_span[tool_span.parent_id]
|
|
assert tool_parent_span.name == "tools"
|
|
assert tool_parent_span.span_type == SpanType.CHAIN
|
|
|
|
graph_span = span_id_to_span[tool_parent_span.parent_id]
|
|
assert graph_span.name == "LangGraph"
|
|
assert graph_span.span_type == SpanType.CHAIN
|
|
|
|
root_span = span_id_to_span[graph_span.parent_id]
|
|
assert root_span.name == "parent"
|
|
assert root_span.span_type == SpanType.UNKNOWN
|
|
|
|
|
|
@skip_when_testing_trace_sdk
|
|
def test_langgraph_chat_agent_trace():
|
|
input_example = {"messages": [{"role": "user", "content": "hi"}]}
|
|
|
|
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)
|
|
# No trace should be created for loading it in
|
|
assert mlflow.get_trace(mlflow.get_last_active_trace_id()) is None
|
|
|
|
loaded_model.predict(input_example)
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert traces[0].info.status == "OK"
|
|
assert traces[0].info.request_metadata[TraceMetadataKey.MODEL_ID] == model_info.model_id
|
|
assert traces[0].data.spans[0].name == "LangGraph"
|
|
assert traces[0].data.spans[0].inputs == input_example
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
list(loaded_model.predict_stream(input_example))
|
|
traces = get_traces()
|
|
assert len(traces) == 2
|
|
assert traces[0].info.status == "OK"
|
|
assert traces[0].info.request_metadata[TraceMetadataKey.MODEL_ID] == model_info.model_id
|
|
assert traces[0].data.spans[0].name == "LangGraph"
|
|
assert traces[0].data.spans[0].inputs == input_example
|
|
|
|
|
|
@skip_when_testing_trace_sdk
|
|
def test_langgraph_autolog_with_update_current_span():
|
|
model_info = mlflow.langchain.log_model(
|
|
lc_model="tests/langgraph/sample_code/langgraph_with_autolog.py",
|
|
input_example={"status": "done"},
|
|
)
|
|
assert model_info.signature is not None
|
|
assert model_info.signature.inputs is not None
|
|
assert model_info.signature.outputs is not None
|