delightfulrachel commited on
Commit
30d2304
·
verified ·
1 Parent(s): cbf016f

Update api_client.py

Browse files
Files changed (1) hide show
  1. api_client.py +275 -126
api_client.py CHANGED
@@ -1,16 +1,38 @@
1
- """API client functions for LLM interactions"""
 
 
2
 
3
  import os
4
- import time
5
- import requests
6
- import hashlib
7
- from functools import lru_cache
8
- from typing import Optional
9
  import logging
 
10
 
 
 
 
 
 
 
 
 
 
11
  logger = logging.getLogger(__name__)
12
 
13
- # Model lists
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  together_models = [
15
  "Qwen/Qwen2.5-Coder-32B-Instruct",
16
  "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
@@ -18,153 +40,280 @@ together_models = [
18
  "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
19
  ]
20
 
21
- anthropic_models = [
22
- "claude-3-7-sonnet-20250219",
23
- "claude-3-haiku-20240307",
24
- "claude-opus-4-20250514",
25
- "claude-sonnet-4-20250514"
26
- ]
27
-
28
- all_models = together_models + anthropic_models
29
 
30
- def get_api_key(provider: str) -> str:
31
- """Securely retrieve API key for the specified provider."""
 
 
 
 
 
 
 
 
 
 
 
32
  try:
33
- if provider == "together":
34
- api_key = os.getenv("TOGETHER_API_KEY")
35
- if not api_key:
36
- raise ValueError("API key not configured. Please contact administrator.")
37
- return api_key
38
- elif provider == "anthropic":
39
- api_key = os.getenv("ANTHROPIC_API_KEY")
40
- if not api_key:
41
- raise ValueError("API key not configured. Please contact administrator.")
42
- return api_key
43
  else:
44
- raise ValueError(f"Unknown provider: {provider}")
45
  except Exception as e:
46
- logger.error(f"Error retrieving API key: {e}")
47
- raise
48
-
49
- def get_provider(model: str) -> str:
50
- """Determine the provider for a given model."""
51
- if model in together_models:
52
- return "together"
53
- elif model in anthropic_models:
54
- return "anthropic"
55
- else:
56
- raise ValueError(f"Unknown model: {model}")
57
 
58
- def call_api_with_retry(api_func, *args, max_retries: int = 3, timeout: int = 30, **kwargs):
59
- """Call API with retry logic and timeout."""
60
- from utils import handle_api_error
61
-
62
- for attempt in range(max_retries):
63
- try:
64
- kwargs['timeout'] = timeout
65
- return api_func(*args, **kwargs)
66
- except requests.Timeout:
67
- if attempt == max_retries - 1:
68
- return "Request timed out. Please try again with a shorter input."
69
- except requests.ConnectionError:
70
- if attempt == max_retries - 1:
71
- return "Connection error. Please check your internet connection."
72
- except Exception as e:
73
- if attempt == max_retries - 1:
74
- return f"Error: {str(e)}"
75
- time.sleep(2 ** attempt) # Exponential backoff
76
-
77
- def call_together_api(model: str, prompt: str, temperature: float = 0.7, max_tokens: int = 1500) -> str:
78
- """Call Together AI API with enhanced error handling."""
79
- from utils import handle_api_error
80
 
81
- api_key = get_api_key("together")
82
- system_message = (
83
- "You are a Salesforce B2B Commerce expert. Be CONCISE and PRECISE. "
84
- "Focus on CODE QUALITY over explanations. Use structured formats when requested. "
85
- "Always check for syntax errors, security issues, and performance problems."
86
- )
 
 
 
 
 
 
 
 
 
87
 
88
- def make_request():
89
  headers = {
90
- "Authorization": f"Bearer {api_key}",
91
- "Content-Type": "application/json"
 
92
  }
 
93
  payload = {
94
  "model": model,
95
- "messages": [
96
- {"role": "system", "content": system_message},
97
- {"role": "user", "content": prompt}
98
- ],
99
- "temperature": temperature,
100
  "max_tokens": max_tokens,
101
- "top_p": 0.9
 
 
 
 
 
 
102
  }
103
- resp = requests.post(
104
- "https://api.together.xyz/v1/chat/completions",
 
105
  headers=headers,
106
  json=payload,
107
- timeout=30
108
  )
109
- if resp.status_code != 200:
110
- return handle_api_error(resp.status_code, resp.text)
111
- data = resp.json()
112
- return data["choices"][0]["message"]["content"]
113
-
114
- return call_api_with_retry(make_request)
 
 
 
 
115
 
116
- def call_anthropic_api(model: str, prompt: str, temperature: float = 0.7, max_tokens: int = 1500) -> str:
117
- """Call Anthropic API with enhanced error handling."""
118
- from utils import handle_api_error
119
 
120
- api_key = get_api_key("anthropic")
121
- system_message = (
122
- "You are a Salesforce B2B Commerce expert. Be CONCISE and PRECISE. "
123
- "Focus on CODE QUALITY over explanations. Use structured formats when requested. "
124
- "Always check for syntax errors, security issues, and performance problems."
125
- )
 
 
 
 
 
 
 
 
 
126
 
127
- def make_request():
128
  headers = {
129
- "x-api-key": api_key,
130
- "anthropic-version": "2023-06-01",
131
- "content-type": "application/json"
132
  }
 
133
  payload = {
134
  "model": model,
135
- "system": system_message,
136
- "messages": [
137
- {"role": "user", "content": prompt}
138
- ],
139
  "temperature": temperature,
140
- "max_tokens": max_tokens
 
 
 
 
 
141
  }
142
- resp = requests.post(
143
- "https://api.anthropic.com/v1/messages",
 
144
  headers=headers,
145
  json=payload,
146
- timeout=30
147
  )
148
- if resp.status_code != 200:
149
- return handle_api_error(resp.status_code, resp.text)
150
- data = resp.json()
151
- return data["content"][0]["text"]
 
 
 
 
 
 
 
 
 
 
152
 
153
- return call_api_with_retry(make_request)
154
-
155
- @lru_cache(maxsize=100)
156
- def cached_llm_call(model_hash: str, prompt_hash: str, model: str, prompt: str, temperature: float = 0.7, max_tokens: int = 1500) -> str:
157
- """Cached LLM call to avoid repeated API calls for same inputs."""
158
- provider = get_provider(model)
159
- if provider == "together":
160
- return call_together_api(model, prompt, temperature, max_tokens)
161
- elif provider == "anthropic":
162
- return call_anthropic_api(model, prompt, temperature, max_tokens)
 
 
 
 
 
 
 
 
163
  else:
164
- return f"Error: Unknown provider for model {model}"
 
 
165
 
166
- def call_llm(model: str, prompt: str, temperature: float = 0.7, max_tokens: int = 1500) -> str:
167
- """Call LLM with caching support."""
168
- model_hash = hashlib.md5(model.encode()).hexdigest()
169
- prompt_hash = hashlib.md5(prompt.encode()).hexdigest()
170
- return cached_llm_call(model_hash, prompt_hash, model, prompt, temperature, max_tokens)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ API client for LLM providers (Anthropic, Together.ai, etc.)
3
+ """
4
 
5
  import os
 
 
 
 
 
6
  import logging
7
+ from typing import Dict, Any
8
 
9
+ # Try to import requests, but gracefully handle if not available
10
+ try:
11
+ import requests
12
+ REQUESTS_AVAILABLE = True
13
+ except ImportError:
14
+ REQUESTS_AVAILABLE = False
15
+ print("Warning: requests module not available. API calls will use mock responses.")
16
+
17
+ # Configure logging
18
  logger = logging.getLogger(__name__)
19
 
20
+ # Model configurations
21
+ anthropic_models = [
22
+ "claude-3-5-sonnet-20241022",
23
+ "claude-3-sonnet-20240229",
24
+ "claude-3-haiku-20240307",
25
+ "claude-opus-4-20250514",
26
+ "claude-sonnet-4-20250514"
27
+ ]
28
+
29
+ # together_models = [
30
+ # "meta-llama/Llama-2-70b-chat-hf",
31
+ # "mistralai/Mixtral-8x7B-Instruct-v0.1",
32
+ # "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
33
+ # "teknium/OpenHermes-2.5-Mistral-7B",
34
+ # "microsoft/DialoGPT-medium"
35
+ # ]
36
  together_models = [
37
  "Qwen/Qwen2.5-Coder-32B-Instruct",
38
  "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
 
40
  "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
41
  ]
42
 
43
+ # Combined list of all available models
44
+ all_models = anthropic_models + together_models
 
 
 
 
 
 
45
 
46
+ def call_llm(model: str, prompt: str, temperature: float = 0.3, max_tokens: int = 4000) -> str:
47
+ """
48
+ Call the specified LLM model with the given prompt.
49
+
50
+ Args:
51
+ model: Model name to use
52
+ prompt: Input prompt
53
+ temperature: Sampling temperature
54
+ max_tokens: Maximum tokens to generate
55
+
56
+ Returns:
57
+ Model response as string
58
+ """
59
  try:
60
+ if model in anthropic_models:
61
+ return call_anthropic(model, prompt, temperature, max_tokens)
62
+ elif model in together_models:
63
+ return call_together_ai(model, prompt, temperature, max_tokens)
 
 
 
 
 
 
64
  else:
65
+ return f"Unsupported model: {model}"
66
  except Exception as e:
67
+ logger.error(f"Error calling model {model}: {str(e)}")
68
+ return f"Error calling model {model}: {str(e)}"
 
 
 
 
 
 
 
 
 
69
 
70
+ def call_anthropic(model: str, prompt: str, temperature: float = 0.3, max_tokens: int = 4000) -> str:
71
+ """
72
+ Call Anthropic Claude API.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
73
 
74
+ Args:
75
+ model: Claude model name
76
+ prompt: Input prompt
77
+ temperature: Sampling temperature
78
+ max_tokens: Maximum tokens to generate
79
+
80
+ Returns:
81
+ Model response
82
+ """
83
+ if not REQUESTS_AVAILABLE:
84
+ return mock_llm_response(model, prompt)
85
+
86
+ api_key = os.getenv("ANTHROPIC_API_KEY")
87
+ if not api_key:
88
+ return mock_llm_response(model, prompt)
89
 
90
+ try:
91
  headers = {
92
+ "Content-Type": "application/json",
93
+ "x-api-key": api_key,
94
+ "anthropic-version": "2023-06-01"
95
  }
96
+
97
  payload = {
98
  "model": model,
 
 
 
 
 
99
  "max_tokens": max_tokens,
100
+ "temperature": temperature,
101
+ "messages": [
102
+ {
103
+ "role": "user",
104
+ "content": prompt
105
+ }
106
+ ]
107
  }
108
+
109
+ response = requests.post(
110
+ "https://api.anthropic.com/v1/messages",
111
  headers=headers,
112
  json=payload,
113
+ timeout=60
114
  )
115
+
116
+ if response.status_code == 200:
117
+ result = response.json()
118
+ return result["content"][0]["text"]
119
+ else:
120
+ return f"Anthropic API error: {response.status_code} - {response.text}"
121
+
122
+ except Exception as e:
123
+ logger.warning(f"Anthropic API call failed: {str(e)}, using mock response")
124
+ return mock_llm_response(model, prompt)
125
 
126
+ def call_together_ai(model: str, prompt: str, temperature: float = 0.3, max_tokens: int = 4000) -> str:
127
+ """
128
+ Call Together.ai API.
129
 
130
+ Args:
131
+ model: Together.ai model name
132
+ prompt: Input prompt
133
+ temperature: Sampling temperature
134
+ max_tokens: Maximum tokens to generate
135
+
136
+ Returns:
137
+ Model response
138
+ """
139
+ if not REQUESTS_AVAILABLE:
140
+ return mock_llm_response(model, prompt)
141
+
142
+ api_key = os.getenv("TOGETHER_API_KEY")
143
+ if not api_key:
144
+ return mock_llm_response(model, prompt)
145
 
146
+ try:
147
  headers = {
148
+ "Authorization": f"Bearer {api_key}",
149
+ "Content-Type": "application/json"
 
150
  }
151
+
152
  payload = {
153
  "model": model,
154
+ "max_tokens": max_tokens,
 
 
 
155
  "temperature": temperature,
156
+ "messages": [
157
+ {
158
+ "role": "user",
159
+ "content": prompt
160
+ }
161
+ ]
162
  }
163
+
164
+ response = requests.post(
165
+ "https://api.together.xyz/v1/chat/completions",
166
  headers=headers,
167
  json=payload,
168
+ timeout=60
169
  )
170
+
171
+ if response.status_code == 200:
172
+ result = response.json()
173
+ return result["choices"][0]["message"]["content"]
174
+ else:
175
+ return f"Together.ai API error: {response.status_code} - {response.text}"
176
+
177
+ except Exception as e:
178
+ logger.warning(f"Together.ai API call failed: {str(e)}, using mock response")
179
+ return mock_llm_response(model, prompt)
180
+
181
+ def test_model_connectivity() -> Dict[str, str]:
182
+ """
183
+ Test connectivity to different model providers.
184
 
185
+ Returns:
186
+ Dictionary with model provider status
187
+ """
188
+ results = {}
189
+
190
+ # Test Anthropic
191
+ anthropic_key = os.getenv("ANTHROPIC_API_KEY")
192
+ if anthropic_key:
193
+ test_response = call_anthropic(anthropic_models[0], "Hello", 0.1, 10)
194
+ results["anthropic"] = "Connected" if not test_response.startswith("Error") else f"Failed: {test_response}"
195
+ else:
196
+ results["anthropic"] = "API key not set"
197
+
198
+ # Test Together.ai
199
+ together_key = os.getenv("TOGETHER_API_KEY")
200
+ if together_key:
201
+ test_response = call_together_ai(together_models[0], "Hello", 0.1, 10)
202
+ results["together_ai"] = "Connected" if not test_response.startswith("Error") else f"Failed: {test_response}"
203
  else:
204
+ results["together_ai"] = "API key not set"
205
+
206
+ return results
207
 
208
+ def get_model_info(model: str) -> Dict[str, Any]:
209
+ """
210
+ Get information about a specific model.
211
+
212
+ Args:
213
+ model: Model name
214
+
215
+ Returns:
216
+ Dictionary with model information
217
+ """
218
+ if model in anthropic_models:
219
+ return {
220
+ "provider": "Anthropic",
221
+ "model": model,
222
+ "type": "Chat",
223
+ "max_tokens": 4096,
224
+ "supports_functions": True
225
+ }
226
+ elif model in together_models:
227
+ return {
228
+ "provider": "Together.ai",
229
+ "model": model,
230
+ "type": "Chat",
231
+ "max_tokens": 4096,
232
+ "supports_functions": False
233
+ }
234
+ else:
235
+ return {
236
+ "provider": "Unknown",
237
+ "model": model,
238
+ "error": "Model not found"
239
+ }
240
+
241
+ def validate_api_keys() -> Dict[str, bool]:
242
+ """
243
+ Validate that required API keys are set.
244
+
245
+ Returns:
246
+ Dictionary with API key validation status
247
+ """
248
+ return {
249
+ "anthropic": bool(os.getenv("ANTHROPIC_API_KEY")),
250
+ "together_ai": bool(os.getenv("TOGETHER_API_KEY"))
251
+ }
252
+
253
+ # Mock functions for testing when API keys are not available
254
+ def mock_llm_response(model: str, prompt: str) -> str:
255
+ """
256
+ Generate a mock response for testing purposes.
257
+
258
+ Args:
259
+ model: Model name
260
+ prompt: Input prompt
261
+
262
+ Returns:
263
+ Mock response
264
+ """
265
+ return f"""## CORRECTED CODE
266
+ ```apex
267
+ // This is a mock response for model: {model}
268
+ trigger MockTrigger on Account (before insert, before update) {{
269
+ // Mock corrected trigger logic
270
+ for (Account acc : Trigger.new) {{
271
+ if (Trigger.isInsert) {{
272
+ // Insert logic
273
+ }}
274
+ if (Trigger.isUpdate) {{
275
+ // Update logic
276
+ }}
277
+ }}
278
+ }}
279
+ ```
280
+
281
+ ## KEY CHANGES
282
+ - Added proper trigger context checks
283
+ - Implemented bulkification patterns
284
+ - Added error handling
285
+
286
+ ## CRITICAL ISSUES FIXED
287
+ 1. Missing trigger context: Added Trigger.isInsert/isUpdate checks
288
+ 2. Governor limits: Implemented proper bulkification
289
+ 3. Error handling: Added try-catch blocks
290
+
291
+ ## REMAINING WARNINGS
292
+ - Test coverage needed for all scenarios
293
+ - Consider adding custom metadata for configuration
294
+ """
295
+
296
+ # Use mock responses if API keys are not available
297
+ def call_llm_with_fallback(model: str, prompt: str, temperature: float = 0.3, max_tokens: int = 4000) -> str:
298
+ """
299
+ Call LLM with fallback to mock response if API keys are not available.
300
+
301
+ Args:
302
+ model: Model name
303
+ prompt: Input prompt
304
+ temperature: Sampling temperature
305
+ max_tokens: Maximum tokens to generate
306
+
307
+ Returns:
308
+ Model response or mock response
309
+ """
310
+ api_keys = validate_api_keys()
311
+
312
+ if model in anthropic_models and not api_keys["anthropic"]:
313
+ logger.warning("Anthropic API key not available, using mock response")
314
+ return mock_llm_response(model, prompt)
315
+ elif model in together_models and not api_keys["together_ai"]:
316
+ logger.warning("Together.ai API key not available, using mock response")
317
+ return mock_llm_response(model, prompt)
318
+
319
+ return call_llm(model, prompt, temperature, max_tokens)