feat: wrap LLM clients with TrackedLLMClient in LLMFactory

This commit is contained in:
wangwei
2026-07-02 15:03:12 +08:00
parent d460397dda
commit 4fea159f5b
2 changed files with 56 additions and 5 deletions
+12 -5
View File
@@ -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()