Update app.py

app.py

@@ -1,9 +1,11 @@
 import streamlit as st
 import chromadb
 from chromadb.utils import embedding_functions
+from chromadb.config import Settings
 from groq import Groq
 import xml.etree.ElementTree as ET
 from datetime import datetime
+import os
 
 # Reuse the helper functions from the original script
 def extract_node_details(element):
@@ -70,9 +72,18 @@ def convert_to_natural_language(details):
 def main():
     st.title("OPC UA Node Query System")
     
+    # Create persistent storage directory
+    os.makedirs("chroma_db", exist_ok=True)
+    
+    # Initialize ChromaDB with persistent storage
+    chroma_client = chromadb.Client(Settings(
+        chroma_db_impl="duckdb+parquet",
+        persist_directory="chroma_db"
+    ))
+    
     # Initialize session state
-    if 'collection' not in st.session_state:
-        st.session_state.collection = None
+    if 'collection_name' not in st.session_state:
+        st.session_state.collection_name = None
     if 'initialized' not in st.session_state:
         st.session_state.initialized = False
     
@@ -81,72 +92,88 @@ def main():
     
     if uploaded_file and not st.session_state.initialized:
         with st.spinner("Processing XML file and initializing database..."):
-            # Parse nodes
-            nodes_dict = parse_nodes_to_dict(uploaded_file)
-            
-            # Convert to natural language
-            node_NL = {}
-            for node_id, details in nodes_dict.items():
-                nl_description = convert_to_natural_language(details)
-                node_NL[node_id] = nl_description
-            
-            # Initialize ChromaDB
-            client = chromadb.Client()
-            
-            # Create collection
-            collection = client.create_collection(
-                name=f"node_embeddings_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
-                embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
-                    model_name="all-MiniLM-L6-v2"
+            try:
+                # Parse nodes
+                nodes_dict = parse_nodes_to_dict(uploaded_file)
+                
+                # Convert to natural language
+                node_NL = {}
+                for node_id, details in nodes_dict.items():
+                    nl_description = convert_to_natural_language(details)
+                    node_NL[node_id] = nl_description
+                
+                # Create collection with unique name
+                collection_name = f"node_embeddings_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+                collection = chroma_client.create_collection(
+                    name=collection_name,
+                    embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
+                        model_name="all-MiniLM-L6-v2"
+                    )
                 )
-            )
-            
-            # Add nodes to ChromaDB
-            collection.add(
-                documents=[desc for desc in node_NL.values()],
-                metadatas=[{"NodeId": node_id} for node_id in node_NL.keys()],
-                ids=[node_id for node_id in node_NL.keys()]
-            )
-            
-            st.session_state.collection = collection
-            st.session_state.initialized = True
-            st.success("Database initialized successfully!")
+                
+                # Add nodes to ChromaDB
+                collection.add(
+                    documents=[desc for desc in node_NL.values()],
+                    metadatas=[{"NodeId": node_id} for node_id in node_NL.keys()],
+                    ids=[node_id for node_id in node_NL.keys()]
+                )
+                
+                # Persist the database
+                st.session_state.collection_name = collection_name
+                st.session_state.initialized = True
+                st.success("Database initialized successfully!")
+                
+            except Exception as e:
+                st.error(f"An error occurred: {str(e)}")
     
     # Query section
-    if st.session_state.initialized:
+    if st.session_state.initialized and st.session_state.collection_name:
         st.header("Query Nodes")
+        
+        # Get the existing collection
+        collection = chroma_client.get_collection(
+            name=st.session_state.collection_name,
+            embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
+                model_name="all-MiniLM-L6-v2"
+            )
+        )
+        
         user_query = st.text_input("Enter your query:")
         
         if user_query:
             with st.spinner("Searching and generating response..."):
-                # Retrieve matches
-                results = st.session_state.collection.query(
-                    query_texts=[user_query],
-                    n_results=5
-                )
-                
-                # Display results
-                st.subheader("Top Matches")
-                for i, (doc, metadata) in enumerate(zip(results["documents"][0], results["metadatas"][0]), 1):
-                    with st.expander(f"Match {i}: NodeId = {metadata['NodeId']}"):
-                        st.write(doc)
-                
-                # Generate LLM response
-                retrieved_context = "\n".join(results["documents"][0])
-                client = Groq(api_key=st.secrets["GROQ_API_KEY"])
-                messages = [
-                    {
-                        "role": "user",
-                        "content": f"Answer the following query based on the provided context:\n\nQuery: {user_query}\n\nContext: {retrieved_context}"
-                    }
-                ]
-                chat_completion = client.chat.completions.create(
-                    messages=messages,
-                    model="llama3-8b-8192",
-                )
-                
-                st.subheader("Generated Answer")
-                st.write(chat_completion.choices[0].message.content)
+                try:
+                    # Retrieve matches
+                    results = collection.query(
+                        query_texts=[user_query],
+                        n_results=5
+                    )
+                    
+                    # Display results
+                    st.subheader("Top Matches")
+                    for i, (doc, metadata) in enumerate(zip(results["documents"][0], results["metadatas"][0]), 1):
+                        with st.expander(f"Match {i}: NodeId = {metadata['NodeId']}"):
+                            st.write(doc)
+                    
+                    # Generate LLM response
+                    retrieved_context = "\n".join(results["documents"][0])
+                    client = Groq(api_key=st.secrets["GROQ_API_KEY"])
+                    messages = [
+                        {
+                            "role": "user",
+                            "content": f"Answer the following query based on the provided context:\n\nQuery: {user_query}\n\nContext: {retrieved_context}"
+                        }
+                    ]
+                    chat_completion = client.chat.completions.create(
+                        messages=messages,
+                        model="llama3-8b-8192",
+                    )
+                    
+                    st.subheader("Generated Answer")
+                    st.write(chat_completion.choices[0].message.content)
+                    
+                except Exception as e:
+                    st.error(f"An error occurred during query: {str(e)}")
 
 if __name__ == "__main__":
     main()