import streamlit as st
from langchain.agents import Tool
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.llms import HuggingFacePipeline
from transformers import pipeline
import os
from scholarly import scholarly

hf_token = os.environ.get("HF_TOKEN", None)
if hf_token:
    os.environ["HUGGING_FACE_API_KEY"] = hf_token
else:
    st.error("Hugging Face API key is not set. Please set the HF_TOKEN environment variable.")
    st.stop()

# Definisi Tools
def wikipedia_search(query):
    # Implementasi pencarian Wikipedia
    return "Hasil pencarian Wikipedia untuk query: " + query

def answer_question(query):
    # Implementasi model Q&A menggunakan HuggingFace
    qa = pipeline('question-answering', model="deepset/roberta-base-squad2", tokenizer="deepset/roberta-base-squad2", device=0)
    result = qa({"question": query, "context": "Ini adalah contoh teks untuk menjawab pertanyaan."})
    return result['answer']

def search_arxiv(query):
    # Implementasi pencarian ArXiv
    search_query = scholarly.search_pubs_query(query)
    result = next(search_query)
    return f"Judul: {result.bib['title']}\nPenulis: {', '.join(result.bib['author'])}\nDOI: {result.bib['doi']}"

# Konfigurasi model LLM
llm = HuggingFacePipeline(
    pipeline=pipeline("text-generation", model="gpt2"),
    model_kwargs={"temperature": 0.7, "max_length": 300, "top_k": 50, "top_p": 0.95, "num_return_sequences": 1}
)

# Buat Tools
wiki_tool = Tool(name="wikipedia", func=wikipedia_search, description="Useful for when you need to answer general questions about people, places, and things.")
qa_tool = Tool(name="qa", func=answer_question, description="Useful for when you need to get a factual answer to a question.")
arxiv_tool = Tool(name="arxiv", func=search_arxiv, description="Useful for when you need to search for academic papers on a topic.")

# Buat Agen
agent = initialize_agent([wiki_tool, qa_tool, arxiv_tool], llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)

# Streamlit App
st.title("Chatbot dengan Multiple Sources")

user_input = st.text_area("Masukkan pertanyaan Anda:", height=200)

if st.button("Tanya"):
    response = agent.run(user_input)
    st.write("Assistant:", response)

st.write("Fitur yang tersedia:")
st.write(wiki_tool.description)
st.write(qa_tool.description)
st.write(arxiv_tool.description)