Implement tool standard for Groq tracing (mlflow#14632)

Signed-off-by: Tomu Hirata <[email protected]>

Author: TomeHirata

mlflow/groq/_groq_autolog.py

@@ -2,6 +2,7 @@
 
 import mlflow
 from mlflow.entities import SpanType
+from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
 from mlflow.utils.autologging_utils.config import AutoLoggingConfig
 
 _logger = logging.getLogger(__name__)
@@ -23,6 +24,8 @@ def _get_span_type(resource: type) -> str:
 
 
 def patched_call(original, self, *args, **kwargs):
+    from groq.types.chat.chat_completion import ChatCompletion
+
     config = AutoLoggingConfig.init(flavor_name=mlflow.groq.FLAVOR_NAME)
 
     if config.log_traces:
@@ -31,6 +34,23 @@ def patched_call(original, self, *args, **kwargs):
             span_type=_get_span_type(self.__class__),
         ) as span:
             span.set_inputs(kwargs)
+
+            if tools := kwargs.get("tools"):
+                try:
+                    set_span_chat_tools(span, tools)
+                except Exception:
+                    _logger.debug(f"Failed to set tools for {span}.", exc_info=True)
+
             outputs = original(self, *args, **kwargs)
             span.set_outputs(outputs)
+
+            if isinstance(outputs, ChatCompletion):
+                try:
+                    messages = kwargs.get("messages", [])
+                    set_span_chat_messages(
+                        span, [*messages, outputs.choices[0].message.model_dump()]
+                    )
+                except Exception:
+                    _logger.debug(f"Failed to set chat messages for {span}.", exc_info=True)
+
             return outputs

mlflow/ml-package-versions.yml

@@ -967,7 +967,7 @@ groq:
       pip install git+https://github.com/groq/groq-python
   autologging:
     minimum: "0.13.0"
-    maximum: "0.15.0"
+    maximum: "0.18.0"
     requirements:
     run: pytest tests/groq
 

mlflow/ml_package_versions.py

