llm.py

from dataclasses import dataclass, field
from openai import OpenAI
import os
from dotenv import load_dotenv
from PIL import Image
from pydantic import BaseModel
from typing import Type, cast
import base64
import io

load_dotenv()

GEMINI_API_KEY =os.getenv("GEMINI_API_KEY")


def encode_image_from_path(image_path:str) -> str:
    """ Encodes an image to Base64

    Args:
        image_path (str): path to image on disk

    Returns:
        str: Base64 encoded image
    """
    with open(file=image_path, mode= "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


def encode_pil_image(image: Image.Image, format: str = "PNG") -> str:
    """
    Encode a PIL Image.Image to a Base64 string.

    Args:
        image (PIL.Image.Image): The PIL image to encode.
        format (str): Image format for encoding (default: "PNG").

    Returns:
        str: Base64-encoded image.
    """
    buffer = io.BytesIO()
    image.save(buffer, format=format)
    buffer.seek(0)
    return base64.b64encode(buffer.read()).decode("utf-8")


@dataclass
class LLM:
    api_key: str | None = GEMINI_API_KEY
    base_url: str | None = "https://generativelanguage.googleapis.com/v1beta/openai/" 
    client:OpenAI = field(init=False)
    model: str = "gemini-3-flash-preview"
    sys_prompt:str | None = None
    user_prompt: str = ""
    image:Image.Image| None = None
    output_model: Type[BaseModel] | None = None

    def __post_init__(self) -> None:
        self.client = OpenAI(api_key=self.api_key, 
                             base_url= self.base_url)


    def get_messages(self) -> list:
        
        messages = []

        # System message (optional)
        if self.sys_prompt:
            messages.append({
                "role": "system",
                "content": self.sys_prompt
            })

        # User message
        if self.image:
            encoded_image = encode_pil_image(self.image)
            user_content = [
                {"type": "text", "text": self.user_prompt},
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{encoded_image}"}
                }
            ]
        else:
            user_content = self.user_prompt

        messages.append({
            "role": "user",
            "content": user_content
        })

        return messages


    def invoke(self) -> BaseModel | str | None:
        try:
            messages = self.get_messages()

            if self.output_model:
                response = self.client.chat.completions.parse( # type: ignore
                    model=self.model,
                    messages=messages,
                    response_format=cast(Type[BaseModel], self.output_model),
                )
                return response.choices[0].message.parsed

            else:
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                )
                return response.choices[0].message.content

        except Exception as e:
            print(f"LLM.invoke() error: {e}")
            return None


sys_prompt_invoice = """You are an invoice extraction assistant. Extract all fields according to the provided schema.
- If a field is missing in the invoice, return null.
- For list fields like "items", return an empty list if there are no items.
- Do not invent or guess values.
- Output must strictly follow the schema."""
user_prompt_invoice = "Extract all invoice information from the attached image."



if __name__ == "__main__":
    
    pass
    

    # print("Loading invoice")
    # if (invoices := load_invoice_images(path="invoices/invoice_2.pdf")):
    #     print("normalizing invoice")
    #     invoice = normalize_invoice_image(img= invoices[0])

    # print("Calling llm")
    
    # # response = LLM(sys_prompt=sys_prompt,
    # #                user_prompt=user_prompt,
    # #                image=invoice,
    # #                output_model=Invoice).invoke()
    
    # response = LLM(user_prompt="What is the capital of France?").invoke()
    # print(response)