mgbam commited on
Commit
aac4d20
·
verified ·
1 Parent(s): d2ff1b7

Update modules/gemini_handler.py

Browse files
Files changed (1) hide show
  1. modules/gemini_handler.py +45 -0
modules/gemini_handler.py CHANGED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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}"