app/api/gemini/main.py - master - durovcode/backend

app/api/gemini/main.py

47 строк1 KB

quyxishi

implemented wishes feature for ai/query method

22 сен 2024, 06:21

22 сен 2024, 06:21bffb994

import json
from typing import Any, Dict, Optional

import google.generativeai as genai
from google.generativeai.types.generation_types import (
    GenerateContentResponse,
)

from app.env import (
    GEMINI_API_KEY,
    GEMINI_MODEL,
    GEMINI_SYSTEM_PROMPT,
    # GIGA_CLIENT_ID,
    # GIGA_CLIENT_SECRET,
    # GIGA_MODEL_NAME,
    # GIGA_MODEL_SCOPE,
    # GIGA_SYSTEM_PROMPT,
)


SYSTEM_PROMPT: Dict[str, str] = {'role': 'system', 'content': GEMINI_SYSTEM_PROMPT}
USER_PROMPT_POSTFIX: str = ''


class GeminiInterface:
    def __init__(self) -> None:
        genai.configure(transport='rest', api_key=GEMINI_API_KEY)

        self.model = genai.GenerativeModel(
            model_name=GEMINI_MODEL,
            generation_config=genai.GenerationConfig(
                response_mime_type='application/json'
            ),
            system_instruction=GEMINI_SYSTEM_PROMPT,
        )

    async def query(
        self, message: str, commentary: Optional[str] = None
    ) -> Dict[str, Any]:
        gemini_response: GenerateContentResponse = self.model.generate_content(
            contents='\n\n'.join(
                ([commentary.upper()] if isinstance(commentary, str) else [])
                + [message]
            )
        )

        return json.loads(gemini_response.text)