app/llm/gemini_client.py

gistfile1.txt

from __future__ import annotations

from vertexai.preview.generative_models import GenerativeModel, Part

from app.common.config import Settings
from app.llm.vertex_init import init_vertex


def generate(
    settings: Settings,
    parts: list[Part] | Part,
    *,
    temperature: float = 0.3,
    max_output_tokens: int = 1024,
) -> str:
    init_vertex(settings)
    model = GenerativeModel(settings.gemini_model)
    resp = model.generate_content(
        parts,
        generation_config={
            "temperature": temperature,
            "max_output_tokens": max_output_tokens,
        },
    )
    return getattr(resp, "text", "") or ""