import json
import base64
import os
from google import genai
from google.genai import types
import gradio as gr

def generate(prompt):
    client = genai.Client(
        api_key=os.environ.get("GEMINI_API_KEY"),
    )

    model = "gemini-2.0-pro-exp-02-05"
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text=f"""prompt: {prompt}              
ordne den prompt in einer der folgenden kategorien

sitzungen
z.b. sitzungen von gemeinderäten,stadträdten,marktgemeinderäten, ausschüssen etc.

versammlungen
z.b. mitglieder- general-, jahreshauptversammlungen

ereignisse
z.b. sperrung einer straße

andere
wenn keiner der anderen kategorien passt
"""),
            ],
        ),
        types.Content(
            role="model",
            parts=[
                types.Part.from_text(text="""{
  \"kategorie\": \"sitzungen\"
}"""),
            ],
        ),
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text=prompt),
            ],
        ),
    ]

    generate_content_config = types.GenerateContentConfig(
        temperature=0.4,
        top_p=0.95,
        top_k=40,
        max_output_tokens=8192,
        response_mime_type="application/json",
        response_schema=genai.types.Schema(
            type=genai.types.Type.OBJECT,
            required=["kategorie"],
            properties={
                "kategorie": genai.types.Schema(
                    type=genai.types.Type.STRING,
                ),
            },
        ),
    )

    result = ""
    for chunk in client.models.generate_content_stream(
        model=model,
        contents=contents,
        config=generate_content_config,
    ):
        result += chunk.text
    json_data=json.loads(result)
    kategorie =json_data["kategorie"]
    return kategorie

iface = gr.Interface(
    fn=generate,
    inputs="text",
    outputs="text",
    title="Kategorie-Zuordnung",
    description="Ordnet einen gegebenen Prompt in eine der vorgegebenen Kategorien ein."
)

iface.launch()