improve agent metrics integration test and cleanup fixtures

- simplified test to use telemetry.query_metrics for verification
- test now validates actual queryable metrics data
- verified by query metrics functionality added in #3074

tests/integration/agents/conftest.py

@@ -0,0 +1,170 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from collections.abc import AsyncGenerator, Callable
+from pathlib import Path
+from typing import Any
+from unittest.mock import Mock, patch
+
+import pytest
+
+from llama_stack.apis.inference import ToolDefinition
+from llama_stack.apis.tools import ToolInvocationResult
+from llama_stack.providers.inline.agents.meta_reference.agent_instance import ChatAgent
+from llama_stack.providers.inline.telemetry.meta_reference.config import (
+    TelemetryConfig,
+    TelemetrySink,
+)
+from llama_stack.providers.inline.telemetry.meta_reference.telemetry import (
+    TelemetryAdapter,
+)
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore.sqlite.sqlite import SqliteKVStoreImpl
+from llama_stack.providers.utils.telemetry import tracing as telemetry_tracing
+
+
+@pytest.fixture
+def make_agent_fixture():
+    def _make(telemetry, kvstore) -> ChatAgent:
+        agent = ChatAgent(
+            agent_id="test-agent",
+            agent_config=Mock(),
+            inference_api=Mock(),
+            safety_api=Mock(),
+            tool_runtime_api=Mock(),
+            tool_groups_api=Mock(),
+            vector_io_api=Mock(),
+            telemetry_api=telemetry,
+            persistence_store=kvstore,
+            created_at="2025-01-01T00:00:00Z",
+            policy=[],
+        )
+        agent.agent_config.client_tools = []
+        agent.agent_config.max_infer_iters = 5
+        agent.input_shields = []
+        agent.output_shields = []
+        agent.tool_defs = [
+            ToolDefinition(tool_name="web_search", description="", parameters={}),
+            ToolDefinition(tool_name="knowledge_search", description="", parameters={}),
+        ]
+        agent.tool_name_to_args = {}
+
+        # Stub tool runtime invoke_tool
+        async def _mock_invoke_tool(
+            *args: Any,
+            tool_name: str | None = None,
+            kwargs: dict | None = None,
+            **extra: Any,
+        ):
+            return ToolInvocationResult(content="Tool execution result")
+
+        agent.tool_runtime_api.invoke_tool = _mock_invoke_tool
+        return agent
+
+    return _make
+
+
+def _chat_stream(tool_name: str | None, content: str = ""):
+    from llama_stack.apis.common.content_types import (
+        TextDelta,
+        ToolCallDelta,
+        ToolCallParseStatus,
+    )
+    from llama_stack.apis.inference import (
+        ChatCompletionResponseEvent,
+        ChatCompletionResponseEventType,
+        ChatCompletionResponseStreamChunk,
+        StopReason,
+    )
+    from llama_stack.models.llama.datatypes import ToolCall
+
+    async def gen():
+        # Start
+        yield ChatCompletionResponseStreamChunk(
+            event=ChatCompletionResponseEvent(
+                event_type=ChatCompletionResponseEventType.start,
+                delta=TextDelta(text=""),
+            )
+        )
+
+        # Content
+        if content:
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=TextDelta(text=content),
+                )
+            )
+
+        # Tool call if specified
+        if tool_name:
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=ToolCallDelta(
+                        tool_call=ToolCall(call_id="call_0", tool_name=tool_name, arguments={}),
+                        parse_status=ToolCallParseStatus.succeeded,
+                    ),
+                )
+            )
+
+        # Complete
+        yield ChatCompletionResponseStreamChunk(
+            event=ChatCompletionResponseEvent(
+                event_type=ChatCompletionResponseEventType.complete,
+                delta=TextDelta(text=""),
+                stop_reason=StopReason.end_of_turn,
+            )
+        )
+
+    return gen()
+
+
+@pytest.fixture
+async def telemetry(tmp_path: Path) -> AsyncGenerator[TelemetryAdapter, None]:
+    db_path = tmp_path / "trace_store.db"
+    cfg = TelemetryConfig(
+        sinks=[TelemetrySink.CONSOLE, TelemetrySink.SQLITE],
+        sqlite_db_path=str(db_path),
+    )
+    telemetry = TelemetryAdapter(cfg, deps={})
+    telemetry_tracing.setup_logger(telemetry)
+    try:
+        yield telemetry
+    finally:
+        await telemetry.shutdown()
+
+
+@pytest.fixture
+async def kvstore(tmp_path: Path) -> SqliteKVStoreImpl:
+    kv_path = tmp_path / "agent_kvstore.db"
+    kv = SqliteKVStoreImpl(SqliteKVStoreConfig(db_path=str(kv_path)))
+    await kv.initialize()
+    return kv
+
+
+@pytest.fixture
+def span_patch():
+    with (
+        patch("llama_stack.providers.inline.agents.meta_reference.agent_instance.get_current_span") as mock_span,
+        patch(
+            "llama_stack.providers.utils.telemetry.tracing.generate_span_id",
+            return_value="0000000000000abc",
+        ),
+    ):
+        mock_span.return_value = Mock(get_span_context=Mock(return_value=Mock(trace_id=0x123, span_id=0xABC)))
+        yield
+
+
+@pytest.fixture
+def make_completion_fn() -> Callable[[str | None, str], Callable]:
+    def _factory(tool_name: str | None = None, content: str = "") -> Callable:
+        async def chat_completion(*args: Any, **kwargs: Any):
+            return _chat_stream(tool_name, content)
+
+        return chat_completion
+
+    return _factory

