mgbam commited on
Commit
aa3b805
·
verified ·
1 Parent(s): e008cd7

Update modules/gemini_handler.py

Browse files
Files changed (1) hide show
  1. modules/gemini_handler.py +64 -21
modules/gemini_handler.py CHANGED
@@ -1,45 +1,88 @@
1
  # modules/gemini_handler.py
2
  """
3
  Dedicated module for all interactions with the Google Gemini API.
 
 
4
  """
5
  import google.generativeai as genai
 
6
  from .config import GEMINI_API_KEY
7
 
8
- # Configure the API key
9
  if not GEMINI_API_KEY:
10
- raise ValueError("GEMINI_API_KEY not found. Please set it in your environment.")
11
  genai.configure(api_key=GEMINI_API_KEY)
12
 
13
- # Set up the model with safety settings to be less restrictive
14
- # This is important for medical contexts, but use with caution.
15
  generation_config = {
16
- "temperature": 0.2,
17
- "top_p": 1,
18
- "top_k": 1,
19
- "max_output_tokens": 4096,
20
  }
21
 
 
 
22
  safety_settings = [
23
  {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
24
  {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
25
- {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
26
  {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
27
  ]
28
 
29
- model = genai.GenerativeModel(
30
- model_name="gemini-pro",
31
- generation_config=generation_config,
32
- safety_settings=safety_settings
33
- )
 
 
 
34
 
35
- async def generate_gemini_response(prompt: str) -> str:
36
- """Generic function to call the Gemini API and get a response."""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
37
  try:
38
- response = await model.generate_content_async(prompt)
39
- # Handle cases where the response might be blocked
40
  if not response.parts:
41
- return "The AI response was blocked due to safety settings. Please rephrase your query."
42
  return response.text
43
  except Exception as e:
44
- print(f"An error occurred with the Gemini API: {e}")
45
- return f"Error: Could not get a response from the AI model. Details: {e}"
 
1
  # modules/gemini_handler.py
2
  """
3
  Dedicated module for all interactions with the Google Gemini API.
4
+ It initializes and configures the models, and provides clean, async functions
5
+ for both text and multi-modal (vision) generation, including robust error handling.
6
  """
7
  import google.generativeai as genai
8
+ from PIL import Image
9
  from .config import GEMINI_API_KEY
10
 
11
+ # --- Configuration ---
12
  if not GEMINI_API_KEY:
13
+ raise ValueError("FATAL: GEMINI_API_KEY not found. Please set it in your environment/secrets.")
14
  genai.configure(api_key=GEMINI_API_KEY)
15
 
16
+ # Configuration for predictable, factual responses
 
17
  generation_config = {
18
+ "temperature": 0.1,
19
+ "top_p": 0.95,
20
+ "top_k": 40,
21
+ "max_output_tokens": 8192, # Increased for detailed reports
22
  }
23
 
24
+ # Safety settings are relaxed slightly for medical context, but still block high-risk content.
25
+ # This prevents the model from refusing to discuss health topics.
26
  safety_settings = [
27
  {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
28
  {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
29
+ {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_ONLY_HIGH"},
30
  {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
31
  ]
32
 
33
+ # --- Model Initialization ---
34
+ try:
35
+ # Model for text-based tasks (synthesis, extraction)
36
+ text_model = genai.GenerativeModel(
37
+ model_name="gemini-1.5-flash-latest", # Use a fast and capable model
38
+ generation_config=generation_config,
39
+ safety_settings=safety_settings
40
+ )
41
 
42
+ # Model for multi-modal tasks (analyzing medical images)
43
+ vision_model = genai.GenerativeModel(
44
+ model_name="gemini-1.5-flash-latest", # Vision models are also great at text
45
+ generation_config=generation_config,
46
+ safety_settings=safety_settings
47
+ )
48
+ except Exception as e:
49
+ raise RuntimeError(f"Failed to initialize Gemini models: {e}")
50
+
51
+ async def generate_text_response(prompt: str) -> str:
52
+ """
53
+ Calls the Gemini text model with a given prompt.
54
+
55
+ Args:
56
+ prompt (str): The detailed prompt for the AI.
57
+
58
+ Returns:
59
+ str: The AI's generated text response, or an error message.
60
+ """
61
+ try:
62
+ response = await text_model.generate_content_async(prompt)
63
+ if not response.parts:
64
+ return "Error: The AI response was blocked, possibly due to safety filters or lack of content. Please rephrase."
65
+ return response.text
66
+ except Exception as e:
67
+ print(f"Gemini API Error (Text): {e}")
68
+ return f"An error occurred while communicating with the AI model. Details: {e}"
69
+
70
+ async def analyze_image_with_text(prompt: str, image: Image.Image) -> str:
71
+ """
72
+ Calls the Gemini vision model with a text prompt and an image.
73
+
74
+ Args:
75
+ prompt (str): The text prompt to guide the image analysis.
76
+ image (Image.Image): The PIL image to be analyzed.
77
+
78
+ Returns:
79
+ str: The AI's generated text analysis, or an error message.
80
+ """
81
  try:
82
+ response = await vision_model.generate_content_async([prompt, image])
 
83
  if not response.parts:
84
+ return "Error: The AI response for the image was blocked. The image may violate safety policies."
85
  return response.text
86
  except Exception as e:
87
+ print(f"Gemini API Error (Vision): {e}")
88
+ return f"An error occurred while analyzing the image with the AI. Details: {e}"