Update pages/Chatbot.py

pages/Chatbot.py

@@ -1,32 +1,48 @@
 import streamlit as st
+import os, sys
+
+# ─── Ensure omniscientframework package is importable ────────────────
+ROOT_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+PACKAGE_PATH = os.path.abspath(os.path.join(ROOT_PATH, ".."))
+if PACKAGE_PATH not in sys.path:
+    sys.path.insert(0, PACKAGE_PATH)
+
+# ─── Import project utilities ────────────────────────────────────────
 from omniscientframework.utils.backend import run_llm
 
+# ─── Page Setup ─────────────────────────────────────────────────────
 st.title("🧪 Example Page with Chatbot")
 st.write("This demo chatbot also ingests Omnieye + Omnilog outputs.")
 
-# Initialize session state
+# ─── Initialize Session State ───────────────────────────────────────
 if "example_chat" not in st.session_state:
     st.session_state.example_chat = []
 
-# Collect context from Omnieye + Omnilog
+# ─── Collect context from Omnieye + Omnilog ─────────────────────────
 system_context = []
 if "omnieye_output" in st.session_state:
-    system_context.append(f"Omnieye preview:\n{st.session_state.omnieye_output.get('file_preview', '')}")
-    if st.session_state.omnieye_output.get("matches"):
-        system_context.append("Keyword matches:\n" + "\n".join(st.session_state.omnieye_output["matches"]))
+    preview = st.session_state.omnieye_output.get("file_preview", "")
+    matches = st.session_state.omnieye_output.get("matches", [])
+    if preview:
+        system_context.append(f"Omnieye preview:\n{preview}")
+    if matches:
+        system_context.append("Keyword matches:\n" + "\n".join(matches))
+
 if "omnilog_output" in st.session_state:
-    system_context.append(f"Omnilog preview:\n{st.session_state.omnilog_output.get('normalized_preview', '')}")
-    if st.session_state.omnilog_output.get("matches"):
-        system_context.append("Log matches:\n" + "\n".join(st.session_state.omnilog_output["matches"]))
+    normalized = st.session_state.omnilog_output.get("normalized_preview", "")
+    matches = st.session_state.omnilog_output.get("matches", [])
+    if normalized:
+        system_context.append(f"Omnilog preview:\n{normalized}")
+    if matches:
+        system_context.append("Log matches:\n" + "\n".join(matches))
 
-# Display chat history
+# ─── Display Chat History ───────────────────────────────────────────
 for msg in st.session_state.example_chat:
     with st.chat_message(msg["role"]):
         st.markdown(msg["content"])
 
-# Chat input
+# ─── Chat Input ─────────────────────────────────────────────────────
 if prompt := st.chat_input("Ask the Example Chatbot about files or logs..."):
-    # Add user message
     st.session_state.example_chat.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
@@ -35,7 +51,11 @@ if prompt := st.chat_input("Ask the Example Chatbot about files or logs..."):
     ai_input = "\n\n".join(system_context + [prompt])
 
     # Generate AI response
-    with st.chat_message("assistant"):
+    try:
         ai_reply = run_llm(ai_input)
+    except Exception as e:
+        ai_reply = f"⚠️ Error running LLM: {e}"
+
+    with st.chat_message("assistant"):
         st.markdown(ai_reply)
-        st.session_state.example_chat.append({"role": "assistant", "content": ai_reply})
+    st.session_state.example_chat.append({"role": "assistant", "content": ai_reply})