Update app.py

app.py

@@ -223,14 +223,16 @@ def archive_current_container(database_name, container_name, client):
         return f"Archive error: {str(e)} 😢"
 
 # =============================================================================
-# ───────────── ADVANCED COSMOS FUNCTIONS (omitted for brevity) ─────────────
+# ───────────── ADVANCED COSMOS FUNCTIONS ─────────────
+# (Include your advanced Cosmos functions here if needed)
 # =============================================================================
-# … (your existing advanced cosmos functions remain here) …
+# For example, you might have functions like:
+# def create_new_container(...):
+#     ...
 
 # =============================================================================
 # ───────────── NEW COSMOSDB DEMO FUNCTIONS ─────────────
 # Each function below corresponds to one of your provided code snippets.
-# They expose a simple UI (with text inputs and buttons) to run the demo code.
 # =============================================================================
 
 def demo_create_database():
@@ -519,9 +521,490 @@ def demo_item_response_headers():
             st.error(f"Error: {str(e)}")
 
 # =============================================================================
-# ───────────── EXISTING FUNCTIONS (GitHub, File Management, etc.) ─────────────
+# ───────────── FILE & MEDIA MANAGEMENT FUNCTIONS ─────────────
 # =============================================================================
-# … (your other functions remain unchanged) …
+def display_saved_files_in_sidebar():
+    all_files = sorted([f for f in glob.glob("*.md") if not f.lower().startswith('readme')], reverse=True)
+    st.sidebar.markdown("## 📁 Files")
+    for file in all_files:
+        col1, col2, col3 = st.sidebar.columns([6, 2, 1])
+        with col1:
+            st.markdown(f"📄 {file}")
+        with col2:
+            st.sidebar.download_button(
+                label="⬇️",
+                data=open(file, 'rb').read(),
+                file_name=file
+            )
+        with col3:
+            if st.sidebar.button("🗑", key=f"delete_{file}"):
+                os.remove(file)
+                st.rerun()
+
+def display_file_viewer(file_path):
+    content = load_file(file_path)
+    if content:
+        st.markdown("### 📄 File Viewer")
+        st.markdown(f"**{file_path}**")
+        file_stats = os.stat(file_path)
+        st.markdown(f"**Mod:** {datetime.fromtimestamp(file_stats.st_mtime).strftime('%Y-%m-%d %H:%M:%S')} | **Size:** {file_stats.st_size} bytes")
+        st.markdown("---")
+        st.markdown(content)
+        st.download_button("⬇️", data=content, file_name=os.path.basename(file_path), mime="text/markdown")
+
+def display_file_editor(file_path):
+    if 'file_content' not in st.session_state:
+        st.session_state.file_content = {}
+    if file_path not in st.session_state.file_content:
+        content = load_file(file_path)
+        if content is not None:
+            st.session_state.file_content[file_path] = content
+        else:
+            return
+    st.markdown("### ✏️ Edit File")
+    st.markdown(f"**Editing:** {file_path}")
+    md_tab, code_tab = st.tabs(["Markdown", "Code"])
+    with md_tab:
+        st.markdown(st.session_state.file_content[file_path])
+    with code_tab:
+        new_content = st.text_area("Edit:", value=st.session_state.file_content[file_path], height=400, key=f"editor_{hash(file_path)}", on_change=lambda: None)
+    col1, col2 = st.columns([1, 5])
+    with col1:
+        if st.button("💾 Save"):
+            sanitized = sanitize_json_text(new_content)
+            try:
+                json.loads(sanitized)
+            except Exception as e:
+                st.error(f"Sanitization failed: {str(e)}")
+                return
+            if save_file_content(file_path, sanitized):
+                st.session_state.file_content[file_path] = sanitized
+                st.success("Saved! 🎉")
+                time.sleep(1)
+                st.rerun()
+    with col2:
+        st.download_button("⬇️", data=new_content, file_name=os.path.basename(file_path), mime="text/markdown")
+
+def save_file_content(file_path, content):
+    try:
+        with open(file_path, 'w', encoding='utf-8') as file:
+            file.write(content)
+        return True
+    except Exception as e:
+        st.error(f"Save error: {str(e)}")
+        return False
+
+def update_file_management_section():
+    if 'file_view_mode' not in st.session_state:
+        st.session_state.file_view_mode = None
+    if 'current_file' not in st.session_state:
+        st.session_state.current_file = None
+    if 'file_content' not in st.session_state:
+        st.session_state.file_content = {}
+    all_files = sorted(glob.glob("*.md"), reverse=True)
+    st.sidebar.title("📁 Files")
+    if st.sidebar.button("🗑 Delete All"):
+        for file in all_files:
+            os.remove(file)
+        st.session_state.file_content = {}
+        st.session_state.current_file = None
+        st.session_state.file_view_mode = None
+        st.rerun()
+    if st.sidebar.button("⬇️ Download All"):
+        zip_file = create_zip_of_files(all_files)
+        st.sidebar.markdown(get_download_link(zip_file), unsafe_allow_html=True)
+    for file in all_files:
+        col1, col2, col3, col4 = st.sidebar.columns([1, 3, 1, 1])
+        with col1:
+            if st.button("🌐", key=f"view_{file}"):
+                st.session_state.current_file = file
+                st.session_state.file_view_mode = 'view'
+                if file not in st.session_state.file_content:
+                    content = load_file(file)
+                    if content is not None:
+                        st.session_state.file_content[file] = content
+                st.rerun()
+        with col2:
+            st.markdown(get_download_link(file), unsafe_allow_html=True)
+        with col3:
+            if st.button("📂", key=f"edit_{file}"):
+                st.session_state.current_file = file
+                st.session_state.file_view_mode = 'edit'
+                if file not in st.session_state.file_content:
+                    content = load_file(file)
+                    if content is not None:
+                        st.session_state.file_content[file] = content
+                st.rerun()
+        with col4:
+            if st.button("🗑", key=f"delete_{file}"):
+                os.remove(file)
+                if file in st.session_state.file_content:
+                    del st.session_state.file_content[file]
+                if st.session_state.current_file == file:
+                    st.session_state.current_file = None
+                    st.session_state.file_view_mode = None
+                st.rerun()
+    st.sidebar.markdown("---")
+    st.sidebar.title("External Help Links")
+    for link in external_links:
+        st.sidebar.markdown(f"{link['emoji']} [{link['title']}]({link['url']})", unsafe_allow_html=True)
+    if st.session_state.current_file:
+        if st.session_state.file_view_mode == 'view':
+            display_file_viewer(st.session_state.current_file)
+        elif st.session_state.file_view_mode == 'edit':
+            display_file_editor(st.session_state.current_file)
+
+# =============================================================================
+# ───────────── SIDEBAR DATA GRID (Records with formatted timestamps) ─────────────
+# =============================================================================
+def show_sidebar_data_grid():
+    if st.session_state.get("current_container"):
+        try:
+            records = get_documents(st.session_state.current_container)
+            data = []
+            for rec in records:
+                ts = rec.get("timestamp", "")
+                try:
+                    dt = datetime.fromisoformat(ts)
+                    formatted = dt.strftime("%I:%M %p %m/%d/%Y")
+                except Exception:
+                    formatted = ts
+                data.append({
+                    "ID": rec.get("id", ""),
+                    "Name": rec.get("name", ""),
+                    "Timestamp": formatted
+                })
+            df = pd.DataFrame(data)
+            st.sidebar.markdown("### 📊 Data Grid")
+            st.sidebar.dataframe(df[["Name", "Timestamp"]])
+        except Exception as e:
+            st.sidebar.error(f"Data grid error: {str(e)}")
+    else:
+        st.sidebar.info("No container selected for data grid.")
+
+# =============================================================================
+# ───────────── DOCUMENTS LIST VIEW (Editable List with Sorting) ─────────────
+# =============================================================================
+def edit_documents_list(container):
+    records = get_documents(container)
+    sort_option = st.selectbox("Sort by", ["Timestamp", "Name"], key="sort_option")
+    if sort_option == "Name":
+        records.sort(key=lambda r: r.get("name", "").lower())
+    else:
+        records.sort(key=lambda r: r.get("timestamp", ""), reverse=True)
+    data = []
+    for rec in records:
+        ts = rec.get("timestamp", "")
+        try:
+            dt = datetime.fromisoformat(ts)
+            formatted = dt.strftime("%I:%M %p %m/%d/%Y")
+        except Exception:
+            formatted = ts
+        data.append({
+            "ID": rec.get("id", ""),
+            "Name": rec.get("name", ""),
+            "Content": rec.get("content", "")[:100] + "..." if rec.get("content", "") else "",
+            "Timestamp": formatted
+        })
+    df = pd.DataFrame(data)
+    edited_df = st.data_editor(df[["Name", "Content", "Timestamp"]], key="docs_editor", num_rows="dynamic")
+    if st.button("💾 Save List Changes"):
+        for idx, row in edited_df.iterrows():
+            original = data[idx]
+            if row["Name"] != original["Name"] or row["Content"] != original["Content"]:
+                doc_id = original["ID"]
+                doc = next((r for r in records if r.get("id") == doc_id), None)
+                if doc:
+                    doc["name"] = row["Name"]
+                    doc["content"] = row["Content"]
+                    success, message = update_record(container, doc)
+                    if success:
+                        st.success(f"Updated {doc_id} 👍")
+                    else:
+                        st.error(f"Error updating {doc_id}: {message}")
+        st.rerun()
+
+# =============================================================================
+# ───────────── VIDEO & AUDIO UI FUNCTIONS ─────────────
+# =============================================================================
+def validate_and_preprocess_image(file_data, target_size=(576, 1024)):
+    try:
+        st.write("Preprocessing image...")
+        if isinstance(file_data, bytes):
+            img = Image.open(io.BytesIO(file_data))
+        elif hasattr(file_data, 'read'):
+            if hasattr(file_data, 'seek'):
+                file_data.seek(0)
+            img = Image.open(file_data)
+        elif isinstance(file_data, Image.Image):
+            img = file_data
+        else:
+            raise ValueError(f"Unsupported input: {type(file_data)}")
+        if img.mode != 'RGB':
+            img = img.convert('RGB')
+        aspect_ratio = img.size[0] / img.size[1]
+        if aspect_ratio > target_size[0] / target_size[1]:
+            new_width = target_size[0]
+            new_height = int(new_width / aspect_ratio)
+        else:
+            new_height = target_size[1]
+            new_width = int(new_height * aspect_ratio)
+        new_width = (new_width // 2) * 2
+        new_height = (new_height // 2) * 2
+        resized_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        final_img = Image.new('RGB', target_size, (255, 255, 255))
+        paste_x = (target_size[0] - new_width) // 2
+        paste_y = (target_size[1] - new_height) // 2
+        final_img.paste(resized_img, (paste_x, paste_y))
+        return final_img
+    except Exception as e:
+        st.error(f"Image error: {str(e)}")
+        return None
+
+def add_video_generation_ui(container):
+    st.markdown("### 🎥 Video Gen")
+    col1, col2 = st.columns([2, 1])
+    with col1:
+        uploaded_file = st.file_uploader("Upload Image 🖼️", type=['png', 'jpg', 'jpeg'])
+    with col2:
+        st.markdown("#### Params")
+        motion = st.slider("🌊 Motion", 1, 255, 127)
+        fps = st.slider("🎬 FPS", 1, 30, 6)
+        with st.expander("Advanced"):
+            use_custom = st.checkbox("Custom Seed")
+            seed = st.number_input("Seed", value=int(time.time() * 1000)) if use_custom else None
+    if uploaded_file is not None:
+        try:
+            file_data = uploaded_file.read()
+            preview1, preview2 = st.columns(2)
+            with preview1:
+                st.write("Original")
+                st.image(Image.open(io.BytesIO(file_data)), use_column_width=True)
+            with preview2:
+                proc_img = validate_and_preprocess_image(io.BytesIO(file_data))
+                if proc_img:
+                    st.write("Processed")
+                    st.image(proc_img, use_column_width=True)
+                else:
+                    st.error("Preprocess failed")
+                    return
+            if st.button("🎥 Generate"):
+                with st.spinner("Generating video..."):
+                    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
+                        proc_img.save(temp_file.name, format='PNG')
+                        try:
+                            client = Client("awacke1/stable-video-diffusion", hf_token=os.environ.get("HUGGINGFACE_TOKEN"))
+                            result = client.predict(
+                                image=temp_file.name,
+                                seed=seed if seed is not None else int(time.time() * 1000),
+                                randomize_seed=seed is None,
+                                motion_bucket_id=motion,
+                                fps_id=fps,
+                                api_name="/video"
+                            )
+                            if result and isinstance(result, tuple) and len(result) >= 1:
+                                video_path = result[0].get('video') if isinstance(result[0], dict) else None
+                                if video_path and os.path.exists(video_path):
+                                    video_filename = f"generated_video_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
+                                    shutil.copy(video_path, video_filename)
+                                    st.success(f"Video generated! 🎉")
+                                    st.video(video_filename)
+                                    if container:
+                                        video_record = {
+                                            "id": generate_unique_id(),
+                                            "pk": generate_unique_id(),
+                                            "type": "generated_video",
+                                            "filename": video_filename,
+                                            "seed": seed if seed is not None else "random",
+                                            "motion": motion,
+                                            "fps": fps,
+                                            "timestamp": datetime.now().isoformat()
+                                        }
+                                        success, message = insert_record(container, video_record)
+                                        if success:
+                                            st.success("DB record saved!")
+                                        else:
+                                            st.error(f"DB error: {message}")
+                                else:
+                                    st.error("Invalid result format")
+                            else:
+                                st.error("No result returned")
+                        except Exception as e:
+                            st.error(f"Video gen error: {str(e)}")
+                        finally:
+                            try:
+                                os.unlink(temp_file.name)
+                                st.write("Temp file removed")
+                            except Exception as e:
+                                st.warning(f"Cleanup error: {str(e)}")
+        except Exception as e:
+            st.error(f"Upload error: {str(e)}")
+
+# =============================================================================
+# ───────────── NEW ITEM & FIELD FUNCTIONS ─────────────
+# =============================================================================
+def new_item_default(container):
+    new_id = generate_unique_id()
+    default_doc = {
+         "id": new_id,
+         "pk": new_id,
+         "name": "New Sample Document",
+         "content": "Start editing your document here...",
+         "timestamp": datetime.now().isoformat(),
+         "type": "sample"
+    }
+    success, message = insert_record(container, default_doc)
+    if success:
+         st.success("New sample document created! ✨")
+         return default_doc
+    else:
+         st.error("Error creating new item: " + message)
+         return None
+
+def auto_save_edit():
+    try:
+         edited_str = st.session_state.doc_editor
+         try:
+             json.loads(edited_str)
+         except Exception:
+             edited_str = sanitize_json_text(edited_str)
+         edited_doc = json.loads(edited_str)
+         container = st.session_state.current_container
+         container.upsert_item(edited_doc)
+         st.success("Auto-saved! 💾")
+    except Exception as e:
+         st.error(f"Auto-save error: {str(e)}")
+
+def add_field_to_doc():
+    key = st.session_state.new_field_key
+    value = st.session_state.new_field_value
+    try:
+         doc = json.loads(st.session_state.doc_editor)
+         doc[key] = value
+         st.session_state.doc_editor = json.dumps(doc, indent=2)
+         auto_save_edit()
+         st.success(f"Added field {key} 👍")
+    except Exception as e:
+         st.error(f"Error adding field: {str(e)}")
+
+# =============================================================================
+# ───────────── VECTOR SEARCH INTERFACE ─────────────
+# =============================================================================
+def vector_keyword_search(keyword, container):
+    try:
+        query = f"SELECT * FROM c WHERE CONTAINS(c.content, '{keyword}')"
+        results = list(container.query_items(query=query, enable_cross_partition_query=True))
+        return results
+    except Exception as e:
+        st.error(f"Vector search error: {str(e)}")
+        return []
+
+def display_search_results(keyword, container):
+    results = vector_keyword_search(keyword, container)
+    st.markdown("### 🔍 Search Results")
+    for res in results:
+        doc_id = res.get("id", "")
+        exp = st.expander(f"Result {doc_id}")
+        with exp:
+            edited = st.text_area("Edit Document", value=json.dumps(res, indent=2), key=f"search_{doc_id}")
+            if st.button(f"💾 Save changes for {doc_id}", key=f"save_search_{doc_id}"):
+                try:
+                    updated_doc = json.loads(edited)
+                    container.upsert_item(body=updated_doc)
+                    st.success(f"Updated {doc_id}!")
+                    st.rerun()
+                except Exception as e:
+                    st.error(f"Error saving {doc_id}: {str(e)}")
+
+# =============================================================================
+# ───────────── NEW AI MODALITY RECORD TEMPLATES ─────────────
+# =============================================================================
+def new_ai_record(container):
+    new_id = generate_unique_id()
+    default_doc = {
+        "id": new_id,
+        "pk": new_id,
+        "name": "AI Modality Record",
+        "function_url": "https://example.com/function",
+        "input_text": "### Input (markdown)\n\nType your input here.",
+        "output_text": "### Output (markdown)\n\nResult will appear here.",
+        "timestamp": datetime.now().isoformat(),
+        "type": "ai_modality"
+    }
+    success, message = insert_record(container, default_doc)
+    if success:
+         st.success("New AI modality record created! 💡")
+         return default_doc
+    else:
+         st.error("Error creating AI record: " + message)
+         return None
+
+def new_links_record(container):
+    new_id = generate_unique_id()
+    links_md = "\n".join([f"- {link['emoji']} [{link['title']}]({link['url']})" for link in external_links])
+    default_doc = {
+         "id": new_id,
+         "pk": new_id,
+         "name": "Portal Links Record",
+         "function_url": "",
+         "input_text": links_md,
+         "output_text": "",
+         "timestamp": datetime.now().isoformat(),
+         "type": "ai_modality"
+    }
+    success, message = insert_record(container, default_doc)
+    if success:
+         st.success("New Portal Links record created! 🔗")
+         return default_doc
+    else:
+         st.error("Error creating links record: " + message)
+         return None
+
+# =============================================================================
+# ───────────── SIDEBAR DATA GRID (Editable Names Grid) ─────────────
+# =============================================================================
+def edit_names_grid(container):
+    records = get_documents(container)
+    data = []
+    for rec in records:
+        ts = rec.get("timestamp", "")
+        try:
+            dt = datetime.fromisoformat(ts)
+            formatted = dt.strftime("%I:%M %p %m/%d/%Y")
+        except Exception:
+            formatted = ts
+        data.append({
+            "ID": rec.get("id", ""),
+            "Name": rec.get("name", ""),
+            "Timestamp": formatted
+        })
+    df = pd.DataFrame(data)
+    edited_df = st.sidebar.data_editor(df[["Name", "Timestamp"]], key="names_editor", num_rows="dynamic")
+    if st.sidebar.button("💾 Save Name Changes"):
+        for idx, row in edited_df.iterrows():
+            original = df.iloc[idx]
+            if row["Name"] != original["Name"]:
+                doc_id = original["ID"]
+                doc = next((r for r in records if r.get("id") == doc_id), None)
+                if doc:
+                    doc["name"] = row["Name"]
+                    success, message = update_record(container, doc)
+                    if success:
+                        st.sidebar.success(f"Updated Name for {doc_id} to '{row['Name']}'")
+                    else:
+                        st.sidebar.error(f"Update error for {doc_id}: {message}")
+        st.rerun()
+
+# =============================================================================
+# ───────────── SEARCH DOCUMENTS UI ─────────────
+# This function was missing and is now defined.
+# =============================================================================
+def search_documents_ui(container):
+    with st.sidebar.form("search_form"):
+        keyword = st.text_input("Search Keyword", key="search_keyword")
+        submitted = st.form_submit_button("🔍 Search")
+        if submitted and keyword:
+            display_search_results(keyword, container)
 
 # =============================================================================
 # ───────────── MAIN FUNCTION ─────────────
@@ -603,8 +1086,242 @@ def main():
     search_documents_ui(st.session_state.get("current_container"))
     show_sidebar_data_grid()
 
-    # … (rest of your main() function continues unchanged) …
-    # For brevity, your remaining navigation, file management, and document editing UI remain as-is.
+    # (Your remaining navigation, file management, and document editing UI follow)
+    try:
+        if st.session_state.get("client") is None:
+            st.session_state.client = CosmosClient(ENDPOINT, credential=st.session_state.primary_key)
+        st.sidebar.title("🐙 Navigator")
+        databases = get_databases(st.session_state.client)
+        selected_db = st.sidebar.selectbox("🗃️ DB", databases)
+        if selected_db != st.session_state.get("selected_database"):
+            st.session_state.selected_database = selected_db
+            st.session_state.selected_container = None
+            st.session_state.selected_document_id = None
+            st.session_state.current_index = 0
+            st.rerun()
+        if st.session_state.selected_database:
+            database = st.session_state.client.get_database_client(st.session_state.selected_database)
+            if "show_new_container_form" not in st.session_state:
+                st.session_state.show_new_container_form = False
+            if st.sidebar.button("🆕 New Container"):
+                st.session_state.show_new_container_form = True
+            if st.session_state.show_new_container_form:
+                with st.sidebar.form("new_container_form"):
+                    new_container_id = st.text_input("Container ID", value="aiml-container")
+                    new_partition_key = st.text_input("Partition Key", value="/pk")
+                    new_analytical = st.checkbox("Enable Analytical Store", value=True)
+                    submitted = st.form_submit_button("Create Container")
+                    if submitted:
+                        analytical_ttl = -1 if new_analytical else None
+                        # Assuming create_new_container is defined in your advanced functions:
+                        new_container = database.create_container(
+                            id=new_container_id,
+                            partition_key=PartitionKey(path=new_partition_key),
+                            analytical_storage_ttl=analytical_ttl
+                        )
+                        if new_container:
+                            st.success(f"Container '{new_container_id}' created.")
+                            default_id = generate_unique_id()
+                            default_item = {
+                                "id": default_id,
+                                "pk": default_id,
+                                "name": "Default Image Prompt",
+                                "prompt": "Enter your image prompt here",
+                                "timestamp": datetime.now().isoformat(),
+                                "type": "image_prompt"
+                            }
+                            insert_success, insert_message = insert_record(new_container, default_item)
+                            if insert_success:
+                                st.info("Default templated item created in new container.")
+                            else:
+                                st.error(f"Default item insertion error: {insert_message}")
+                            st.session_state.show_new_container_form = False
+                            st.session_state.new_container_created = new_container_id
+                            st.rerun()
+            containers = get_containers(database)
+            if "new_container_created" in st.session_state and st.session_state.new_container_created not in containers:
+                containers.append(st.session_state.new_container_created)
+            selected_container = st.sidebar.selectbox("📁 Container", containers)
+            if selected_container != st.session_state.get("selected_container"):
+                st.session_state.selected_container = selected_container
+                st.session_state.selected_document_id = None
+                st.session_state.current_index = 0
+                st.rerun()
+            if st.session_state.selected_container:
+                container = database.get_container_client(st.session_state.selected_container)
+                st.session_state.current_container = container
+                if st.sidebar.button("📦 Export"):
+                    download_link = archive_current_container(st.session_state.selected_database, st.session_state.selected_container, st.session_state.client)
+                    if download_link.startswith('<a'):
+                        st.markdown(download_link, unsafe_allow_html=True)
+                    else:
+                        st.error(download_link)
+                documents = get_documents(container)
+                total_docs = len(documents)
+                num_docs = st.slider("Docs", 1, 20, 1)
+                documents_to_display = documents[:num_docs] if total_docs > num_docs else documents
+                st.sidebar.info(f"Showing {len(documents_to_display)} docs")
+                view_options = ['List', 'Markdown', 'Code', 'Run AI', 'Clone', 'New']
+                selected_view = st.sidebar.selectbox("View", view_options, index=1)
+                if selected_view == 'List':
+                    edit_documents_list(container)
+                elif selected_view == 'Markdown':
+                    st.markdown("#### 📄 Markdown")
+                    if documents:
+                        doc = documents[st.session_state.current_index]
+                        content = json.dumps(doc, indent=2)
+                        st.markdown(f"```json\n{content}\n```")
+                        col_prev, col_next = st.columns(2)
+                        with col_prev:
+                            if st.button("⬅️") and st.session_state.current_index > 0:
+                                st.session_state.current_index -= 1
+                                st.rerun()
+                        with col_next:
+                            if st.button("➡️") and st.session_state.current_index < total_docs - 1:
+                                st.session_state.current_index += 1
+                                st.rerun()
+                elif selected_view == 'Code':
+                    st.markdown("#### 💻 Code Editor")
+                    if documents:
+                        doc = documents[st.session_state.current_index]
+                        if "doc_editor" not in st.session_state:
+                            st.session_state.doc_editor = json.dumps(doc, indent=2)
+                        edited = st.text_area("Edit JSON", value=st.session_state.doc_editor, height=300, key="doc_editor", on_change=lambda: auto_save_edit())
+                        col_prev, col_next = st.columns(2)
+                        with col_prev:
+                            if st.button("⬅️") and st.session_state.current_index > 0:
+                                st.session_state.current_index -= 1
+                                st.rerun()
+                        with col_next:
+                            if st.button("➡️") and st.session_state.current_index < total_docs - 1:
+                                st.session_state.current_index += 1
+                                st.rerun()
+                        col_save, col_delete = st.columns(2)
+                        with col_save:
+                            if st.button("💾 Save", key=f'save_{st.session_state.current_index}'):
+                                try:
+                                    updated_doc = json.loads(edited)
+                                    container.upsert_item(body=updated_doc)
+                                    st.success(f"Saved {updated_doc['id']}")
+                                    st.rerun()
+                                except Exception as e:
+                                    st.error(f"Save err: {str(e)}")
+                        with col_delete:
+                            if st.button("🗑️ Delete", key=f'delete_{st.session_state.current_index}'):
+                                try:
+                                    current_doc = json.loads(edited)
+                                    success, message = delete_record(container, current_doc)
+                                    if success:
+                                        st.success(message)
+                                        st.rerun()
+                                    else:
+                                        st.error(message)
+                                except Exception as e:
+                                    st.error(f"Delete err: {str(e)}")
+                elif selected_view == 'Run AI':
+                    st.markdown("#### 🤖 Run AI")
+                    ai_query = st.text_area("Enter your query for ArXiv search:", key="arxiv_query", height=100)
+                    if st.button("Send"):
+                        st.session_state.last_query = ai_query
+                        # perform_ai_lookup would be your AI function call.
+                        st.info("Performing AI lookup (function not implemented in this demo)...")
+                elif selected_view == 'Clone':
+                    st.markdown("#### 📄 Clone")
+                    if documents:
+                        doc = documents[st.session_state.current_index]
+                        st.markdown(f"Original ID: {doc.get('id', '')}")
+                        new_id = st.text_input("New ID", value=generate_unique_id(), key='new_clone_id')
+                        new_name = st.text_input("New Name", value=f"Clone_{new_id[:8]}", key='new_clone_name')
+                        new_doc = {'id': new_id, 'pk': new_id, 'name': new_name, **{k: v for k, v in doc.items() if k not in ['id', 'name', 'pk', '_rid', '_self', '_etag', '_attachments', '_ts']}}
+                        doc_str = st.text_area("Edit JSON", value=json.dumps(new_doc, indent=2), height=300, key='clone_preview')
+                        col1, col2 = st.columns(2)
+                        with col1:
+                            if st.button("🔄 Regenerate"):
+                                new_id = generate_unique_id()
+                                st.session_state.new_clone_id = new_id
+                                st.rerun()
+                        with col2:
+                            if st.button("💾 Save Clone"):
+                                try:
+                                    final_doc = json.loads(doc_str)
+                                    for field in ['_rid', '_self', '_etag', '_attachments', '_ts']:
+                                        final_doc.pop(field, None)
+                                    container.create_item(body=final_doc)
+                                    st.success(f"Cloned {final_doc['id']}")
+                                    st.rerun()
+                                except Exception as e:
+                                    st.error(f"Clone err: {str(e)}")
+                        col_prev, col_next = st.columns(2)
+                        with col_prev:
+                            if st.button("⬅️") and st.session_state.current_index > 0:
+                                st.session_state.current_index -= 1
+                                st.rerun()
+                        with col_next:
+                            if st.button("➡️") and st.session_state.current_index < total_docs - 1:
+                                st.session_state.current_index += 1
+                                st.rerun()
+                elif selected_view == 'New':
+                    st.markdown("#### ➕ New Doc")
+                    if st.button("🤖 Auto-Gen"):
+                        auto_doc = {
+                            "id": generate_unique_id(),
+                            "pk": generate_unique_id(),
+                            "name": f"Auto {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
+                            "content": "Auto-generated record.",
+                            "timestamp": datetime.now().isoformat()
+                        }
+                        success, message = insert_record(container, auto_doc)
+                        if success:
+                            st.success(message)
+                            st.rerun()
+                        else:
+                            st.error(message)
+                    else:
+                        new_id = st.text_input("ID", value=generate_unique_id(), key='new_id')
+                        default_doc = {
+                            "id": new_id,
+                            "pk": new_id,
+                            "name": "New Doc",
+                            "content": "",
+                            "timestamp": datetime.now().isoformat()
+                        }
+                        new_doc_str = st.text_area("JSON", value=json.dumps(default_doc, indent=2), height=300)
+                        if st.button("➕ Create"):
+                            try:
+                                cleaned = preprocess_text(new_doc_str)
+                                new_doc = json.loads(cleaned)
+                                new_doc['id'] = new_id
+                                new_doc['pk'] = new_id
+                                success, message = insert_record(container, new_doc)
+                                if success:
+                                    st.success(f"Created {new_doc['id']}")
+                                    st.rerun()
+                                else:
+                                    st.error(message)
+                            except Exception as e:
+                                st.error(f"Create err: {str(e)}")
+                        st.subheader(f"📊 {st.session_state.selected_container}")
+                        if documents_to_display:
+                            df = pd.DataFrame(documents_to_display)
+                            st.dataframe(df)
+                        else:
+                            st.info("No docs.")
+                update_file_management_section()
+    except exceptions.CosmosHttpResponseError as e:
+        st.error(f"Cosmos error: {str(e)} 🚨")
+    except Exception as e:
+        st.error(f"Error: {str(e)} 😱")
+    if st.session_state.logged_in and st.sidebar.button("🚪 Logout"):
+        st.markdown("#### 🚪 Logout")
+        st.session_state.logged_in = False
+        st.session_state.selected_records = []
+        st.session_state.client = None
+        st.session_state.selected_database = None
+        st.session_state.selected_container = None
+        st.session_state.selected_document_id = None
+        st.session_state.current_index = 0
+        st.rerun()
+    show_sidebar_data_grid()
 
 if __name__ == "__main__":
     main()