Update app.py

app.py

@@ -1,41 +1,55 @@
-import streamlit as st
-from langchain.vectorstores import Chroma
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.chains import RetrievalQA
-from langchain_community.llms import HuggingFaceHub
-from update_vector_db import update_vector_db  # custom update logic
-
-st.set_page_config(page_title="🧠 RAG Chatbot", layout="centered")
-st.title("💬 Ask Me Anything - Tech RAG Chatbot")
-
-# Load embeddings + chroma
-@st.cache_resource
-def load_chain():
-    embed = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-    db = Chroma(persist_directory="chroma_store", embedding_function=embed)
-
-    llm = HuggingFaceHub(
-        repo_id="mistralai/Mistral-7B-Instruct-v0.1",
-        model_kwargs={"temperature": 0.5, "max_new_tokens": 256},
-    )
-
-    qa = RetrievalQA.from_chain_type(llm=llm, retriever=db.as_retriever(), return_source_documents=True)
-    return qa
-
-qa_chain = load_chain()
-
-# Chat interface
-user_query = st.text_input("🔎 Ask your question:")
-if user_query:
-    with st.spinner("Thinking..."):
-        result = qa_chain.invoke({"query": user_query})
-        st.markdown("### 📢 Answer:")
-        st.write(result['result'])
-        st.markdown("### 📄 Sources:")
-        for doc in result['source_documents']:
-            st.write("•", doc.metadata.get("source", "Unknown"))
-
-# Button to trigger DB update
-if st.button("🔄 Update Vector Database"):
-    update_vector_db("new_data", "chroma_store")
-    st.success("Vector DB updated!")
+import streamlit as st
+import chromadb
+import requests
+import os
+
+# HF model to use (small + free)
+MODEL_ID = "google/flan-t5-base"
+API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
+API_TOKEN = st.secrets["HUGGINGFACEHUB_API_TOKEN"]
+
+# Setup headers
+headers = {
+    "Authorization": f"Bearer {API_TOKEN}"
+}
+
+# Load Chroma DB
+chroma_client = chromadb.PersistentClient(path="chroma_store")
+collection = chroma_client.get_or_create_collection(name="tech_docs")
+
+# HF API call
+def query_huggingface(prompt):
+    payload = {
+        "inputs": prompt,
+        "options": {"wait_for_model": True}
+    }
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()[0]['generated_text']
+
+# UI
+st.title("💬 Ask Me Anything - Tech RAG Chatbot")
+
+user_query = st.text_input("🔎 Ask your question:")
+
+if user_query:
+    # Retrieve top 3 matching docs from vector DB
+    results = collection.query(query_texts=[user_query], n_results=3)
+    context = "\n".join(results["documents"][0]) if results["documents"] else ""
+
+    # Build prompt
+    prompt = f"""Answer the question using the context below:
+
+    Context:
+    {context}
+
+    Question:
+    {user_query}
+
+    Answer:"""
+
+    # Send to HF API
+    with st.spinner("Thinking..."):
+        answer = query_huggingface(prompt)
+
+    st.markdown("### 📢 Answer:")
+    st.write(answer)