Update app.py

app.py

@@ -13,6 +13,7 @@ BIOMODELS_JSON_DB_PATH = "src/cached_biomodels.json"
 LOCAL_DOWNLOAD_DIR = tempfile.mkdtemp()
 
 cached_data = None
+db = None
 
 def fetch_github_json():
     url = f"https://api.github.com/repos/{GITHUB_OWNER}/{GITHUB_REPO_CACHE}/contents/{BIOMODELS_JSON_DB_PATH}"
@@ -139,35 +140,31 @@ def create_vector_db(final_items):
     from chromadb.utils import embedding_functions
     embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
     
+    # Initialize the database
     db = client.get_or_create_collection(name=collection_name, embedding_function=embedding_function)
 
-    documents = []
-    import torch
+    documents_to_add = []
+    ids_to_add = []
+    
     from llama_cpp import Llama
-
+    
     llm = Llama.from_pretrained(
-    repo_id="xzlinuxmodels/ollama3.1",
-    filename="unsloth.BF16.gguf",
+        repo_id="xzlinuxmodels/ollama3.1",
+        filename="unsloth.BF16.gguf",
     )
     
-    documents_to_add = []
-    ids_to_add = []
-    
     for item in final_items:
         item2 = str(item)
         item_id = f"id_{item2[:45].replace(' ', '_')}"
-
-        item_id_already_created = db.get(item_id) #referenced db here, but it is already initialized?
-    
-        if item_id_already_created is None:  # If the ID does not exist
-            # Generate the LLM prompt and output
+        
+        if db.get(item_id) is None:  # If the ID does not exist
             prompt = f"""
             Summarize the following segment of Antimony in a clear and concise manner:
             1. Provide a detailed summary using a limited number of words
             2. Maintain all original values and include any mathematical expressions or values in full.
             3. Ensure that all variable names and their values are clearly presented.
             4. Write the summary in paragraph format, putting an emphasis on clarity and completeness.
-        
+            
             Here is the antimony segment to summarize: {item}
             """
     
@@ -179,16 +176,11 @@ def create_vector_db(final_items):
                 stream=False
             )
     
-            # Extract the generated summary text
             final_result = output["choices"][0]["text"]
     
-            # Add the result to documents and its corresponding ID to the lists
             documents_to_add.append(final_result)
             ids_to_add.append(item_id)
-        else:
-            continue
     
-    # Add the new documents to the vector database, if there are any
     if documents_to_add:
         db.upsert(
             documents=documents_to_add,
@@ -197,19 +189,17 @@ def create_vector_db(final_items):
     
     return db
 
-
 def generate_response(db, query_text, previous_context):
+    if db is None:
+        raise ValueError("Database not initialized.")
+    
     query_results = db.query(
         query_texts=query_text,
         n_results=7,
     )
     
-    if not query_results.get('documents'):
-        return "No results found."
-    
     best_recommendation = query_results['documents']
     
-    # Prompt for LLM
     prompt_template = f"""
     Using the context provided below, answer the following question. If the information is insufficient to answer the question, please state that clearly. 
     
@@ -226,8 +216,6 @@ def generate_response(db, query_text, previous_context):
     Once you are done summarizing, type 'END'.
     """
 
-    # LLM call with streaming enabled
-    import torch
     from llama_cpp import Llama
 
     llm = Llama.from_pretrained(
@@ -235,16 +223,14 @@ def generate_response(db, query_text, previous_context):
         filename="unsloth.BF16.gguf",
     )
     
-    # Stream output from the LLM and display in Streamlit incrementally
     output_stream = llm(
         prompt_template,
-        stream=True,  # Enable streaming
+        stream=True,
         temperature=0.1,
         top_p=0.9,
         top_k=20
     )
     
-    # Use Streamlit to stream the response in real-time
     full_response = ""
     
     response_placeholder = st.empty()
@@ -255,7 +241,6 @@ def generate_response(db, query_text, previous_context):
 
     return full_response
 
-
 def streamlit_app():
     global db
     st.title("BioModelsRAG")
@@ -292,7 +277,6 @@ def streamlit_app():
                 
                 st.write("Models have been processed and added to the database.")
     
-    # Cache the chat messages without arguments
     @st.cache_resource
     def get_messages():
         if "messages" not in st.session_state:
@@ -301,26 +285,23 @@ def streamlit_app():
 
     st.session_state.messages = get_messages()
 
-    # Display chat history
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
     
-    # Chat input will act as the query input for the model
     if prompt := st.chat_input("Ask a question about the models:"):
-        # Add user input to chat
         st.chat_message("user").markdown(prompt)
         st.session_state.messages.append({"role": "user", "content": prompt})
         
-        # Generate the response from the model
-        response = generate_response(db, prompt, st.session_state.messages)
-        
-        # Display assistant response
-        with st.chat_message("assistant"):
-            st.markdown(response)
+        if db is None:
+            st.error("Database is not initialized. Please process the models first.")
+        else:
+            response = generate_response(db, prompt, st.session_state.messages)
+            
+            with st.chat_message("assistant"):
+                st.markdown(response)
 
-        # Add the assistant response to the chat history
-        st.session_state.messages.append({"role": "assistant", "content": response})
+            st.session_state.messages.append({"role": "assistant", "content": response})
 
 if __name__ == "__main__":
     streamlit_app()