Coverage for src/local_deep_research/metrics/pricing/pricing_fetcher.py: 98%

1""" 

2LLM Pricing Data Fetcher 

3 

4Fetches real-time pricing data from various LLM providers. 

5Supports multiple providers and fallback to static pricing. 

6""" 

7 

8from typing import Dict, Optional 

9 

10import aiohttp 

11from loguru import logger 

12 

13from ...llm.providers.base import normalize_provider 

14 

15 

16class PricingFetcher: 

17 """Fetches LLM pricing data from various sources.""" 

18 

19 def __init__(self): 

20 self.session = None 

21 self.static_pricing = self._load_static_pricing() 

22 

23 async def __aenter__(self): 

24 self.session = aiohttp.ClientSession() 

25 return self 

26 

27 async def __aexit__(self, exc_type, exc_val, exc_tb): 

28 if self.session: 

29 await self.session.close() 

30 

31 def _load_static_pricing(self) -> Dict[str, Dict[str, float]]: 

32 """Load static pricing as fallback (per 1K tokens in USD).""" 

33 return { 

34 # OpenAI Models 

35 "gpt-4": {"prompt": 0.03, "completion": 0.06}, 

36 "gpt-4-turbo": {"prompt": 0.01, "completion": 0.03}, 

37 "gpt-4o": {"prompt": 0.005, "completion": 0.015}, 

38 "gpt-4o-mini": {"prompt": 0.00015, "completion": 0.0006}, 

39 "gpt-3.5-turbo": {"prompt": 0.001, "completion": 0.002}, 

40 # Anthropic Models 

41 "claude-3-opus": {"prompt": 0.015, "completion": 0.075}, 

42 "claude-3-sonnet": {"prompt": 0.003, "completion": 0.015}, 

43 "claude-3-haiku": {"prompt": 0.00025, "completion": 0.00125}, 

44 "claude-3-5-sonnet": {"prompt": 0.003, "completion": 0.015}, 

45 # Google Models 

46 "gemini-pro": {"prompt": 0.0005, "completion": 0.0015}, 

47 "gemini-pro-vision": {"prompt": 0.0005, "completion": 0.0015}, 

48 "gemini-1.5-pro": {"prompt": 0.0035, "completion": 0.0105}, 

49 "gemini-1.5-flash": {"prompt": 0.00035, "completion": 0.00105}, 

50 # Local/Open Source (free) 

51 "ollama": {"prompt": 0.0, "completion": 0.0}, 

52 "llama": {"prompt": 0.0, "completion": 0.0}, 

53 "mistral": {"prompt": 0.0, "completion": 0.0}, 

54 "gemma": {"prompt": 0.0, "completion": 0.0}, 

55 "qwen": {"prompt": 0.0, "completion": 0.0}, 

56 "codellama": {"prompt": 0.0, "completion": 0.0}, 

57 "vicuna": {"prompt": 0.0, "completion": 0.0}, 

58 "alpaca": {"prompt": 0.0, "completion": 0.0}, 

59 "lmstudio": {"prompt": 0.0, "completion": 0.0}, 

60 "llamacpp": {"prompt": 0.0, "completion": 0.0}, 

61 } 

62 

63 async def get_model_pricing( 

64 self, model_name: str, provider: str = None 

65 ) -> Optional[Dict[str, float]]: 

66 """Get pricing for a specific model and provider.""" 

67 return self.get_static_model_pricing(model_name, provider) 

68 

69 def get_static_model_pricing( 

70 self, model_name: str, provider: Optional[str] = None 

71 ) -> Optional[Dict[str, float]]: 

72 """Resolve provider-aware fallback pricing without network I/O.""" 

73 # Normalize inputs 

74 model_name = model_name.lower() if model_name else "" 

75 provider = normalize_provider(provider) or "" 

76 

77 # Provider-first approach: Check if provider indicates local/free models 

78 local_providers = ["ollama", "lmstudio", "llamacpp"] 

79 if provider in local_providers: 

80 logger.debug( 

81 f"Local provider '{provider}' detected - returning zero cost" 

82 ) 

83 return {"prompt": 0.0, "completion": 0.0} 

84 

85 # Fallback to static pricing with provider priority 

86 if provider: 

87 # First try provider-specific lookup with exact matching 

88 provider_models = self._get_models_by_provider(provider) 

89 # Try exact match 

90 if model_name in provider_models: 

91 return provider_models[model_name] 

92 # Try exact match without provider prefix 

93 if "/" in model_name: 

94 model_only = model_name.split("/")[-1] 

95 if model_only in provider_models: 

96 return provider_models[model_only] 

97 

98 # Exact model name matching only 

99 # First try exact match 

100 if model_name in self.static_pricing: 

101 return self.static_pricing[model_name] 

102 

103 # Try exact match without provider prefix (e.g., "openai/gpt-4o-mini" -> "gpt-4o-mini") 

104 if "/" in model_name: 

105 model_only = model_name.split("/")[-1] 

106 if model_only in self.static_pricing: 

107 return self.static_pricing[model_only] 

108 

109 # No pricing found - return None instead of default pricing 

110 logger.warning( 

111 f"No pricing found for model: {model_name}, provider: {provider}" 

112 ) 

113 return None 

114 

115 def _get_models_by_provider( 

116 self, provider: str 

117 ) -> Dict[str, Dict[str, float]]: 

118 """Get models for a specific provider.""" 

119 provider = normalize_provider(provider) or "" 

120 provider_models = {} 

121 

122 if provider == "openai": 

123 provider_models = { 

124 k: v 

125 for k, v in self.static_pricing.items() 

126 if k.startswith("gpt") 

127 } 

128 elif provider == "anthropic": 

129 provider_models = { 

130 k: v 

131 for k, v in self.static_pricing.items() 

132 if k.startswith("claude") 

133 } 

134 elif provider == "google": 

135 provider_models = { 

136 k: v 

137 for k, v in self.static_pricing.items() 

138 if k.startswith("gemini") 

139 } 

140 elif provider in ["ollama", "lmstudio", "llamacpp"]: 

141 # All local models are free 

142 provider_models = { 

143 k: v 

144 for k, v in self.static_pricing.items() 

145 if v["prompt"] == 0.0 and v["completion"] == 0.0 

146 } 

147 

148 return provider_models 

149 

150 def get_provider_from_model(self, model_name: str) -> str: 

151 """Determine the provider from model name.""" 

152 model_name = model_name.lower() 

153 

154 if "gpt" in model_name or "openai" in model_name: 

155 return "openai" 

156 if "claude" in model_name or "anthropic" in model_name: 

157 return "anthropic" 

158 if "gemini" in model_name or "google" in model_name: 

159 return "google" 

160 if "llama" in model_name or "meta" in model_name: 

161 return "meta" 

162 if "mistral" in model_name: 

163 return "mistral" 

164 if "ollama" in model_name: 164 ↛ 165line 164 didn't jump to line 165 because the condition on line 164 was never true

165 return "ollama" 

166 return "unknown"