Kackle commited on
Commit
420ce0f
·
verified ·
1 Parent(s): 39d0c5d

Update gemini_agent.py

Browse files
Files changed (1) hide show
  1. gemini_agent.py +21 -42
gemini_agent.py CHANGED
@@ -19,7 +19,7 @@ class GeminiAgent:
19
  api_key = os.getenv('GOOGLE_API_KEY')
20
  genai.configure(api_key=api_key)
21
 
22
- self.model = genai.GenerativeModel('gemini-2.5-flash')
23
  self.last_request_time = 0
24
  self.min_request_interval = 8.0 # 7 seconds between requests (10 per minute limit, with margin)
25
 
@@ -163,38 +163,27 @@ Provide only the direct answer. If it's a quote, give just the quoted text. If i
163
  async def _handle_text_question(self, question: str) -> str:
164
  """Handle regular text-based questions"""
165
  prompt = ""
166
- # Use retrieval for explicit web/Wikipedia questions and factual queries
167
  def is_explicit_retrieval_question(question):
168
  q = question.lower()
169
- # Explicit retrieval requests
170
- explicit_triggers = [
171
- "according to wikipedia", "from wikipedia", "search the web",
172
- "duckduckgo", "web search"
173
- ]
174
- # Factual question patterns that benefit from Wikipedia
175
- factual_triggers = [
176
- "who played", "actor who", "actress who", "cast of", "starred in",
177
- "tv show", "television", "movie", "film", "character in",
178
- "version of", "adaptation of", "remake of"
179
- ]
180
- return any(trigger in q for trigger in explicit_triggers + factual_triggers)
181
  wiki_context = ""
182
  ddg_context = ""
183
  if is_explicit_retrieval_question(question):
184
- try:
185
- wiki_context = self.wiki_tool.run(question)
186
- print(f"Wiki context: {wiki_context[:200]}...") # Debug
187
- except Exception as e:
188
- print(f"Wikipedia tool failed: {e}")
189
-
190
- # Use DuckDuckGo as fallback if Wikipedia context is poor or for web search requests
191
- def is_poor_context(context):
192
- return not context or any(x in context.lower() for x in ["not found", "no results", "does not contain", "disambiguation"])
193
-
194
- if "duckduckgo" in question.lower() or "web search" in question.lower() or is_poor_context(wiki_context):
195
  try:
196
  ddg_context = self.ddg_tool.run(question)
197
- print(f"DDG context: {ddg_context[:200]}...") # Debug
198
  except Exception as e:
199
  print(f"DuckDuckGo tool failed: {e}")
200
  # Handle attached file questions with enhanced prompts
@@ -236,12 +225,9 @@ Provide only the direct answer. If it's a quote, give just the quoted text. If i
236
 
237
  # Prepend context to the prompt if available and likely relevant
238
  def is_good_context(context):
239
- return context and not any(x in context.lower() for x in ["not found", "no results", "does not contain", "disambiguation"])
240
  if wiki_context and is_good_context(wiki_context):
241
- if "first name" in question.lower():
242
- prompt = f"Based on this Wikipedia information, find the actor who played Ray/Roman and identify what character they played in Magda M., then give only their first name:\n{wiki_context}\n\n{question}\n\nAnswer:"
243
- else:
244
- prompt = f"Use the following Wikipedia context to answer the question:\n{wiki_context}\n\n{prompt}"
245
  elif ddg_context and is_good_context(ddg_context):
246
  prompt = f"Use the following web search context to answer the question:\n{ddg_context}\n\n{prompt}"
247
 
@@ -275,17 +261,10 @@ Provide only the direct answer. If it's a quote, give just the quoted text. If i
275
 
276
  # If the question expects a single value, extract it
277
  if any(kw in question.lower() for kw in ["how many", "what is the", "who", "where", "give only", "provide only"]):
278
- # Special handling for "give only the first name"
279
- if "give only the first name" in question.lower() or "first name" in question.lower():
280
- # Extract first name from full name
281
- name_match = re.search(r'\b([A-ZĄĆĘŁŃÓŚŹŻ][a-ząćęłńóśźż]+)', answer)
282
- if name_match:
283
- answer = name_match.group(1)
284
- else:
285
- # Extract the first number, word, or phrase (tweak regex as needed)
286
- match = re.search(r'^[A-Za-z0-9 ,+-]+', answer)
287
- if match:
288
- answer = match.group(0).strip()
289
 
290
  # Post-processing for chess move extraction
291
  if 'chess position' in question.lower() and 'image' in question.lower():
 
19
  api_key = os.getenv('GOOGLE_API_KEY')
20
  genai.configure(api_key=api_key)
21
 
22
+ self.model = genai.GenerativeModel('gemini-2.0-flash')
23
  self.last_request_time = 0
24
  self.min_request_interval = 8.0 # 7 seconds between requests (10 per minute limit, with margin)
25
 
 
163
  async def _handle_text_question(self, question: str) -> str:
164
  """Handle regular text-based questions"""
165
  prompt = ""
166
+ # Only use retrieval for explicit web/Wikipedia questions
167
  def is_explicit_retrieval_question(question):
168
  q = question.lower()
169
+ return (
170
+ "according to wikipedia" in q or
171
+ "from wikipedia" in q or
172
+ "search the web" in q or
173
+ "duckduckgo" in q or
174
+ "web search" in q
175
+ )
 
 
 
 
 
176
  wiki_context = ""
177
  ddg_context = ""
178
  if is_explicit_retrieval_question(question):
179
+ if "wikipedia" in question.lower():
180
+ try:
181
+ wiki_context = self.wiki_tool.run(question)
182
+ except Exception as e:
183
+ print(f"Wikipedia tool failed: {e}")
184
+ if "duckduckgo" in question.lower() or "web search" in question.lower():
 
 
 
 
 
185
  try:
186
  ddg_context = self.ddg_tool.run(question)
 
187
  except Exception as e:
188
  print(f"DuckDuckGo tool failed: {e}")
189
  # Handle attached file questions with enhanced prompts
 
225
 
226
  # Prepend context to the prompt if available and likely relevant
227
  def is_good_context(context):
228
+ return context and not any(x in context.lower() for x in ["not found", "no results", "does not contain information"])
229
  if wiki_context and is_good_context(wiki_context):
230
+ prompt = f"Use the following Wikipedia context to answer the question:\n{wiki_context}\n\n{prompt}"
 
 
 
231
  elif ddg_context and is_good_context(ddg_context):
232
  prompt = f"Use the following web search context to answer the question:\n{ddg_context}\n\n{prompt}"
233
 
 
261
 
262
  # If the question expects a single value, extract it
263
  if any(kw in question.lower() for kw in ["how many", "what is the", "who", "where", "give only", "provide only"]):
264
+ # Extract the first number, word, or phrase (tweak regex as needed)
265
+ match = re.search(r'^[A-Za-z0-9 ,+-]+', answer)
266
+ if match:
267
+ answer = match.group(0).strip()
 
 
 
 
 
 
 
268
 
269
  # Post-processing for chess move extraction
270
  if 'chess position' in question.lower() and 'image' in question.lower():