import gradio as gr
import requests
import subprocess
import time
import sys

# Start the Ollama server and capture its URL output
process = subprocess.Popen(['bash', 'run_ollama.sh'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = process.communicate()

# Check for errors or retrieve the URL from stdout
if process.returncode != 0:
    print(f"Error starting Ollama: {stderr.decode()}")
    sys.exit(1)

ollama_url = stdout.decode().strip().split('\n')[-1].split()[-1]  # Extract the URL

console_output = []

def call_ollama_api(prompt):
    try:
        response = requests.post(
            f'{ollama_url}/generate',
            json={"prompt": prompt}
        )
        response.raise_for_status()  # Raise an exception for HTTP errors
        result = response.json()
        output = result.get('text', '')
        console_output.append(f"Prompt: {prompt}\nResponse: {output}")
        return output
    except requests.exceptions.RequestException as e:
        error_message = f"Error: {e}"
        console_output.append(error_message)
        return error_message

def chat(prompt):
    output = call_ollama_api(prompt)
    return output

def show_console_output():
    return "\n\n".join(console_output)

iface = gr.Interface(
    fn=chat,
    inputs="text",
    outputs="text",
    title="Ollama Chat",
    description="Chat with Ollama API and see the responses."
)

console_iface = gr.Interface(
    fn=show_console_output,
    inputs=None,
    outputs="text",
    title="Console Output",
    description="See the console outputs here."
)

# Launch a tabbed interface
gr.TabbedInterface([iface, console_iface], ["Chat", "Console"]).launch()