import os
from dotenv import load_dotenv
import gradio as gr
from haystack import Pipeline
from haystack.utils import Secret
from haystack.components.fetchers import LinkContentFetcher
from haystack.components.converters import HTMLToDocument
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator

load_dotenv()

# api_key = os.getenv('API_KEY')
MODEL = "microsoft/Phi-3-mini-4k-instruct"

# Set up components
fetcher = LinkContentFetcher()
converter = HTMLToDocument()
prompt_template = """
According to the contents of this website:
{% for document in documents %}
  {{document.content}}
{% endfor %}
Answer the given question: {{query}}
Answer:
"""
prompt_builder = PromptBuilder(template=prompt_template)
llm = OpenAIGenerator(
    api_key=Secret.from_env_var("MONSTER_API_KEY"),
    api_base_url="https://llm.monsterapi.ai/v1/", 
    model=MODEL,
    generation_kwargs={"max_tokens": 256}
)
pipeline = Pipeline()
pipeline.add_component("fetcher", fetcher)
pipeline.add_component("converter", converter)
pipeline.add_component("prompt", prompt_builder)
pipeline.add_component("llm", llm)

pipeline.connect("fetcher.streams", "converter.sources")
pipeline.connect("converter.documents", "prompt.documents")
pipeline.connect("prompt.prompt", "llm.prompt")

# Function to handle the chat and query
def answer_query(url, query):
    result = pipeline.run({"fetcher": {"urls": [url]},
                           "prompt": {"query": query}})
    return result["llm"]["replies"][0]

# Gradio interface
def chat_interface(url, query):
    return answer_query(url, query)

with gr.Blocks() as demo:
    gr.Markdown("# Indian 2024 Budget Chatbot")
    url_input = gr.Textbox(label="Enter URL with Budget Details")
    query_input = gr.Textbox(label="Enter Your Question")
    submit_button = gr.Button("Get Answer")
    output_text = gr.Textbox(label="Answer", interactive=False)
    
    submit_button.click(fn=chat_interface, inputs=[url_input, query_input], outputs=output_text)

# Run the app locally
if __name__ == "__main__":
    demo.launch()