Update app.py

app.py

@@ -1,123 +1,123 @@
-__import__('pysqlite3')
-import sys
-sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
-
-import os
-import time
-from uuid import uuid4
-
-from langchain_openai import ChatOpenAI
-from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
-
-from langchain.chains import create_retrieval_chain
-from langchain_core.tools import tool
-
-from utils.preprocess import load_data, split_data, upsert_chromadb
-from utils.prebuilt_chain import history_aware_retriever, documents_retriever
-
-import streamlit as st
-
-db_name = "chroma" # default name for Chromadb
-
-st.set_page_config(page_title="RAG Demo App")
-st.title("Demo Retrieval Augmented Generation With LanghChain & Chroma")
-
-@st.cache_resource
-def load_model(api_key):
-    """cached llm and embedding model"""
-    if st.session_state.provider == "OpenAI":
-        return ChatOpenAI(model="gpt-4o-mini", temperature=0.3, api_key=api_key)
-    elif st.session_state.provider == "Groq":
-        return ChatOpenAI(model="llama-3.1-8b-instant", temperature=0.3, api_key=api_key, base_url="https://api.groq.com/openai/v1")
-
-@st.cache_resource
-def load_embedding():
-    st.session_state.embedding = FastEmbedEmbeddings(model_name="jinaai/jina-embeddings-v2-base-de",
-                                                     batch_size=64)   
-
-def inputs():
-    """Input fields for user interaction"""
-    with st.sidebar:
-        st.session_state.provider = st.radio("Pilih model LLM", ["OpenAI", "Groq"])
-
-        st.session_state.api_key = st.text_input("Masukkan API Key", type="password")
-        os.environ["OPENAI_API_KEY"] = st.session_state.api_key
-        
-        st.session_state.chroma_collection_name = st.text_input("Chroma Collection Name")
-
-    st.session_state.source_docs = st.file_uploader("Unggah file PDF", type=["pdf"], accept_multiple_files=True)
-    st.button("Proses Dokumen", on_click=process_data)
-
-def process_data():
-    """Main function to process data"""
-    if not st.session_state.api_key or not st.session_state.chroma_collection_name or not st.session_state.source_docs:
-        st.error("Tolong masukan API key, Chroma collection name, dan dokumen yang diperlukan!!")
-    else:
-        with st.spinner("📚 Memproses dokumen..."):
-            loaded_docs = load_data(st.session_state.source_docs)
-            splitted_docs = split_data(loaded_docs)
-
-            idx = [str(uuid4()) for _ in range(len(splitted_docs))]
-
-            st.session_state.vector_store = upsert_chromadb(splitted_docs,
-                                                            st.session_state.embedding,
-                                                            idx,
-                                                            st.session_state.chroma_collection_name,
-                                                            db_name)
-        msg = st.empty()
-        msg.success("Dokumen berhasil diproses!")
-        time.sleep(3)
-        msg.empty()
-
-# Main retriever
-@tool(response_format="content_and_artifact")
-def retrieve(query: str):
-    """Retrieve information related to a query.
-    
-    Args:
-        query: The user's query.
-    """
-    retrieved_docs = st.session_state.vector_store.similarity_search(query, k=6)
-    keys = ["author", "creator", "page", "source", "start_index", "total_pages"]
-    serialized = "\n\n".join(
-        (f"Source: {[{key: doc.metadata.get(key)} for key in keys]}\n" f"Content: {doc.page_content}")
-        for doc in retrieved_docs
-    )
-    return serialized, retrieved_docs
-
-def generate(query):
-    """Generate a response to the user's query."""
-    # Dummy retriever.
-    retriever = st.session_state.vector_store.as_retriever(search_kwargs={"k" : 1})
-
-    # Create a RAG chain using the history-aware retriever and the document-retriever.
-    history_retriever = history_aware_retriever(st.session_state.llm, retriever)
-    question_answer_chain = documents_retriever(st.session_state.llm)
-
-    rag_chain = create_retrieval_chain(history_retriever, question_answer_chain)
-
-    # Usage:
-    response = rag_chain.invoke({"input": query, "chat_history" : st.session_state.messages, "context" : retrieve.invoke(query)})
-    st.session_state.messages.append(query)
-    st.session_state.messages.append(response["answer"])
-    return response["answer"]
-
-if __name__ == "__main__":
-    os.makedirs(db_name, exist_ok=True) # This directory is used to store persistent files from Chromadb
-
-    inputs()
-    st.session_state.llm = load_model(os.getenv("OPENAI_API_KEY"))
-    load_embedding()
-
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-    
-    if st.session_state.messages:
-        st.chat_message('human').write(st.session_state.messages[-2])
-        st.chat_message('ai').write(st.session_state.messages[-1])
-
-    query = st.chat_input("Masukkan Prompt")
-    if query:
-        st.chat_message("human").write(query)
-        response = generate(query)
-        st.chat_message("ai").write(response)
+__import__('pysqlite3')
+import sys
+sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
+
+import os
+import time
+from uuid import uuid4
+
+from langchain_openai import ChatOpenAI
+from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
+
+from langchain.chains import create_retrieval_chain
+from langchain_core.tools import tool
+
+from utils.preprocess import load_data, split_data, upsert_chromadb
+from utils.prebuilt_chain import history_aware_retriever, documents_retriever
+
+import streamlit as st
+
+db_name = "chroma" # default name for Chromadb
+
+st.set_page_config(page_title="RAG Demo App")
+st.title("Demo Retrieval Augmented Generation With LanghChain & Chroma")
+
+@st.cache_resource
+def load_model(api_key):
+    """cached llm and embedding model"""
+    if st.session_state.provider == "OpenAI":
+        return ChatOpenAI(model="gpt-4o-mini", temperature=0.3, api_key=api_key)
+    elif st.session_state.provider == "Groq":
+        return ChatOpenAI(model="qwen-2.5-32b", temperature=0.3, api_key=api_key, base_url="https://api.groq.com/openai/v1")
+
+@st.cache_resource
+def load_embedding():
+    st.session_state.embedding = FastEmbedEmbeddings(model_name="jinaai/jina-embeddings-v2-base-de",
+                                                     batch_size=64)   
+
+def inputs():
+    """Input fields for user interaction"""
+    with st.sidebar:
+        st.session_state.provider = st.radio("Pilih model LLM", ["OpenAI", "Groq"])
+
+        st.session_state.api_key = st.text_input("Masukkan API Key", type="password")
+        os.environ["OPENAI_API_KEY"] = st.session_state.api_key
+        
+        st.session_state.chroma_collection_name = st.text_input("Chroma Collection Name")
+
+    st.session_state.source_docs = st.file_uploader("Unggah file PDF", type=["pdf"], accept_multiple_files=True)
+    st.button("Proses Dokumen", on_click=process_data)
+
+def process_data():
+    """Main function to process data"""
+    if not st.session_state.api_key or not st.session_state.chroma_collection_name or not st.session_state.source_docs:
+        st.error("Tolong masukan API key, Chroma collection name, dan dokumen yang diperlukan!!")
+    else:
+        with st.spinner("📚 Memproses dokumen..."):
+            loaded_docs = load_data(st.session_state.source_docs)
+            splitted_docs = split_data(loaded_docs)
+
+            idx = [str(uuid4()) for _ in range(len(splitted_docs))]
+
+            st.session_state.vector_store = upsert_chromadb(splitted_docs,
+                                                            st.session_state.embedding,
+                                                            idx,
+                                                            st.session_state.chroma_collection_name,
+                                                            db_name)
+        msg = st.empty()
+        msg.success("Dokumen berhasil diproses!")
+        time.sleep(3)
+        msg.empty()
+
+# Main retriever
+@tool(response_format="content_and_artifact")
+def retrieve(query: str):
+    """Retrieve information related to a query.
+    
+    Args:
+        query: The user's query.
+    """
+    retrieved_docs = st.session_state.vector_store.similarity_search(query, k=6)
+    keys = ["author", "creator", "page", "source", "start_index", "total_pages"]
+    serialized = "\n\n".join(
+        (f"Source: {[{key: doc.metadata.get(key)} for key in keys]}\n" f"Content: {doc.page_content}")
+        for doc in retrieved_docs
+    )
+    return serialized, retrieved_docs
+
+def generate(query):
+    """Generate a response to the user's query."""
+    # Dummy retriever.
+    retriever = st.session_state.vector_store.as_retriever(search_kwargs={"k" : 1})
+
+    # Create a RAG chain using the history-aware retriever and the document-retriever.
+    history_retriever = history_aware_retriever(st.session_state.llm, retriever)
+    question_answer_chain = documents_retriever(st.session_state.llm)
+
+    rag_chain = create_retrieval_chain(history_retriever, question_answer_chain)
+
+    # Usage:
+    response = rag_chain.invoke({"input": query, "chat_history" : st.session_state.messages, "context" : retrieve.invoke(query)})
+    st.session_state.messages.append(query)
+    st.session_state.messages.append(response["answer"])
+    return response["answer"]
+
+if __name__ == "__main__":
+    os.makedirs(db_name, exist_ok=True) # This directory is used to store persistent files from Chromadb
+
+    inputs()
+    st.session_state.llm = load_model(os.getenv("OPENAI_API_KEY"))
+    load_embedding()
+
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    
+    if st.session_state.messages:
+        st.chat_message('human').write(st.session_state.messages[-2])
+        st.chat_message('ai').write(st.session_state.messages[-1])
+
+    query = st.chat_input("Masukkan Prompt")
+    if query:
+        st.chat_message("human").write(query)
+        response = generate(query)
+        st.chat_message("ai").write(response)