omni_agent_v8.py

import os
import time
from google import genai
from google.genai import types, errors
from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type
from fhirpy import SyncFHIRClient
from fhir.resources.patient import Patient
from fhir.resources.observation import Observation

class OMNIOrchestratorV8:
    def __init__(self, api_key, fhir_url="https://local-sovereign-fhir.et/v1"):
        # Unified 2026 SDK Client
        self.client = genai.Client(api_key=api_key)
        self.fhir_client = SyncFHIRClient(fhir_url)
        
        # 2026 Model IDs
        self.pro_model = "gemini-3-pro-preview"
        self.flash_model = "gemini-3-flash-preview"

    @retry(
        retry=retry_if_exception_type(errors.ClientError), 
        wait=wait_exponential(multiplier=1, min=4, max=10),
        stop=stop_after_attempt(2)
    )
    def scribe_audio(self, audio_bytes):
        """
        Ambient Scribe: Converts clinical dialogue to SOAP notes.
        Uses Flash for high-quota and low latency.
        """
        instruction = (
            "You are an Abyssinia Intelligence Scribe. Listen to this clinical encounter. "
            "Extract a SOAP note. Use English for technical fields and Amharic for the patient summary."
        )
        try:
            res = self.client.models.generate_content(
                model=self.flash_model,
                contents=[
                    instruction, 
                    types.Part.from_bytes(data=audio_bytes, mime_type="audio/wav")
                ]
            )
            return res.text
        except Exception as e:
            return f"Scribe Error: {str(e)}"

    def map_genomics(self, vcf_text):
        """
        Genomic Specialist: Analyzes VCF data for drug-gene interactions.
        Implements Model Fallback to handle 429 Resource Exhausted errors.
        """
        instruction = "Identify Ethiopian-specific pharmacogenomic risks from this VCF data."
        
        try:
            # Primary: Try the high-reasoning Pro model
            res = self.client.models.generate_content(
                model=self.pro_model,
                contents=[instruction, vcf_text]
            )
        except errors.ClientError as e:
            if "429" in str(e):
                # Secondary: Fallback to Flash if Pro quota is hit
                res = self.client.models.generate_content(
                    model=self.flash_model,
                    contents=[instruction, vcf_text]
                )
            else:
                raise e
        return res.text

    def synthesize_digital_twin(self, soap_text, genomic_text):
        """
        Final Orchestrator: Merges all data into a FHIR-compliant Digital Twin.
        """
        prompt = f"Merge these clinical notes and genomic risks into a unified Precision Health Report:\n\nNOTES: {soap_text}\nGENOMICS: {genomic_text}"
        
        # V8 default uses 'Thinking' mode for synthesis
        res = self.client.models.generate_content(
            model=self.flash_model,
            contents=prompt,
            config=types.GenerateContentConfig(
                thinking_config=types.ThinkingConfig(include_thoughts=True)
            )
        )
        return res.text