bpHigh commited on
Commit
5cde3d1
·
verified ·
1 Parent(s): f3d66a8

Update utils/google_genai_llm.py

Browse files
Files changed (1) hide show
  1. utils/google_genai_llm.py +80 -1
utils/google_genai_llm.py CHANGED
@@ -9,4 +9,83 @@ def get_response(prompt: str) -> str:
9
  response = client.models.generate_content(
10
  model="gemini-2.5-flash-preview-05-20", contents=prompt
11
  )
12
- return(response.text)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9
  response = client.models.generate_content(
10
  model="gemini-2.5-flash-preview-05-20", contents=prompt
11
  )
12
+ return(response.text)
13
+
14
+ def generate_with_gemini(
15
+ prompt: str,
16
+ purpose: str,
17
+ max_retries: int = 5,
18
+ retry_delay: float = 1.0
19
+ ) -> Optional[str]:
20
+ """
21
+ Generate text using Gemini API with fallback and retry logic.
22
+
23
+ Args:
24
+ prompt (str): The input prompt for text generation
25
+ api_key (str): Your Gemini API key
26
+ max_retries (int): Maximum number of retry attempts (default: 5)
27
+ retry_delay (float): Initial delay between retries in seconds (default: 1.0)
28
+
29
+ Returns:
30
+ Optional[str]: Generated text response or None if all attempts fail
31
+ """
32
+
33
+ print(f"Purpose: {purpose}")
34
+ # Model configurations with their respective max output tokens
35
+ models_config = [
36
+ {
37
+ "name": "gemini-2.5-flash-preview-05-20",
38
+ "max_output_tokens": 65536,
39
+ "description": "Gemini 2.5 Flash Preview"
40
+ },
41
+ {
42
+ "name": "gemini-2.0-flash",
43
+ "max_output_tokens": 8192,
44
+ "description": "Gemini 2.0 Flash"
45
+ }
46
+ ]
47
+
48
+ # Try each model with retry logic
49
+ for model_config in models_config:
50
+ model_name = model_config["name"]
51
+ max_tokens = model_config["max_output_tokens"]
52
+ model_desc = model_config["description"]
53
+
54
+ print(f"Attempting to use {model_desc} ({model_name})...")
55
+
56
+ for attempt in range(max_retries):
57
+ try:
58
+ # Create generation config
59
+ config = types.GenerateContentConfig(
60
+ max_output_tokens=max_tokens,
61
+ temperature=0.0
62
+ )
63
+
64
+ # Make the API call
65
+ response = client.models.generate_content(
66
+ model=model_name,
67
+ contents=[prompt],
68
+ config=config
69
+ )
70
+
71
+ # Check if response has content
72
+ if response and response.text:
73
+ print(f"✅ Success with {model_desc} on attempt {attempt + 1}")
74
+ return response.text
75
+ else:
76
+ print(f"⚠️ Empty response from {model_desc} on attempt {attempt + 1}")
77
+
78
+ except Exception as e:
79
+ print(f"❌ Error with {model_desc} on attempt {attempt + 1}: {str(e)}")
80
+
81
+ # If this is not the last attempt, wait before retrying
82
+ if attempt < max_retries - 1:
83
+ wait_time = retry_delay * (2 ** attempt) # Exponential backoff
84
+ print(f"⏳ Waiting {wait_time:.1f} seconds before retry...")
85
+ time.sleep(wait_time)
86
+ else:
87
+ print(f"💥 All {max_retries} attempts failed for {model_desc}")
88
+ break
89
+
90
+ print("❌ All models and retry attempts exhausted. Unable to generate response.")
91
+ return None