Update main.py

main.py

@@ -1,33 +1,22 @@
-from fastapi import FastAPI
+import gradio as gr
 import requests
-import os
-from fastapi.middleware.cors import CORSMiddleware
 
-app = FastAPI()
+FASTAPI_URL = "http://localhost:7860/ask_ollama/"
 
-# Allow frontend to connect
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-OLLAMA_BASE_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
+def ask_ollama(question):
+    response = requests.get(f"{FASTAPI_URL}?question={question}")
+    if response.status_code == 200:
+        return response.json().get("response", "No response from Ollama.")
+    return "Error: Could not connect to API."
 
-@app.get("/hi")
-def read_root():
-    return {"message": "FastAPI with Ollama is running!"}
+# Create Gradio interface
+iface = gr.Interface(
+    fn=ask_ollama,
+    inputs=gr.Textbox(placeholder="Ask a question..."),
+    outputs="text",
+    title="Ollama Chat",
+    description="Ask anything and get responses from the Ollama model."
+)
 
-@app.get("/ask_ollama/")
-def ask_ollama(question: str):
-    payload = {
-        "model": "mistral",
-        "messages": [
-            {"role": "system", "content": "You are an AI assistant."},
-            {"role": "user", "content": question}
-        ]
-    }
-    response = requests.post(f"{OLLAMA_BASE_URL}/api/generate", json=payload)
-    return response.json()
+# Run on port 7861 (so FastAPI and Gradio can run together)
+iface.launch(server_name="0.0.0.0", server_port=7861, share=True)