import streamlit as st
from langchain.callbacks import StreamlitCallbackHandler
from langchain.agents import AgentType, Tool
from typing import Literal
from langchain.agents import initialize_agent, load_tools, AgentType
from langchain.chains.base import Chain
from langchain.chat_models import ChatOpenAI
from langchain_experimental.plan_and_execute import (
load_chat_planner, load_agent_executor, PlanAndExecute
)
from langchain.utilities import PythonREPL
ReasoningStrategies = Literal["zero-shot-react", "plan-and-solve"]

def load_agent(
    tool_names: list[str],
    strategy: ReasoningStrategies = "zero-shot-react",
    custom_tools: list = [],
    openai_key:str = ''
    ) -> Chain:

    llm = ChatOpenAI(temperature=0.5, streaming=True,openai_api_key = openai_key ,model = 'gpt-3.5-turbo-0125')
    tools = load_tools(
    tool_names=tool_names,
    llm = llm)

    tools = tools + custom_tools

    if strategy == "plan-and-solve":
        planner = load_chat_planner(llm)
        executor = load_agent_executor(llm, tools, verbose=True)
        return PlanAndExecute(planner=planner, executor=executor,
        verbose=True)
    
    return initialize_agent(tools=tools, llm=llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,verbose=True)

python_repl = PythonREPL()
repl_tool  = Tool(
        "python_repl",
        func = python_repl.run,
        description="A Python command shell. Do not add markdown text. Do not use recursion and print every step"
    )
strategy = st.radio("Reasoning strategy",("plan-and-solve", "zero-shot-react"))

tool_names = st.multiselect('Which tools do you want to use?',["google-search", "ddg-search", "arxiv","wikipedia", "python_repl", "pal-math", "llm-math"],
["ddg-search", "wikipedia"])

def run_agent(prompt,st_callback, openai_key):
    agent_chain = load_agent(tool_names=tool_names, strategy=strategy,custom_tools = [repl_tool], openai_key = openai_key )
    agent_chain.run(prompt, callbacks=[st_callback])

st_callback = StreamlitCallbackHandler(st.container())
openai_key = st.text_input(label="Your OpenAI API key", type="password", placeholder="Enter API Key Here")
if prompt := st.chat_input():
    st.chat_message("user").write(prompt)
    
    with st.chat_message("assistant"):
        st_callback = StreamlitCallbackHandler(st.container())
        response = run_agent(prompt, st_callback, openai_key)
        st.write(response)