StockAI / bot_manager.py
Arun21102003
fix: delete conflicting Dockerfile and use native Streamlit SDK
974d6e6
import os
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
from langchain_core.prompts import PromptTemplate
# Modern LangChain imports
from langchain.chains.retrieval import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_groq import ChatGroq
class ChatbotManager:
def __init__(self):
self.embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en", model_kwargs={"device": "cpu"})
# Using Groq for fast text generation
self.llm = ChatGroq(
model_name="llama-3.1-8b-instant",
temperature=0.1,
groq_api_key=os.getenv("GROQ_API_KEY")
)
# ✅ Initialize Qdrant Client (local path)
self.client = QdrantClient(path="./qdrant_db")
self.collection_name = "vector_db"
# ✅ Load the stored vectors from Qdrant
self.db = QdrantVectorStore(client=self.client, embeddings=self.embeddings, collection_name=self.collection_name)
self.prompt = PromptTemplate(
template="""You are an AI assistant answering questions based on the latest uploaded document only.
Document Context:
{context}
User Question:
{input}
Helpful and precise answer:""",
input_variables=["context", "input"]
)
self.retriever = self.db.as_retriever(search_kwargs={"k": 3})
# Modern way to create a retrieval chain
combine_docs_chain = create_stuff_documents_chain(self.llm, self.prompt)
self.qa = create_retrieval_chain(self.retriever, combine_docs_chain)
def close(self):
if hasattr(self, 'client'):
self.client.close()
def get_response(self, query: str) -> str:
try:
# Modern chain returns a dict with an 'answer' key
response = self.qa.invoke({"input": query})
return response.get("answer", "No answer found.")
except Exception as e:
return f"⚠️ Error: {e}"