diff --git a/backend/app/services/llm/llm_factory.py b/backend/app/services/llm/llm_factory.py index 093439d..7b7e118 100644 --- a/backend/app/services/llm/llm_factory.py +++ b/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() diff --git a/backend/tests/observability/test_llm_factory_tracking.py b/backend/tests/observability/test_llm_factory_tracking.py new file mode 100644 index 0000000..9cdb5c0 --- /dev/null +++ b/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