Spaces:
Sleeping
Sleeping
Final
Browse files- gemini_agent.py +1 -64
gemini_agent.py
CHANGED
@@ -339,66 +339,7 @@ class GeminiAgent:
|
|
339 |
# Load answer bank
|
340 |
self._load_answer_bank()
|
341 |
|
342 |
-
|
343 |
-
"""Load the answer bank from JSON file."""
|
344 |
-
try:
|
345 |
-
ans_bank_path = os.path.join(os.path.dirname(__file__), 'kndg_info.json')
|
346 |
-
with open(ans_bank_path, 'r') as f:
|
347 |
-
self.answer_bank = json.load(f)
|
348 |
-
except Exception as e:
|
349 |
-
print(f"Warning: Could not load answer bank: {e}")
|
350 |
-
self.answer_bank = []
|
351 |
-
|
352 |
-
def _check_answer_bank(self, query: str) -> Optional[str]:
|
353 |
-
"""Check if query matches any question in answer bank using LLM with retries."""
|
354 |
-
max_retries = 5
|
355 |
-
base_sleep = 1
|
356 |
-
|
357 |
-
for attempt in range(max_retries):
|
358 |
-
try:
|
359 |
-
if not self.answer_bank:
|
360 |
-
return None
|
361 |
-
|
362 |
-
# Filter questions with answer_score = 1
|
363 |
-
valid_questions = [entry for entry in self.answer_bank if entry.get('answer_score', 0) == 1]
|
364 |
-
if not valid_questions:
|
365 |
-
return None
|
366 |
-
|
367 |
-
# Create a prompt for the LLM to compare the query with answer bank questions
|
368 |
-
prompt = f"""Given a user query and a list of reference questions, determine if the query is semantically similar to any of the reference questions.
|
369 |
-
Consider them similar if they are asking for the same information, even if phrased differently.
|
370 |
-
|
371 |
-
User Query: {query}
|
372 |
-
|
373 |
-
Reference Questions:
|
374 |
-
{json.dumps([{'id': i, 'question': q['question']} for i, q in enumerate(valid_questions)], indent=2)}
|
375 |
-
|
376 |
-
Instructions:
|
377 |
-
1. Compare the user query with each reference question
|
378 |
-
2. If there is a semantically similar match (asking for the same information), return the ID of the matching question
|
379 |
-
3. If no good match is found, return -1
|
380 |
-
4. Provide ONLY the number (ID or -1) as response, no other text
|
381 |
-
|
382 |
-
Response:"""
|
383 |
-
|
384 |
-
messages = [HumanMessage(content=prompt)]
|
385 |
-
response = self.llm.invoke(messages)
|
386 |
-
match_id = int(response.content.strip())
|
387 |
-
|
388 |
-
if match_id >= 0 and match_id < len(valid_questions):
|
389 |
-
print(f"Hmm! Match found for query: {query}")
|
390 |
-
return valid_questions[match_id]['answer']
|
391 |
-
|
392 |
-
return None
|
393 |
-
|
394 |
-
except Exception as e:
|
395 |
-
sleep_time = base_sleep * (attempt + 1)
|
396 |
-
if attempt < max_retries - 1:
|
397 |
-
print(f"Answer bank check attempt {attempt + 1} failed. Retrying in {sleep_time} seconds...")
|
398 |
-
time.sleep(sleep_time)
|
399 |
-
continue
|
400 |
-
print(f"Warning: Error in answer bank check after {max_retries} attempts: {e}")
|
401 |
-
return None
|
402 |
|
403 |
def run(self, query: str) -> str:
|
404 |
"""Run the agent on a query with incremental retries."""
|
@@ -407,10 +348,6 @@ Response:"""
|
|
407 |
|
408 |
for attempt in range(max_retries):
|
409 |
try:
|
410 |
-
# First check answer bank
|
411 |
-
cached_answer = self._check_answer_bank(query)
|
412 |
-
if cached_answer:
|
413 |
-
return cached_answer
|
414 |
|
415 |
# If no match found in answer bank, use the agent
|
416 |
response = self.agent.run(query)
|
|
|
339 |
# Load answer bank
|
340 |
self._load_answer_bank()
|
341 |
|
342 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
343 |
|
344 |
def run(self, query: str) -> str:
|
345 |
"""Run the agent on a query with incremental retries."""
|
|
|
348 |
|
349 |
for attempt in range(max_retries):
|
350 |
try:
|
|
|
|
|
|
|
|
|
351 |
|
352 |
# If no match found in answer bank, use the agent
|
353 |
response = self.agent.run(query)
|