backend/app/services/llm/llm_factory.py

@@ -7,6 +7,8 @@ from functools import lru_cache
 from .base_client import BaseLLMClient, LLMConfig, LLMProvider, LLMResponse
 from .deepseek_client import DeepSeekClient
 from .qwen_client import QwenClient, QwenVLClient
+from .tracked_client import TrackedLLMClient
+from app.shared.model_usage_tracker import get_model_usage_tracker
 # Keep provider-specific behavior explicit so debugging stays straightforward.
 
 
@@ -45,7 +47,7 @@ class LLMFactory:
         max_tokens: int = 4096,
         temperature: float = 0.7,
         **kwargs
-    ) -> BaseLLMClient:
+    ) -> "BaseLLMClient | TrackedLLMClient":
         """Handle create for the L L M Factory instance."""
         provider_enum = self._parse_provider(provider)
 
@@ -76,11 +78,16 @@ class LLMFactory:
         # Keep provider-specific behavior explicit so debugging stays straightforward.
         client = self._create_client(config)
 
+        # Wrap in TrackedLLMClient so every call site (agentic, HyDE, perception,
+        # compliance, document summarization, main answer generation) is recorded
+        # without each of them needing to know about usage tracking.
+        tracked_client = TrackedLLMClient(client, get_model_usage_tracker())
+
         # Keep provider-specific behavior explicit so debugging stays straightforward.
-        LLMFactory._global_instances[cache_key] = client
+        LLMFactory._global_instances[cache_key] = tracked_client
 
         logger.info(f"LLM客户端创建成功并缓存: {provider} - {model}")
-        return client
+        return tracked_client
 
     def _parse_provider(self, provider: str) -> LLMProvider:
         """Handle parse provider for this module for the L L M Factory instance."""
@@ -137,7 +144,7 @@ class LLMFactory:
 
         return client_class(config)
 
-    def get_cached(self, provider: str, model: Optional[str] = None) -> Optional[BaseLLMClient]:
+    def get_cached(self, provider: str, model: Optional[str] = None) -> "BaseLLMClient | TrackedLLMClient | None":
         """Return cached for the L L M Factory instance."""
         provider_enum = self._parse_provider(provider)
         model = model or DEFAULT_MODELS.get(provider_enum)
@@ -200,7 +207,7 @@ def get_llm_client(
     provider: str = "qwen",
     model: Optional[str] = None,
     **kwargs
-) -> BaseLLMClient:
+) -> "BaseLLMClient | TrackedLLMClient":
     """Return llm client."""
     factory = get_llm_factory()
 

backend/tests/observability/test_llm_factory_tracking.py

@@ -0,0 +1,44 @@
+"""Verifies get_llm_client() returns a usage-tracked client end to end."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from app.services.llm.llm_factory import LLMFactory, get_llm_client
+from app.services.llm.tracked_client import TrackedLLMClient
+from app.shared.model_usage_tracker import get_model_usage_tracker
+
+
+@pytest.fixture(autouse=True)
+def _reset_singletons():
+    """Clear the two process-wide singletons this test touches, before and after.
+
+    LLMFactory._global_instances and get_model_usage_tracker() both persist
+    for the life of the process; without this fixture, tests would leak
+    cached clients/usage data into each other and become order-dependent.
+    """
+    LLMFactory._global_instances.clear()
+    get_model_usage_tracker().snapshot()  # no-op read, just documents intent
+    get_model_usage_tracker()._entries.clear()
+    yield
+    LLMFactory._global_instances.clear()
+    get_model_usage_tracker()._entries.clear()
+
+
+def test_get_llm_client_returns_tracked_client():
+    """get_llm_client() must return a TrackedLLMClient, not the raw provider client."""
+    with patch("app.services.llm.llm_factory.DeepSeekClient") as mock_cls:
+        mock_cls.return_value = MagicMock()
+        client = get_llm_client(provider="deepseek", model="deepseek-v4-flash", api_key="test-key")
+    assert isinstance(client, TrackedLLMClient)
+
+
+def test_get_llm_client_caches_the_tracked_instance():
+    """A second call with the same provider/model must return the same TrackedLLMClient."""
+    with patch("app.services.llm.llm_factory.DeepSeekClient") as mock_cls:
+        mock_cls.return_value = MagicMock()
+        first = get_llm_client(provider="deepseek", model="deepseek-v4-flash", api_key="test-key")
+        second = get_llm_client(provider="deepseek", model="deepseek-v4-flash", api_key="test-key")
+    assert first is second