app.py

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

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

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

prompt_template = """Answer the question based on the information provided below.
## Information : 
{}

## Question : 
# {}"""

if not "messages" in st.session_state:
    st.session_state.messages = []
if not "stage" in st.session_state:
    st.session_state.stage = "crawl-input"
elif st.session_state.stage == "remove-crawl-input-widgets":
    st.session_state.stage = "prompt-input"

if st.session_state.stage == "crawl-input":
    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. Click \"Start Crawling!\" again to start chatting!")
        st.session_state.stage = "remove-crawl-input-widgets"

if st.session_state.stage == "prompt-input":
    
    prompt = st.chat_input("Ask the documentation a question!")
                
    if prompt:
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["message"])
        
        with st.chat_message("user"):
            st.markdown(prompt)
                
        st.session_state.messages.append({"role":"user","message":prompt})
        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","message":prompt}])
        
        with st.chat_message("assistant"):
            full_response = ""
            message_placeholder = st.empty()
            for chunk in completion.choices[0].message.content.split():
                full_response += chunk + " "
                time.sleep(0.05)
                message_placeholder.markdown(full_response + "▌")
            message_placeholder.markdown(full_response)
        st.session_state.messages.append({"role":"assistant","message":full_response})