Spaces:
Running
Running
Update modules/gemini_handler.py
Browse files- 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 |
-
#
|
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 |
-
#
|
14 |
-
# This is important for medical contexts, but use with caution.
|
15 |
generation_config = {
|
16 |
-
"temperature": 0.
|
17 |
-
"top_p":
|
18 |
-
"top_k":
|
19 |
-
"max_output_tokens":
|
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": "
|
26 |
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
|
27 |
]
|
28 |
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
|
|
|
|
|
|
34 |
|
35 |
-
|
36 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
37 |
try:
|
38 |
-
response = await
|
39 |
-
# Handle cases where the response might be blocked
|
40 |
if not response.parts:
|
41 |
-
return "The AI response was blocked
|
42 |
return response.text
|
43 |
except Exception as e:
|
44 |
-
print(f"
|
45 |
-
return f"
|
|
|
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}"
|