GEAK-agent_debug/src/models/Gemini.py

from typing import List
from tenacity import retry, stop_after_attempt, wait_random_exponential
import requests
from models.Base import BaseModel


class GeminiModel(BaseModel):
    def __init__(self, 
                 model_id="gemini-2.5-pro-preview-05-06", 
                 api_key=None):
        assert api_key is not None, "no api key is provided."
        self.model_id = model_id
        
        self.SERVER = "https://llm-api.amd.com/vertex/gemini"
        self.HEADERS = {"Ocp-Apim-Subscription-Key": api_key}

    
    @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(5))
    def generate(self, 
                 messages: List, 
                 temperature=1.0, 
                 presence_penalty=0, 
                 frequency_penalty=0, 
                 max_tokens=30000) -> str:
        body = {
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature":temperature,
            "top_P": 0.95,
            "presence_Penalty": presence_penalty,
            "frequency_Penalty": frequency_penalty,
        }
        response_gemine = requests.post(url=f"{self.SERVER}/{self.model_id}/chat", 
                            json=body,
                            headers=self.HEADERS)
        assert response_gemine.status_code == 200
        code_chat_completion_result = response_gemine.json()
        
        return code_chat_completion_result['candidates'][0]['content']['parts'][0]['text']