103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
import asyncio
|
|
import json
|
|
import base64
|
|
from app.config import GEMINI_API_KEY, GEMINI_MODEL
|
|
|
|
|
|
class GeminiClient:
|
|
|
|
def __init__(self):
|
|
self.api_key = GEMINI_API_KEY
|
|
self.model = GEMINI_MODEL or "gemini-1.5-flash"
|
|
|
|
async def extract_from_document(
|
|
self,
|
|
file_bytes: bytes,
|
|
mime_type: str,
|
|
patient_context: dict,
|
|
) -> dict:
|
|
if not self.api_key or self.api_key in ["your_gemini_key", "your_gemini_api_key", ""]:
|
|
return {
|
|
"document_type": "medical_report",
|
|
"extracted_text": "Sample HbA1c 7.2%, Blood Pressure 130/85 mmHg",
|
|
"findings": [
|
|
{"test": "HbA1c", "value": "7.2", "unit": "%", "interpretation": "Elevated"},
|
|
{"test": "Blood Pressure", "value": "130/85", "unit": "mmHg", "interpretation": "Stage 1"},
|
|
],
|
|
"medications_found": ["Metformin 500mg"],
|
|
"allergies_found": [],
|
|
"summary": "Report indicates elevated HbA1c and prehypertension.",
|
|
}
|
|
|
|
try:
|
|
from google import genai
|
|
from google.genai import types
|
|
|
|
client = genai.Client(api_key=self.api_key)
|
|
|
|
prompt = f"""
|
|
You are a medical document extraction AI.
|
|
Analyze this medical report / prescription / lab result for patient {patient_context.get('name', 'Patient')}.
|
|
Extract all lab tests, numeric values, units, diagnosed conditions, medications, and allergies into structured JSON.
|
|
|
|
Return ONLY JSON:
|
|
{{
|
|
"document_type": "lab_report | prescription | discharge_summary | other",
|
|
"extracted_text": "brief summary text",
|
|
"findings": [
|
|
{{
|
|
"test": "test name",
|
|
"value": "numeric or text value",
|
|
"unit": "unit or null",
|
|
"interpretation": "Normal | High | Low | etc."
|
|
}}
|
|
],
|
|
"medications_found": ["medication 1", "medication 2"],
|
|
"allergies_found": ["allergy 1"],
|
|
"summary": "short clinical summary of the document"
|
|
}}
|
|
"""
|
|
part = types.Part.from_bytes(data=file_bytes, mime_type=mime_type)
|
|
response = await asyncio.wait_for(
|
|
client.aio.models.generate_content(
|
|
model=self.model,
|
|
contents=[part, prompt]
|
|
),
|
|
timeout=4.0
|
|
)
|
|
cleaned_text = response.text.replace("```json", "").replace("```", "").strip()
|
|
return json.loads(cleaned_text)
|
|
except Exception as e:
|
|
print(f"[Gemini Extraction Warning]: {e}")
|
|
return {
|
|
"document_type": "unknown",
|
|
"findings": [],
|
|
"medications_found": [],
|
|
"allergies_found": [],
|
|
"summary": f"Could not extract document details: {e}",
|
|
}
|
|
|
|
async def generate_response(
|
|
self,
|
|
system_prompt: str,
|
|
user_message: str,
|
|
) -> str:
|
|
if not self.api_key or self.api_key in ["your_gemini_key", "your_gemini_api_key", ""]:
|
|
raise ValueError("Gemini API key not configured")
|
|
|
|
from google import genai
|
|
from google.genai import types
|
|
|
|
client = genai.Client(api_key=self.api_key)
|
|
full_prompt = f"{system_prompt}\n\nPATIENT SAID: {user_message}"
|
|
config = types.GenerateContentConfig(response_mime_type="application/json")
|
|
response = await asyncio.wait_for(
|
|
client.aio.models.generate_content(
|
|
model=self.model,
|
|
contents=full_prompt,
|
|
config=config,
|
|
),
|
|
timeout=3.5
|
|
)
|
|
return response.text
|