feat: wrap LLM clients with TrackedLLMClient in LLMFactory
This commit is contained in:
@@ -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()
|
||||
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user