import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
from gpt4all import GPT4All

st.set_page_config(page_title="AutoGPT with Streamlit", layout="wide")

st.title("🤖 AutoGPT Agent (Docker + Streamlit)")
st.write("Powered by **DeepSeek + GPT4All**")

# Sidebar setup
st.sidebar.header("Settings")
model_choice = st.sidebar.selectbox("Choose Model", ["DeepSeek", "GPT4All"])
goal = st.text_area("Enter your AI Goal:", placeholder="e.g., Research 5 AI trends and summarize")

if st.button("Run Agent"):
    if model_choice == "DeepSeek":
        model_name = "deepseek-ai/deepseek-coder-1.3b-base"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name)
        inputs = tokenizer(goal, return_tensors="pt")
        output = model.generate(**inputs, max_new_tokens=200)
        result = tokenizer.decode(output[0], skip_special_tokens=True)

    elif model_choice == "GPT4All":
        gpt4all = GPT4All("gpt4all-lora-quantized.bin")
        with gpt4all.chat_session():
            result = gpt4all.generate(goal, max_tokens=200)

    else:
        result = "No model selected."

    st.subheader("Agent Output")
    st.write(result)