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

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 # Normalize inputs 

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

69 provider = normalize_provider(provider) or "" 

70 

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

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

73 if provider in local_providers: 

74 logger.debug( 

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

76 ) 

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

78 

79 # Fallback to static pricing with provider priority 

80 if provider: 

81 # First try provider-specific lookup with exact matching 

82 provider_models = self._get_models_by_provider(provider) 

83 # Try exact match 

84 if model_name in provider_models: 

85 return provider_models[model_name] 

86 # Try exact match without provider prefix 

87 if "/" in model_name: 87 ↛ 94line 87 didn't jump to line 94 because the condition on line 87 was always true

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

89 if model_only in provider_models: 89 ↛ 94line 89 didn't jump to line 94 because the condition on line 89 was always true

90 return provider_models[model_only] 

91 

92 # Exact model name matching only 

93 # First try exact match 

94 if model_name in self.static_pricing: 

95 return self.static_pricing[model_name] 

96 

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

98 if "/" in model_name: 

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

100 if model_only in self.static_pricing: 100 ↛ 104line 100 didn't jump to line 104 because the condition on line 100 was always true

101 return self.static_pricing[model_only] 

102 

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

104 logger.warning( 

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

106 ) 

107 return None 

108 

109 def _get_models_by_provider( 

110 self, provider: str 

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

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

113 provider = normalize_provider(provider) or "" 

114 provider_models = {} 

115 

116 if provider == "openai": 

117 provider_models = { 

118 k: v 

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

120 if k.startswith("gpt") 

121 } 

122 elif provider == "anthropic": 

123 provider_models = { 

124 k: v 

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

126 if k.startswith("claude") 

127 } 

128 elif provider == "google": 

129 provider_models = { 

130 k: v 

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

132 if k.startswith("gemini") 

133 } 

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

135 # All local models are free 

136 provider_models = { 

137 k: v 

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

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

140 } 

141 

142 return provider_models 

143 

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

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

146 model_name = model_name.lower() 

147 

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

149 return "openai" 

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

151 return "anthropic" 

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

153 return "google" 

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

155 return "meta" 

156 if "mistral" in model_name: 

157 return "mistral" 

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

159 return "ollama" 

160 return "unknown"