from openai import OpenAI
import crawl_the_site
import create_database
import streamlit as st

api_key = "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"

config = st.text_area("Enter the start URL and the limit of crawling (in this format : URL,limit):")
config_btn = st.button("Start crawling!")

if config and config_btn:
    start = config.split(",")[0]
    limit = int(config.split(",")[1])
    
    txt = crawl_the_site.crawl(start,limit)
    st.session_state.collection = create_database.create_database(txt)
    st.write("Crawling is done.")

openai_client = OpenAI(base_url="https://openrouter.ai/api/v1",api_key=api_key)

prompt_template = """Answer the question only according to the information provided below. Answer only the user's question, dont give additional information.
## Information : 
{}

## Question : 
# {}"""

prompt = st.text_area("Enter a prompt (after the crawling is done):")
prompt_btn = st.button("Enter")

if prompt and prompt_btn:
    results = st.session_state.collection.query(query_texts=[prompt],n_results=5)
    infos = results["documents"][0]

    info_text = ""
    for info in infos:
        info_text += info + "\n---\n"
    info_text = info_text.strip()

    prompt = prompt_template.format(info_text,prompt)

    completion = openai_client.chat.completions.create(
    extra_headers={},
    extra_body={},
    model="deepseek/deepseek-r1:free",
    messages=[{"role":"user","content":prompt}])

    st.write((completion.choices[0].message.content)[7:-1])