Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import requests
 import os
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 
 api_keys = []
 last_api_key_used = 0
@@ -33,9 +35,7 @@ def get_api_key():
     global last_api_key_used
 
     if num_api_keys == 0:
-        raise Exception(
-            "No API keys found. Format: OLLAMA_<number>"
-        )
+        raise Exception("No API keys found. Format: OLLAMA_<number>")
 
     last_api_key_used += 1
     if last_api_key_used >= num_api_keys:
@@ -63,20 +63,63 @@ def perform_search(query):
     response = requests.post(url, headers=headers, json=payload)
     response.raise_for_status()
 
-    results = response.json()
-    return results
+    return response.json()
 
-    # return "\n".join(
-    #     f"{r['title']}: {r['url']}"
-    #     for r in results
-    # )
 
+# -----------------------
+# FastAPI
+# -----------------------
+
+app = FastAPI(title="Ollama Web Search API")
+
+
+class SearchRequest(BaseModel):
+    query: str
+    max_results: int = 1
+
+
+@app.post("/api/search")
+def api_search(request: SearchRequest):
+
+    try:
+        api_key = get_api_key()
+
+        url = "https://ollama.com/api/web_search"
+
+        headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json"
+        }
+
+        payload = {
+            "query": request.query,
+            "max_results": request.max_results
+        }
+
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+
+        return response.json()
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}
+
+
+# -----------------------
+# Gradio UI
+# -----------------------
 
 demo = gr.Interface(
     fn=perform_search,
     inputs=gr.Textbox(label="Search Query"),
-    outputs=gr.Textbox(label="Results"),
+    outputs=gr.JSON(label="Results"),
     title="Ollama Web Search API"
 )
 
-demo.launch()
+# Mount Gradio at root
+app = gr.mount_gradio_app(app, demo, path="/")