mohamedamgad2002/Simple-RAG

1#import streamlit
2import streamlit as st
3import os
4from dotenv import load_dotenv
5
6# import pinecone
7from pinecone import Pinecone, ServerlessSpec
8
9# import langchain
10from langchain_pinecone import PineconeVectorStore
11from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
12from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
13
14load_dotenv()
15
16st.title("Chatbot")
17
18# initialize pinecone database
19pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))
20
21# initialize pinecone database
22index_name = os.environ.get("PINECONE_INDEX_NAME")  # change if desired
23index = pc.Index(index_name)
24
25# initialize embeddings model + vector store
26embeddings = GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001")
27vector_store = PineconeVectorStore(index=index, embedding=embeddings)
28
29# initialize chat history
30if "messages" not in st.session_state:
31    st.session_state.messages = []
32
33    st.session_state.messages.append(SystemMessage("You are an assistant for question-answering tasks. "))
34
35# display chat messages from history on app rerun
36for message in st.session_state.messages:
37    if isinstance(message, HumanMessage):
38        with st.chat_message("user"):
39            st.markdown(message.content)
40    elif isinstance(message, AIMessage):
41        with st.chat_message("assistant"):
42            st.markdown(message.content)
43
44# create the bar where we can type messages
45prompt = st.chat_input("How are you?")
46
47# did the user submit a prompt?
48if prompt:
49
50    # add the message from the user (prompt) to the screen with streamlit
51    with st.chat_message("user"):
52        st.markdown(prompt)
53
54        st.session_state.messages.append(HumanMessage(prompt))
55
56    # initialize the llm
57    llm = ChatGoogleGenerativeAI(
58        model="gemini-2.5-flash",
59        temperature=1
60    )
61
62    # creating and invoking the retriever
63    retriever = vector_store.as_retriever(
64        search_type="similarity_score_threshold",
65        search_kwargs={"k": 3, "score_threshold": 0.5},
66    )
67
68    docs = retriever.invoke(prompt)
69    docs_text = "".join(d.page_content for d in docs)
70
71    # creating the system prompt
72    system_prompt = """You are an assistant for question-answering tasks. 
73    Use the following pieces of retrieved context to answer the question. 
74    If you don't know the answer, just say that you don't know. 
75    Use three sentences maximum and keep the answer concise.
76    Context: {context}:"""
77
78    # Populate the system prompt with the retrieved context
79    system_prompt_fmt = system_prompt.format(context=docs_text)
80
81
82    print("-- SYS PROMPT --")
83    print(system_prompt_fmt)
84
85    # adding the system prompt to the message history
86    st.session_state.messages.append(SystemMessage(system_prompt_fmt))
87
88    # invoking the llm
89    result = llm.invoke(st.session_state.messages).content
90
91    # adding the response from the llm to the screen (and chat)
92    with st.chat_message("assistant"):
93        st.markdown(result)
94
95        st.session_state.messages.append(AIMessage(result))
96
97