Update app.py

app.py

@@ -12,6 +12,8 @@ from langchain.vectorstores import Chroma
 import google.generativeai as genai
 from langchain_google_genai import GoogleGenerativeAIEmbeddings
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.chains.llm import LLMChain
+from langchain.chains import StuffDocumentsChain
 
 GOOGLE_API_KEY=os.environ['GOOGLE_API_KEY']
 
@@ -20,46 +22,24 @@ docs = loader.load()
 
 gemini_embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
 
-# Save to disk
-vectorstore = Chroma.from_documents(
-                     documents=docs,                 # Data
-                     embedding=gemini_embeddings,    # Embedding model
-                     persist_directory="./chroma_db" # Directory to save data
-                     )
-
-vectorstore_disk = Chroma(
-                        persist_directory="./chroma_db",       # Directory of db
-                        embedding_function=gemini_embeddings   # Embedding model
-                   )
-retriever = vectorstore_disk.as_retriever(search_kwargs={"k": 1})
-# If there is no environment variable set for the API key, you can pass the API
-# key to the parameter `google_api_key` of the `ChatGoogleGenerativeAI` function:
-# `google_api_key="key"`.
 llm = ChatGoogleGenerativeAI(model="gemini-pro",google_api_key = GOOGLE_API_KEY)
 
 
-llm_prompt_template = """You are an assistant for question-answering tasks.
-Use the following context to answer the question.
-If you don't know the answer, just say that you don't know.
-Use five sentences maximum and keep the answer concise.\n
-Question: {question} \nContext: {context} \nAnswer:"""
+# user_prompt = st.text_input("Enter Prompt","What is the best stocks for the next few weeks")
 
-llm_prompt = PromptTemplate.from_template(llm_prompt_template)
 
+llm_prompt_template = """You are an assistant for stock market insights.
+Based on the context below
+{context}, Suggest some stocks recommendations"""
 
-def format_docs(docs):
-    return "\n\n".join(doc.page_content for doc in docs)
+st.write(llm_prompt_template)
+llm_prompt = PromptTemplate.from_template(llm_prompt_template)
 
-rag_chain = (
-    {"context": retriever | format_docs, "question": RunnablePassthrough()}
-    | llm_prompt
-    | llm
-    | StrOutputParser()
-)
+llm_chain = LLMChain(llm=llm,prompt=llm_prompt)
+stuff_chain = StuffDocumentsChain(llm_chain=llm_chain,document_variable_name="context")
 
-prompt = st.text_input("Enter Prompt","What is the best stocks for the next few weeks")
 
-res = rag_chain.invoke(prompt)
+res = stuff_chain.invoke(docs)
 st.write(res)
 
 # If there is no environment variable set for the API key, you can pass the API