Reservationer/assets/gemini_search.py

import json
import google
import google.genai

class gemini_search:
    def __init__(self, apikey: str) -> None:
        self.__apikey = apikey

    @property
    def __client(self) -> google.genai.Client:
        return google.genai.Client(api_key = self.__apikey)

    @property
    def __model(self) -> str:
        return "gemini-2.5-flash-preview-05-20"

    def __user(self, barcode: str) -> google.genai.types.Content:
        return google.genai.types.Content(
            role = "user",
            parts = [
                google.genai.types.Part.from_text(
                    text = barcode
                )
            ]
        )

    @property
    def __instructions(self) -> str:
        return (
            "Find a product with the given EAN code. Write a short " \
            + "description of this product (about 200 characters) in " \
            + "the native language of the country where the product " \
            + "comes from. Write only the truth."
        )

    @property
    def __research(self) -> google.genai.types.GenerateContentConfig:
        return google.genai.types.GenerateContentConfig(
            response_mime_type = "text/plain",
            thinking_config = google.genai.types.ThinkingConfig(
                thinking_budget = 0,
            ),
            tools = [
                google.genai.types.Tool(
                    google_search = google.genai.types.GoogleSearch()
                )
            ],
            system_instruction = [
                google.genai.types.Part.from_text(
                    text = self.__instructions
                )
            ]
        )

    @property
    def __process(self) -> google.genai.types.GenerateContentConfig:
        return google.genai.types.GenerateContentConfig(
            thinking_config = google.genai.types.ThinkingConfig(
                thinking_budget = 0,
            ),
            response_mime_type = "application/json",
            response_schema = google.genai.types.Schema(
                type = google.genai.types.Type.OBJECT,
                required = ["title", "description", "author"],
                properties = {
                    "title": google.genai.types.Schema(
                        type = google.genai.types.Type.STRING,
                    ),
                    "description": google.genai.types.Schema(
                        type = google.genai.types.Type.STRING,
                    ),
                    "author": google.genai.types.Schema(
                        type = google.genai.types.Type.STRING,
                    ),
                },
            ),
            system_instruction = [
                google.genai.types.Part.from_text(
                    text = self.__instructions
                )
            ],
        )

    def __request(
        self, 
        barcode: str,
        conf: google.genai.types.GenerateContentConfig,
        addon: str | None = None
    ) -> str:
        content = ""
        contents = [
            self.__user(barcode)
        ]

        if addon is not None:
            contents.append(
                google.genai.types.Content(
                    role = "model",
                    parts = [
                        google.genai.types.Part.from_text(
                            text = addon
                        )
                    ]
                )
            )

            contents.append(
                self.__user(barcode)
            )

        for count in self.__client.models.generate_content_stream(
            model = self.__model,
            contents = contents,
            config = conf
        ):
            content = content + count.text

        return content

    def search(self, barcode: str) -> dict:
        research = self.__request(barcode, self.__research)
        model = self.__request(barcode, self.__process, research)

        return json.loads(model)