@@ -381,7 +381,7 @@
         },
         "autologging": {
             "minimum": "0.13.0",
-            "maximum": "0.15.0"
+            "maximum": "0.18.0"
         }
     },
     "bedrock": {

tests/groq/test_groq_autolog.py

@@ -1,9 +1,11 @@
-import os
+import json
 from unittest.mock import patch
 
 import groq
+import pytest
 from groq.types.audio.transcription import Transcription
 from groq.types.audio.translation import Translation
+from groq.types.chat import ChatCompletionMessageToolCall
 from groq.types.chat.chat_completion import (
     ChatCompletion,
     ChatCompletionMessage,
@@ -22,6 +24,16 @@
     "messages": [{"role": "user", "content": "test message"}],
 }
 
+DUMMY_COMPLETION_USAGE = CompletionUsage(
+    completion_tokens=648,
+    prompt_tokens=20,
+    total_tokens=668,
+    completion_time=0.54,
+    prompt_time=0.000181289,
+    queue_time=0.012770949,
+    total_time=0.540181289,
+)
+
 DUMMY_CHAT_COMPLETION_RESPONSE = ChatCompletion(
     id="chatcmpl-test-id",
     choices=[
@@ -42,25 +54,24 @@
     model="llama3-8b-8192",
     object="chat.completion",
     system_fingerprint="fp_test",
-    usage=CompletionUsage(
-        completion_tokens=648,
-        prompt_tokens=20,
-        total_tokens=668,
-        completion_time=0.54,
-        prompt_time=0.000181289,
-        queue_time=0.012770949,
-        total_time=0.540181289,
-    ),
+    usage=DUMMY_COMPLETION_USAGE,
     x_groq={"id": "req_test"},
 )
 
 
-@patch.dict(os.environ, {"GROQ_API_KEY": "test_key"})
-@patch("groq._client.Groq.post", return_value=DUMMY_CHAT_COMPLETION_RESPONSE)
-def test_chat_completion_autolog(mock_post):
+@pytest.fixture(autouse=True)
+def init_state(monkeypatch):
+    monkeypatch.setenv("GROQ_API_KEY", "test_key")
+    yield
+    mlflow.groq.autolog(disable=True)
+
+
+def test_chat_completion_autolog():
     mlflow.groq.autolog()
     client = groq.Groq()
-    client.chat.completions.create(**DUMMY_CHAT_COMPLETION_REQUEST)
+
+    with patch("groq._client.Groq.post", return_value=DUMMY_CHAT_COMPLETION_RESPONSE):
+        client.chat.completions.create(**DUMMY_CHAT_COMPLETION_REQUEST)
 
     traces = get_traces()
     assert len(traces) == 1
@@ -74,13 +85,166 @@ def test_chat_completion_autolog(mock_post):
 
     mlflow.groq.autolog(disable=True)
     client = groq.Groq()
-    client.chat.completions.create(**DUMMY_CHAT_COMPLETION_REQUEST)
+
+    with patch("groq._client.Groq.post", return_value=DUMMY_CHAT_COMPLETION_RESPONSE):
+        client.chat.completions.create(**DUMMY_CHAT_COMPLETION_REQUEST)
 
     # No new trace should be created
     traces = get_traces()
     assert len(traces) == 1
 
 
+TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "calculate",
+            "description": "Evaluate a mathematical expression",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "expression": {
+                        "type": "string",
+                        "description": "The mathematical expression to evaluate",
+                    }
+                },
+                "required": ["expression"],
+            },
+        },
+    }
+]
+DUMMY_TOOL_CALL_REQUEST = {
+    "model": "test_model",
+    "max_tokens": 1024,
+    "messages": [{"role": "user", "content": "What is 25 * 4 + 10?"}],
+    "tools": TOOLS,
+}
+DUMMY_TOOL_CALL_RESPONSE = ChatCompletion(
+    id="chatcmpl-test-id",
+    choices=[
+        Choice(
+            finish_reason="stop",
+            index=0,
+            logprobs=None,
+            message=ChatCompletionMessage(
+                content=None,
+                role="assistant",
+                function_call=None,
+                tool_calls=[
+                    ChatCompletionMessageToolCall(
+                        id="tool call id",
+                        function={
+                            "name": "calculate",
+                            "arguments": json.dumps({"expression": "25 * 4 + 10"}),
+                        },
+                        type="function",
+                    )
+                ],
+                reasoning=None,
+            ),
+        )
+    ],
+    created=1733574047,
+    model="llama3-8b-8192",
+    object="chat.completion",
+    system_fingerprint="fp_test",
+    usage=DUMMY_COMPLETION_USAGE,
+    x_groq={"id": "req_test"},
+)
+
+
+def test_tool_calling_autolog():
+    mlflow.groq.autolog()
+    client = groq.Groq()
+
+    with patch("groq._client.Groq.post", return_value=DUMMY_TOOL_CALL_RESPONSE):
+        client.chat.completions.create(**DUMMY_TOOL_CALL_REQUEST)
+
+    traces = get_traces()
+    assert len(traces) == 1
+    assert traces[0].info.status == "OK"
+    assert len(traces[0].data.spans) == 1
+    span = traces[0].data.spans[0]
+    assert span.name == "Completions"
+    assert span.span_type == SpanType.CHAT_MODEL
+    assert span.inputs == DUMMY_TOOL_CALL_REQUEST
+    assert span.outputs == DUMMY_TOOL_CALL_RESPONSE.to_dict()
+    assert span.get_attribute("mlflow.chat.tools") == TOOLS
+    assert span.get_attribute("mlflow.chat.messages") == [
+        *DUMMY_TOOL_CALL_REQUEST["messages"],
+        DUMMY_TOOL_CALL_RESPONSE.choices[0].message.to_dict(),
+    ]
+
+
+DUMMY_TOOL_RESPONSE_REQUEST = {
+    "model": "test_model",
+    "max_tokens": 1024,
+    "messages": [
+        {"role": "user", "content": "What is 25 * 4 + 10?"},
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "id": "tool call id",
+                    "function": {
+                        "name": "calculate",
+                        "arguments": json.dumps({"expression": "25 * 4 + 10"}),
+                    },
+                    "type": "function",
+                }
+            ],
+        },
+        {"role": "tool", "name": "calculate", "content": json.dumps({"result": 110})},
+    ],
+    "tools": TOOLS,
+}
+DUMMY_TOOL_RESPONSE_RESPONSE = ChatCompletion(
+    id="chatcmpl-test-id",
+    choices=[
+        Choice(
+            finish_reason="stop",
+            index=0,
+            logprobs=None,
+            message=ChatCompletionMessage(
+                content="The result of the calculation is 110",
+                role="assistant",
+                function_call=None,
+                reasoning=None,
+                tool_calls=None,
+            ),
+        )
+    ],
+    created=1733574047,
+    model="llama3-8b-8192",
+    object="chat.completion",
+    system_fingerprint="fp_test",
+    usage=DUMMY_COMPLETION_USAGE,
+    x_groq={"id": "req_test"},
+)
+
+
+def test_tool_response_autolog():
+    mlflow.groq.autolog()
+    client = groq.Groq()
+
+    with patch("groq._client.Groq.post", return_value=DUMMY_TOOL_RESPONSE_RESPONSE):
+        client.chat.completions.create(**DUMMY_TOOL_RESPONSE_REQUEST)
+
+    traces = get_traces()
+    assert len(traces) == 1
+    assert traces[0].info.status == "OK"
+    assert len(traces[0].data.spans) == 1
+    span = traces[0].data.spans[0]
+    assert span.name == "Completions"
+    assert span.span_type == SpanType.CHAT_MODEL
+    assert span.inputs == DUMMY_TOOL_RESPONSE_REQUEST
+    assert span.outputs == DUMMY_TOOL_RESPONSE_RESPONSE.to_dict()
+    assert span.get_attribute("mlflow.chat.messages") == [
+        *DUMMY_TOOL_RESPONSE_REQUEST["messages"],
+        DUMMY_TOOL_RESPONSE_RESPONSE.choices[0].message.to_dict(),
+    ]
+
+
 BINARY_CONTENT = b"\x00\x00\x00\x14ftypM4A \x00\x00\x00\x00mdat\x00\x01\x02\x03"
 
 DUMMY_AUDIO_TRANSCRIPTION_REQUEST = {
@@ -91,12 +255,12 @@ def test_chat_completion_autolog(mock_post):
 DUMMY_AUDIO_TRANSCRIPTION_RESPONSE = Transcription(text="Test audio", x_groq={"id": "req_test"})
 
 
-@patch.dict(os.environ, {"GROQ_API_KEY": "test_key"})
-@patch("groq._client.Groq.post", return_value=DUMMY_AUDIO_TRANSCRIPTION_RESPONSE)
-def test_audio_transcription_autolog(mock_post):
+def test_audio_transcription_autolog():
     mlflow.groq.autolog()
     client = groq.Groq()
-    client.audio.transcriptions.create(**DUMMY_AUDIO_TRANSCRIPTION_REQUEST)
+
+    with patch("groq._client.Groq.post", return_value=DUMMY_AUDIO_TRANSCRIPTION_RESPONSE):
+        client.audio.transcriptions.create(**DUMMY_AUDIO_TRANSCRIPTION_REQUEST)
 
     traces = get_traces()
     assert len(traces) == 1
@@ -112,7 +276,9 @@ def test_audio_transcription_autolog(mock_post):
 
     mlflow.groq.autolog(disable=True)
     client = groq.Groq()
-    client.audio.transcriptions.create(**DUMMY_AUDIO_TRANSCRIPTION_REQUEST)
+
+    with patch("groq._client.Groq.post", return_value=DUMMY_AUDIO_TRANSCRIPTION_RESPONSE):
+        client.audio.transcriptions.create(**DUMMY_AUDIO_TRANSCRIPTION_REQUEST)
 
     # No new trace should be created
     traces = get_traces()
@@ -127,12 +293,12 @@ def test_audio_transcription_autolog(mock_post):
 DUMMY_AUDIO_TRANSLATION_RESPONSE = Translation(text="Test audio", x_groq={"id": "req_test"})
 
 
-@patch.dict(os.environ, {"GROQ_API_KEY": "test_key"})
-@patch("groq._client.Groq.post", return_value=DUMMY_AUDIO_TRANSLATION_RESPONSE)
-def test_audio_translation_autolog(mock_post):
+def test_audio_translation_autolog():
     mlflow.groq.autolog()
     client = groq.Groq()
-    client.audio.translations.create(**DUMMY_AUDIO_TRANSLATION_REQUEST)
+
+    with patch("groq._client.Groq.post", return_value=DUMMY_AUDIO_TRANSLATION_RESPONSE):
+        client.audio.translations.create(**DUMMY_AUDIO_TRANSLATION_REQUEST)
 
     traces = get_traces()
     assert len(traces) == 1
@@ -148,7 +314,9 @@ def test_audio_translation_autolog(mock_post):
 
     mlflow.groq.autolog(disable=True)
     client = groq.Groq()
-    client.audio.translations.create(**DUMMY_AUDIO_TRANSLATION_REQUEST)
+
+    with patch("groq._client.Groq.post", return_value=DUMMY_AUDIO_TRANSLATION_RESPONSE):
+        client.audio.translations.create(**DUMMY_AUDIO_TRANSLATION_REQUEST)
 
     # No new trace should be created
     traces = get_traces()