app.py

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


route="""	
how to handle special case "zugverbindung".

Regeln:
Wenn eine Zugverbindung von {Startort} nach {Zielort} angefragt wird, return json object with Startort and Zielort.
always follow json scheme below.

Wichtig: Gib absolut keinen Text vor oder nach dem JSON aus (keine Erklärungen, kein "Hier ist das Ergebnis").

{
  "start_loc": "fill in Startort here",
  "dest_loc": "fill in Zielort here"
}

"""



def clean_json_string(json_str):
    """
    Removes any comments or prefixes before the actual JSON content.
    """
    # Find the first occurrence of '{'
    json_start = json_str.find('{')
    if json_start == -1:
        # If no '{' is found, try with '[' for arrays
        json_start = json_str.find('[')
        if json_start == -1:
            return json_str  # Return original if no JSON markers found
    
    # Extract everything from the first JSON marker
    cleaned_str = json_str[json_start:]
    return cleaned_str
    # Verify it's valid JSON
    try:
        json.loads(cleaned_str)
        return cleaned_str
    except json.JSONDecodeError:
        return json_str  # Return original if cleaning results in invalid JSON

def generate(input_text):
    try:
        client = genai.Client(
            api_key=os.environ.get("GEMINI_API_KEY"),
        )
    except Exception as e:
        return f"Error initializing client: {e}.  Make sure GEMINI_API_KEY is set."

    model = "gemini-flash-latest"
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text=f"{route}  {input_text}"),
            ],
        ),
    ]
    tools = [
        types.Tool(google_search=types.GoogleSearch()),
    ]
    generate_content_config = types.GenerateContentConfig(
        temperature=0.4,
        thinking_config = types.ThinkingConfig(
            thinking_budget=0,
        ),
        tools=tools,
        response_mime_type="text/plain",
    )


    response_text = ""
    try:
      for chunk in client.models.generate_content_stream(
          model=model,
          contents=contents,
          config=generate_content_config,
      ):
          response_text += chunk.text
    except Exception as e:
        return f"Error during generation: {e}"
    data = clean_json_string(response_text)
    data = data[:-1]
    return response_text, ""
   

if __name__ == '__main__':

    with gr.Blocks() as demo:
        title=gr.Markdown("# Gemini 2.0 Flash + Websearch")
        output_textbox = gr.Markdown()
        input_textbox = gr.Textbox(lines=3, label="", placeholder="Enter message here...")
        submit_button = gr.Button("send")
        submit_button.click(fn=generate,inputs=input_textbox,outputs=[output_textbox, input_textbox])
    demo.launch(show_error=True)