tests/integration/agents/test_agent_metrics_integration.py

@@ -5,55 +5,79 @@
 # the root directory of this source tree.
 
 import asyncio
-from unittest.mock import AsyncMock, Mock, patch
+from typing import Any
 
-from llama_stack.providers.inline.agents.meta_reference.agent_instance import ChatAgent
+from llama_stack.providers.utils.telemetry import tracing as telemetry_tracing
 
 
 class TestAgentMetricsIntegration:
-    """Smoke test for agent metrics integration"""
-
-    async def test_agent_metrics_methods_exist_and_work(self):
-        """Test that metrics methods exist and can be called without errors"""
-        # Create a minimal agent instance with mocked dependencies
-        telemetry_api = AsyncMock()
-        telemetry_api.logged_events = []
-
-        async def mock_log_event(event):
-            telemetry_api.logged_events.append(event)
-
-        telemetry_api.log_event = mock_log_event
-
-        agent = ChatAgent(
-            agent_id="test-agent",
-            agent_config=Mock(),
-            inference_api=Mock(),
-            safety_api=Mock(),
-            tool_runtime_api=Mock(),
-            tool_groups_api=Mock(),
-            vector_io_api=Mock(),
-            telemetry_api=telemetry_api,
-            persistence_store=Mock(),
-            created_at="2025-01-01T00:00:00Z",
-            policy=[],
+    async def test_agent_metrics_end_to_end(
+        self: Any,
+        telemetry: Any,
+        kvstore: Any,
+        make_agent_fixture: Any,
+        span_patch: Any,
+        make_completion_fn: Any,
+    ) -> None:
+        from llama_stack.apis.inference import (
+            SamplingParams,
+            UserMessage,
         )
 
-        with patch("llama_stack.providers.inline.agents.meta_reference.agent_instance.get_current_span") as mock_span:
-            mock_span.return_value = Mock(get_span_context=Mock(return_value=Mock(trace_id=123, span_id=456)))
+        agent: Any = make_agent_fixture(telemetry, kvstore)
 
-            # Test all metrics methods work
-            agent._track_step()
-            agent._track_workflow("completed", 2.5)
-            agent._track_tool("web_search")
+        session_id = await agent.create_session("s")
+        sampling_params = SamplingParams(max_tokens=64)
 
-            # Wait for async operations
-            await asyncio.sleep(0.01)
+        # single trace: plain, knowledge_search, web_search
+        await telemetry_tracing.start_trace("agent_metrics")
+        agent.inference_api.chat_completion = make_completion_fn(None, "Hello! I can help you with that.")
+        async for _ in agent.run(
+            session_id,
+            "t1",
+            [UserMessage(content="Hello")],
+            sampling_params,
+            stream=True,
+        ):
+            pass
+        agent.inference_api.chat_completion = make_completion_fn("knowledge_search", "")
+        async for _ in agent.run(
+            session_id,
+            "t2",
+            [UserMessage(content="Please search knowledge")],
+            sampling_params,
+            stream=True,
+        ):
+            pass
+        agent.inference_api.chat_completion = make_completion_fn("web_search", "")
+        async for _ in agent.run(
+            session_id,
+            "t3",
+            [UserMessage(content="Please search web")],
+            sampling_params,
+            stream=True,
+        ):
+            pass
+        await telemetry_tracing.end_trace()
 
-            # Basic verification that telemetry was called
-            assert len(telemetry_api.logged_events) >= 3
+        # Poll briefly to avoid flake with async persistence
+        tool_labels: set[str] = set()
+        for _ in range(10):
+            resp = await telemetry.query_metrics("llama_stack_agent_tool_calls_total", start_time=0, end_time=None)
+            tool_labels.clear()
+            for series in getattr(resp, "data", []) or []:
+                for lbl in getattr(series, "labels", []) or []:
+                    name = getattr(lbl, "name", None) or getattr(lbl, "key", None)
+                    value = getattr(lbl, "value", None)
+                    if name == "tool" and value:
+                        tool_labels.add(value)
 
-            # Verify we can call the methods without exceptions
-            agent._track_tool("knowledge_search")  # Test tool mapping
-            await asyncio.sleep(0.01)
+            # Look for both web_search AND some form of knowledge search
+            if ("web_search" in tool_labels) and ("rag" in tool_labels or "knowledge_search" in tool_labels):
+                break
+            await asyncio.sleep(0.1)
 
-            assert len(telemetry_api.logged_events) >= 4
+        # More descriptive assertion
+        assert bool(tool_labels & {"web_search", "rag", "knowledge_search"}), (
+            f"Expected tool calls not found. Got: {tool_labels}"
+        )