llm/llm.py

from pydantic import BaseModel
from google import genai
from google.genai.types import HttpOptions
from dotenv import load_dotenv
import os
import json
import re
from vertexai.generative_models import Part, Image, GenerativeModel
from utils import setup_logger
import vertexai

logger = setup_logger(__name__)

load_dotenv()

project = os.getenv("GOOGLE_CLOUD_PROJECT")
location = os.getenv("GOOGLE_CLOUD_LOCATION")
vertexai.init(project=project, location=location)


class DocumentLLM(BaseModel):

    def call_llm_api(self, prompt, image_path):

        model: GenerativeModel = GenerativeModel(
            model_name="gemini-2.0-flash-001")

        text_part = Part.from_text(prompt)
        image_part = Part.from_image(Image.load_from_file(image_path))

        response = model.generate_content([
            image_part,
            text_part])

        content = response.text

        try:
            content = json.loads(content)
        except Exception as e:
            logger.info(f"Json is being formatted")
            content = re.sub(r"^```json\s*|\s*```$", "",
                             content, flags=re.MULTILINE)

            # Parse JSON
            content = json.loads(content)

        return content