Newer
Older
cortex-hub / ai-hub / app / core / providers / llm / gemini.py
import httpx
import logging
import json
from typing import final
from app.core.providers.base import LLMProvider
from app.config import settings

@final
class GeminiProvider(LLMProvider):
    """Provider for the Google Gemini API."""
    def __init__(self, api_url: str):
        self.url = api_url

    async def generate_response(self, prompt: str) -> str:
        payload = {"contents": [{"parts": [{"text": prompt}]}]}
        headers = {"Content-Type": "application/json"}
        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(self.url, json=payload, headers=headers)
                response.raise_for_status()
                # Await the async `json` method
                data = response.json()
                return data['candidates'][0]['content']['parts'][0]['text']
        except Exception as e:
            logging.error("Gemini Provider Error", exc_info=True)
            raise