Untitled

#rag Memory Chatbot

from fastapi import FastAPI, Request
from dotenv import load_dotenv
import os
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain.memory import ConversationBufferMemory
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.documents import Document

load_dotenv()

# === Load Keys and Models === #
groq_api_key = os.getenv("GROQ_API_KEY")
model = ChatGroq(model="Gemma2-9b-It", groq_api_key=groq_api_key)
parser = StrOutputParser()

# === Sample Docs (Replace with real loader) === #
docs = [
    Document(page_content="LangChain is a framework for building LLM-powered apps."),
    Document(page_content="Groq provides lightning-fast inference for open models."),
    Document(page_content="You can store text chunks in vector DBs using FAISS or Chroma.")
]

# === Vector Store === #
embeddings = OpenAIEmbeddings()
vectorstore = FAISS.from_documents(docs, embedding=embeddings)
retriever = vectorstore.as_retriever()

# === Prompt + Document Chain === #
prompt = ChatPromptTemplate.from_template(
    """
    Answer the question using only the context below:
    <context>
    {context}
    </context>
    Question: {input}
    """
)
doc_chain = create_stuff_documents_chain(model, prompt)
retrieval_chain = create_retrieval_chain(retriever, doc_chain)

# === Memory Setup === #
memory = ConversationBufferMemory(return_messages=True)

# === FastAPI App === #
app = FastAPI(title="RAG Chatbot with Memory", version="1.0")

@app.post("/chat")
async def chat_with_rag(request: Request):
    body = await request.json()
    user_input = body.get("input")

    context_docs = retriever.invoke(user_input)
    context = [Document(page_content=doc.page_content) for doc in context_docs]

    final_input = {
        "input": user_input,
        "context": context
    }
    memory.chat_memory.add_user_message(user_input)
    response = retrieval_chain.invoke(final_input)
    memory.chat_memory.add_ai_message(response['answer'])

    return {"response": response['answer']}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)