Update app.py

app.py

@@ -2,108 +2,49 @@ import os
 import threading
 import time
 import subprocess
-import gradio as gr
-from ollama import Client
+from flask import Flask, request, Response, stream_with_context
+import requests
 
-# 1. Start Ollama Server (Standard Setup)
-def ollama_service_thread():
-    subprocess.run("ollama serve", shell=True)
+# 1. Start Ollama Server
+def run_ollama():
+    subprocess.Popen(["ollama", "serve"])
 
-threading.Thread(target=ollama_service_thread, daemon=True).start()
+threading.Thread(target=run_ollama, daemon=True).start()
 
-# Wait for server to wake up
-print("Waiting for Ollama to start...")
+# 2. Wait & Pull Model
+print("Waiting for Ollama...")
 time.sleep(10)
-
-# 2. Pull the Model
 MODEL = "llama3.2:3b"
 subprocess.run(f"ollama pull {MODEL}", shell=True)
+print("Model Ready!")
 
-# 3. Connect Client
-client = Client(host='http://127.0.0.1:11434')
-
-# --- FUNCTION A: FOR THE WEB UI (Streaming) ---
-# This uses 'yield' to type out words one by one.
-def chat_ui(message, history, temperature, max_new_tokens, top_p, top_k, penalty):
-    conversation = []
-    for prompt, answer in history:
-        conversation.extend([
-            {"role": "user", "content": prompt},
-            {"role": "assistant", "content": answer},
-        ])
-    conversation.append({"role": "user", "content": message})
-
-    response = client.chat(
-        model=MODEL,
-        messages=conversation,
-        stream=True,  
-        options={
-            "num_predict": max_new_tokens,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "repeat_penalty": penalty,
-        }
-    )
-    
-    partial_message = ""
-    for chunk in response:
-        if 'message' in chunk and 'content' in chunk['message']:
-            partial_message += chunk['message']['content']
-            yield partial_message
+# 3. Create a Proxy App (Flask)
+app = Flask(__name__)
 
-# --- FUNCTION B: FOR n8n (Blocking / No Stream) ---
-# This uses 'return' to send the full text at once.
-# This prevents the 404/Timeout errors in n8n.
-def chat_n8n(message, temperature, max_new_tokens, top_p, top_k, penalty):
-    conversation = [{"role": "user", "content": message}]
+# This forwards everything from /v1/* to local Ollama
+@app.route('/v1/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE'])
+def proxy(path):
+    url = f"http://127.0.0.1:11434/v1/{path}"
     
-    response = client.chat(
-        model=MODEL,
-        messages=conversation,
-        stream=False, # <--- CRITICAL: No streaming for API
-        options={
-            "num_predict": max_new_tokens,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "repeat_penalty": penalty,
-        }
+    # Forward the request to Ollama
+    resp = requests.request(
+        method=request.method,
+        url=url,
+        headers={key: value for (key, value) in request.headers if key != 'Host'},
+        data=request.get_data(),
+        cookies=request.cookies,
+        allow_redirects=False,
+        stream=True 
     )
-    return response['message']['content']
 
-# --- THE INTERFACE ---
-with gr.Blocks() as demo:
-    gr.HTML("<h1>🦙 Llama 3.2 (Dual Mode: UI + API)</h1>")
-    
-    # 1. The Visible Chatbot (Uses Streaming Function A)
-    gr.ChatInterface(
-        fn=chat_ui,
-        additional_inputs=[
-            gr.Slider(0, 1, 0.7, label="Temperature"),
-            gr.Slider(128, 4096, 1024, label="Max Tokens"),
-            gr.Slider(0.0, 1.0, 0.9, label="Top P"),
-            gr.Slider(1, 100, 40, label="Top K"),
-            gr.Slider(0.0, 2.0, 1.1, label="Penalty"),
-        ]
-    )
-    
-    # 2. The Hidden API Route (Uses Blocking Function B)
-    # We create a button that is invisible but defines the API name "n8n"
-    api_btn = gr.Button("API Call", visible=False)
-    api_btn.click(
-        fn=chat_n8n,
-        inputs=[
-            gr.Textbox(),          # Message
-            gr.Number(value=0.7),  # Temp
-            gr.Number(value=1024), # Tokens
-            gr.Number(value=0.9),  # Top P
-            gr.Number(value=40),   # Top K
-            gr.Number(value=1.1)   # Penalty
-        ],
-        outputs=gr.Textbox(),
-        api_name="n8n"  # <--- THIS IS THE KEY. It creates /api/n8n
-    )
+    # Stream the response back to n8n
+    def generate():
+        for chunk in resp.iter_content(chunk_size=4096):
+            if chunk:
+                yield chunk
+
+    return Response(generate(), status=resp.status_code, headers=dict(resp.headers))
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+if __name__ == '__main__':
+    # Run on port 7860 (Hugging Face default)
+    app.run(host='0.0.0.0', port=7860)