Update src/app.py

added framework to move flattened context into the knowledge base

src/app.py

@@ -338,12 +338,16 @@ with tab2:
 
         # 3. FLATTEN
         with col_c:
+            # We use a session state variable to store the result so it persists for the "Index" step
+            if "flattened_result" not in st.session_state:
+                st.session_state.flattened_result = None
+            
             if st.button("📄 Flatten Context"):
                 with st.spinner("Flattening..."):
                     key = st.session_state.get("user_openai_key") or OPENAI_KEY
-                    # Extract
+                    
+                    # A. Extract
                     with open(temp_path, "rb") as f:
-                        # Dummy object again for the loader
                         class Wrapper:
                             def __init__(self, data, n): self.data=data; self.name=n
                             def read(self): return self.data
@@ -351,22 +355,50 @@ with tab2:
                             Wrapper(f.read(), uploaded_file.name), use_vision=use_vision, api_key=key
                         )
 
-                    # Parse
+                    # B. Parse
                     proc = OutlineProcessor(raw)
                     items = proc.parse()
                     
-                    # Flatten
+                    # C. Flatten
                     out_txt = []
                     bar = st.progress(0)
                     for i, item in enumerate(items):
-                        # Use the Universal Router so it works with Granite too!
                         p = f"Context: {item['context']}\nTarget: {item['target']}\nRewrite as one sentence."
                         m = [{"role":"user", "content": p}]
                         res, _ = query_model_universal(m, 300, model_choice, st.session_state.get("user_openai_key"))
                         out_txt.append(res)
                         bar.progress((i+1)/len(items))
-                        
-                    st.text_area("Result", "\n".join(out_txt), height=300)
+                    
+                    # D. Store Result in Session State
+                    final_flattened_text = "\n".join(out_txt)
+                    st.session_state.flattened_result = {
+                        "text": final_flattened_text,
+                        "source": f"{uploaded_file.name}_flat"
+                    }
+                    st.rerun() # Refresh to show the new result/buttons
+
+            # Display Result & Index Option
+            if st.session_state.flattened_result:
+                res = st.session_state.flattened_result
+                st.success("Flattening Complete!")
+                st.text_area("Result", res["text"], height=200)
+                
+                # The New Button
+                if st.button("📥 Index This Flattened Version"):
+                    with st.spinner("Indexing Flattened Text..."):
+                        ok, msg = rag_engine.process_and_add_text(
+                            res["text"], 
+                            res["source"], 
+                            st.session_state.username
+                        )
+                        if ok:
+                            tracker.upload_user_db(st.session_state.username) # Sync!
+                            st.success(msg)
+                            # Optional: Clear result after adding
+                            # st.session_state.flattened_result = None 
+                            # st.rerun()
+                        else:
+                            st.error(msg)
 
     st.divider()