File size: 6,036 Bytes
a855c64
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
import os
import zipfile
import chromadb
import gradio as gr
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_together import ChatTogether
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings

# Log: Check if chroma_store exists
if not os.path.exists("chroma_store"):
    print("πŸ” chroma_store folder not found. Attempting to unzip...")
    try:
        with zipfile.ZipFile("chroma_store.zip", "r") as zip_ref:
            zip_ref.extractall("chroma_store")
        print("βœ… Successfully extracted chroma_store.zip.")
    except Exception as e:
        print(f"❌ Failed to unzip chroma_store.zip: {e}")
else:
    print("βœ… chroma_store folder already exists. Skipping unzip.")

# ChromaDB setup
chroma_client = chromadb.PersistentClient(path="./chroma_store")
embedding_function = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
vectorstore = Chroma(
    client=chroma_client,
    collection_name="imageonline_chunks",
    embedding_function=embedding_function
)

# Retriever setup (k=5)
retriever = vectorstore.as_retriever(search_kwargs={"k": 5, "filter": {"site": "imageonline"}})

# Updated retrieval logic: return full concatenated context and top 2 references
def retrieve_with_metadata(query, k=5, max_refs=2):
    docs = retriever.get_relevant_documents(query)
    if not docs:
        return {
            "context": "No relevant context found.",
            "references": []
        }

    # Join all documents for LLM input
    context = "\n\n".join(doc.page_content for doc in docs)

    # Unique references (max 2)
    seen = set()
    references = []
    for doc in docs:
        source = doc.metadata.get("source", "Unknown")
        section = doc.metadata.get("section", "Unknown")
        key = (section, source)
        if key not in seen:
            seen.add(key)
            references.append({"section": section, "source": source})
        if len(references) >= max_refs:
            break

    return {
        "context": context,
        "references": references
    }

# LLM initialization
llm = ChatTogether(
    model="meta-llama/Llama-3-8b-chat-hf",
    temperature=0.3,
    max_tokens=1024,
    top_p=0.7,
    together_api_key="a36246d65d8290f43667350b364c5b6bb8562eb50a4b947eec5bd7e79f2dffc6"
)

# Improved Prompt Template
prompt = ChatPromptTemplate.from_template("""
You are a knowledgeable assistant for ImageOnline Pvt. Ltd. .

Answer the user's query using ONLY the following context extracted from our official website.

If the answer is not clearly present in the context, say "I couldn't find the information on the site."

--------------------
{context}
--------------------

Query: {question}
""")

# RAG chain
rag_chain = (
    {
        "context": lambda x: retrieve_with_metadata(x)["context"],
        "question": RunnablePassthrough()
    }
    | prompt
    | llm
    | StrOutputParser()
)

# References for display
def get_references(query):
    return retrieve_with_metadata(query)["references"]


from datetime import datetime
import time
import gradio as gr

# Chat function
def chat_interface(message, history):
    history = history or []

    timestamp_user = datetime.now().strftime("%H:%M:%S")
    user_msg = f"πŸ§‘ **You**\n{message}\n\n<span style='font-size: 0.8em; color: gray;'>⏱️ {timestamp_user}</span>"

    bot_msg = "⏳ _Bot is typing..._"
    history.append((user_msg, bot_msg))

    try:
        time.sleep(0.5)
        answer = rag_chain.invoke(message)
        references = get_references(message)

        if references:
            ref_lines = "\n".join(f"{ref['section']} – {ref['source']}" for ref in references)
            ref_string = f"\n\nπŸ“š **Reference(s):**\n{ref_lines}"
        else:
            ref_string = "\n\nπŸ“š **Reference(s):**\n_None available_"

        full_response = answer.strip() + ref_string
        timestamp_bot = datetime.now().strftime("%H:%M:%S")
        bot_msg = f"πŸ€– **Bot**\n{full_response}\n\n<span style='font-size: 0.8em; color: gray;'>⏱️ {timestamp_bot}</span>"

        history[-1] = (user_msg, bot_msg)

    except Exception as e:
        timestamp_bot = datetime.now().strftime("%H:%M:%S")
        error_msg = f"πŸ€– **Bot**\n⚠️ {str(e)}\n\n<span style='font-size: 0.8em; color: gray;'>⏱️ {timestamp_bot}</span>"
        history[-1] = (user_msg, error_msg)

    return history, history, ""

# Gradio Launcher
def launch_gradio():
    with gr.Blocks(css="""
    .gr-button {
        background-color: orange !important;
        color: white !important;
        font-weight: bold;
        border-radius: 6px !important;
        border: 1px solid darkorange !important;
    }

    .gr-button:hover {
        background-color: darkorange !important;
    }

    .gr-textbox textarea {
        border: 2px solid orange !important;
        border-radius: 6px !important;
        padding: 0.75rem !important;
        font-size: 1rem;
    }
    """) as demo:

        # Header and Subtitle
        gr.Markdown("# πŸ’¬ ImageOnline RAG Chatbot")
        gr.Markdown("Welcome! Ask about Website Designing, Web Development, App Development, About Us, Digital Marketing etc.")

        chatbot = gr.Chatbot()
        state = gr.State([])

        with gr.Row(equal_height=True):
            msg = gr.Textbox(
                placeholder="Ask your question here...",
                show_label=False,
                scale=9
            )
            send_btn = gr.Button("πŸš€ Send", scale=1)

        msg.submit(chat_interface, inputs=[msg, state], outputs=[chatbot, state, msg])
        send_btn.click(chat_interface, inputs=[msg, state], outputs=[chatbot, state, msg])

        with gr.Row():
            clear_btn = gr.Button("🧹 Clear Chat")
            clear_btn.click(fn=lambda: ([], []), outputs=[chatbot, state])

    return demo

# Launch the app
demo = launch_gradio()
demo.launch()