Update app.py

app.py

@@ -1,166 +1,66 @@
-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".
-Wichtig: Dies Regeln gelten nur wenn eine zugverbindung angefragt wird, else answer prompt
-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"{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)
-
-
-
-
-
-"""
 import os
 import asyncio
-import gradio as gr
 from google import genai
 from google.genai import types
-
 from mcp import ClientSession
-from mcp.client.sse import sse_client  # Spezifischer Transport für Gradio/HF
-
-async def generate(input_text):
-    # WICHTIG: Gradio MCP Server benötigen oft das Suffix /sse
-    mcp_url = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
-    
-    try:
-        client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
-        
-        # SSE Transport nutzen, um den 'text/html' Fehler zu vermeiden
-        async with sse_client(url=mcp_url) as (read_stream, write_stream):
-            async with ClientSession(read_stream, write_stream) as mcp_session:
-                await mcp_session.initialize()
+from mcp.client.sse import sse_client
+
+# Konfiguration des MCP-Servers (Beispiel DB-Timetable auf Hugging Face)
+MCP_URL = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
+
+async def fetch_train_connections(prompt: str):
+    client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
+    model_id = "gemini-2.0-flash"
+
+    # 1. Verbindung zum MCP-Server über SSE herstellen [3]
+    async with sse_client(url=MCP_URL) as (read_stream, write_stream):
+        async with ClientSession(read_stream, write_stream) as session:
+            # Initialisierung der MCP-Sitzung [4], [5]
+            await session.initialize()
+
+            # 2. MCP-Tools abrufen und für Gemini konvertieren [4], [2]
+            mcp_tools = await session.list_tools()
+            tools = types.Tool(function_declarations=[
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters": tool.inputSchema,
+                }
+                for tool in mcp_tools.tools
+            ])
+
+            # 3. Anfrage an das Modell senden [6]
+            contents = [types.Content(role="user", parts=[types.Part(text=prompt)])]
+            response = await client.aio.models.generate_content(
+                model=model_id,
+                contents=contents,
+                config=types.GenerateContentConfig(
+                    tools=[tools],
+                    temperature=0.4
+                )
+            )
 
-                model_id = "gemini-2.0-flash" 
+            # 4. Tool-Calling Loop: Falls das Modell eine Zugverbindung sucht [7], [8]
+            if response.candidates.content.parts.function_call:
+                fc = response.candidates.content.parts.function_call
+                
+                # Tool auf dem MCP-Server ausführen
+                tool_result = await session.call_tool(fc.name, fc.args)
                 
-                generate_content_config = types.GenerateContentConfig(
-                    temperature=0.4,
-                    tools=[
-                        types.Tool(google_search=types.GoogleSearch()),
-                        mcp_session # Reicht die Tools des DB-Servers an Gemini durch
-                    ],
+                # Ergebnis an das Modell zurückgeben für die finale Antwort [8], [9]
+                tool_response_part = types.Part.from_function_response(
+                    name=fc.name, 
+                    response={"result": tool_result.content.text}
                 )
-
-                response_text = ""
-                async for chunk in client.aio.models.generate_content_stream(
-                    model=model_id,
-                    contents=input_text,
-                    config=generate_content_config,
-                ):
-                    if chunk.text:
-                        response_text += chunk.text
                 
-                return response_text, ""
+                contents.append(response.candidates.content)
+                contents.append(types.Content(role="user", parts=[tool_response_part]))
 
-    except Exception as e:
-        return f"Verbindung zum DB-Fahrplan fehlgeschlagen: {str(e)}", ""
-
-def gradio_wrapper(input_text):
-    return asyncio.run(generate(input_text))
+                final_response = await client.aio.models.generate_content(
+                    model=model_id,
+                    contents=contents,
+                    config=types.GenerateContentConfig(tools=[tools])
+                )
+                return final_response.text
 
-if __name__ == '__main__':
-    with gr.Blocks() as demo:
-        gr.Markdown("# Gemini Flash + DB Timetable")
-        input_tx = gr.Textbox(label="Anfrage", placeholder="Wann fährt der nächste Zug von Berlin nach Hamburg?")
-        btn = gr.Button("Senden")
-        output_md = gr.Markdown()
-        
-        btn.click(fn=gradio_wrapper, inputs=input_tx, outputs=[output_md, input_tx])
-        
-    demo.launch()
-    """
+            return response.text