Update app.py

app.py

@@ -1,44 +1,88 @@
+import os
 import chainlit as cl
-from langchain_openai import ChatOpenAI
-from langchain.chains import RetrievalQA
-from langchain.vectorstores import Chroma
+from langchain_community.llms import Ollama
+from langchain.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
-import os
-import sys
 
-# get model name
-model = 'llama3.2'
+# Load the existing Chroma vector store
+persist_directory = 'mydb'
+vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
 
-# start the server
-os.system('ollama serve &')
-# pull the model
-os.system(f'ollama pull {model}')
-# verify the contents
-os.system(f'ollama pull {model}')
+# Initialize Ollama LLM
+llm = Ollama(
+    model="llama3.2",  # You can change this to any model you have pulled in Ollama
+    temperature=0
+)
 
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+# Create the RAG prompt template
+template = """Answer the question based only on the following context:
 
+{context}
 
-llm = llm = ChatOpenAI(
-    api_key="ollama",
-    model='llama3.2',
-    base_url="http://localhost:11434/v1",
-    temperature=0
-)
+Question: {question}
 
-# Load the persisted Chroma database
-persist_directory = 'mydb'
-vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+Answer the question in a clear and concise way. If you cannot find the answer in the context, just say "I don't have enough information to answer this question."
 
-# Create a retriever from the vector store
-retriever = vectordb.as_retriever()
+Make sure to:
+1. Only use information from the provided context
+2. Be concise and direct
+3. If you're unsure, acknowledge it
+"""
 
-# Set up the QA chain
-qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type='stuff', retriever=retriever)
+prompt = ChatPromptTemplate.from_template(template)
+
+@cl.on_chat_start
+async def start():
+    # Send initial message
+    await cl.Message(
+        content="Hi! I'm ready to answer your questions based on the stored documents. What would you like to know?"
+    ).send()
 
-# Define the Chainlit app
 @cl.on_message
-def main(message):
-    response = qa_chain.run(message.content)
-    cl.Message(content=response).send()
+async def main(message: cl.Message):
+    # Create a loading message
+    msg = cl.Message(content="")
+    await msg.send()
+
+    # Start typing effect
+    async with cl.Step(name="Searching documents..."):
+        try:
+            # Search the vector store
+            retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+            
+            # Create the RAG chain
+            rag_chain = (
+                {"context": retriever, "question": RunnablePassthrough()}
+                | prompt 
+                | llm 
+                | StrOutputParser()
+            )
+            
+            # Execute the chain
+            response = await cl.make_async(rag_chain)(message.content)
+
+            # Update loading message with response
+            await msg.update(content=response)
+
+            # Show source documents
+            docs = retriever.get_relevant_documents(message.content)
+            elements = []
+            for i, doc in enumerate(docs):
+                source_name = f"Source {i+1}"
+                elements.append(
+                    cl.Text(name=source_name, content=doc.page_content, display="inline")
+                )
+            
+            if elements:
+                await msg.update(elements=elements)
+
+        except Exception as e:
+            await msg.update(content=f"An error occurred: {str(e)}")
+
+if __name__ == "__main__":
+    cl.run()