Untitled
unknown
plain_text
a month ago
2.4 kB
3
Indexable
#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)
Editor is loading...
Leave a Comment