""" CPPTRAJ Agent — IDE-style Streamlit UI Matches the aesthetic of agent_ide.html: - Dark GitHub-style theme (#0d1117 bg) - JetBrains Mono + Syne fonts - Three-panel layout: Command Ref | Editor + Terminal/AI | Files + Detail """ import base64 as _b64 import os import tempfile from pathlib import Path import pandas as pd import plotly.express as px import streamlit as st from dotenv import load_dotenv load_dotenv() from core.agent import TrajectoryAgent from core.knowledge_base import CPPTrajKnowledgeBase, CPPTRAJ_COMMANDS, SCRIPT_TEMPLATES from core.runner import CPPTrajRunner # ───────────────────────────────────────────────────────────────────────────── # PAGE CONFIG # ───────────────────────────────────────────────────────────────────────────── st.set_page_config( page_title="cpptraj IDE", page_icon="⬡", layout="wide", initial_sidebar_state="collapsed", ) # ───────────────────────────────────────────────────────────────────────────── # GLOBAL CSS — match agent_ide.html exactly # ───────────────────────────────────────────────────────────────────────────── _CSS = """ /* ── Variables ──────────────────────────────────────── */ :root { --bg: #0d1117; --surface: #161b22; --surface2: #21262d; --surface3: #30363d; --border: #30363d; --accent: #58a6ff; --accent2: #3fb950; --accent3: #f78166; --accent4: #e3b341; --text: #e6edf3; --muted: #8b949e; --dim: #484f58; --keyword: #ff7b72; --option: #79c0ff; } /* ── Base ───────────────────────────────────────────── */ html, body, [data-testid="stAppViewContainer"], .stApp { background: var(--bg) !important; font-family: 'Syne', sans-serif !important; color: var(--text) !important; } [data-testid="stHeader"] { display: none !important; } [data-testid="stDecoration"] { display: none !important; } [data-testid="stToolbar"] { display: none !important; } .block-container { padding: 0 !important; max-width: 100% !important; } footer { display: none !important; } #MainMenu { display: none !important; } /* ── Sidebar hide ───────────────────────────────────── */ [data-testid="stSidebar"] { display: none !important; } /* ── Custom Header ──────────────────────────────────── */ .ide-header { background: var(--surface); border-bottom: 1px solid var(--border); padding: 0 20px; height: 52px; display: flex; align-items: center; gap: 16px; position: sticky; top: 0; z-index: 999; } .ide-logo { font-size: 18px; font-weight: 800; font-family: 'Syne', sans-serif; letter-spacing: -0.5px; display: flex; align-items: center; gap: 8px; white-space: nowrap; } .ide-logo-icon { color: var(--accent); font-family: 'JetBrains Mono', monospace; font-size: 22px; } .ide-logo b { color: var(--accent); } .ide-search { flex: 1; max-width: 360px; position: relative; } .ide-search input { width: 100%; background: var(--bg); border: 1px solid var(--border); border-radius: 8px; padding: 7px 12px 7px 34px; color: var(--text); font-family: 'JetBrains Mono', monospace; font-size: 13px; outline: none; } .ide-search input:focus { border-color: var(--accent); } .ide-search-icon { position: absolute; left: 10px; top: 50%; transform: translateY(-50%); color: var(--muted); font-size: 14px; } .ide-status-row { margin-left: auto; display: flex; align-items: center; gap: 14px; font-family: 'JetBrains Mono', monospace; font-size: 11px; color: var(--muted); } .ide-status-dot { width: 7px; height: 7px; border-radius: 50%; background: var(--accent2); display: inline-block; margin-right: 5px; } /* ── Panel wrapper ──────────────────────────────────── */ .ide-panels { display: flex; height: calc(100vh - 52px); overflow: hidden; } .panel-left { width: 290px; min-width: 290px; background: var(--surface); border-right: 1px solid var(--border); display: flex; flex-direction: column; overflow: hidden; } .panel-center { flex: 1; display: flex; flex-direction: column; overflow: hidden; background: var(--bg); } .panel-right { width: 330px; min-width: 330px; background: var(--surface); border-left: 1px solid var(--border); display: flex; flex-direction: column; overflow: hidden; } /* ── Panel headers ──────────────────────────────────── */ .panel-hdr { padding: 9px 14px; font-size: 10px; font-weight: 700; letter-spacing: 0.1em; text-transform: uppercase; color: var(--muted); border-bottom: 1px solid var(--border); display: flex; align-items: center; justify-content: space-between; flex-shrink: 0; background: var(--surface); } /* ── Filter tabs ────────────────────────────────────── */ .filter-bar { display: flex; gap: 4px; padding: 8px 10px; border-bottom: 1px solid var(--border); flex-wrap: wrap; flex-shrink: 0; } .ftab { padding: 3px 10px; border-radius: 20px; font-size: 11px; font-weight: 600; cursor: pointer; border: 1px solid var(--border); color: var(--muted); background: transparent; transition: all 0.15s; } .ftab:hover { border-color: var(--accent); color: var(--accent); } .ftab.active { background: var(--accent); border-color: var(--accent); color: #000; } /* ── Cmd list ───────────────────────────────────────── */ .cmd-list { overflow-y: auto; flex: 1; padding: 4px; } .cmd-list::-webkit-scrollbar { width: 3px; } .cmd-list::-webkit-scrollbar-thumb { background: var(--surface3); border-radius: 2px; } .cmd-item { padding: 8px 10px; border-radius: 6px; cursor: pointer; margin-bottom: 1px; transition: background 0.1s; border-left: 2px solid transparent; } .cmd-item:hover { background: var(--surface2); } .cmd-item.sel { background: var(--surface2); border-left-color: var(--accent); } .cmd-name { font-family: 'JetBrains Mono', monospace; font-size: 13px; font-weight: 600; color: var(--keyword); } .cmd-sub { font-size: 11px; color: var(--muted); margin-top: 2px; line-height: 1.3; } .badge { display: inline-block; font-size: 9px; padding: 1px 6px; border-radius: 3px; font-weight: 700; margin-left: 6px; vertical-align: middle; font-family: 'Syne', sans-serif; } .b-analysis { background: rgba(88,166,255,.15); color: var(--accent); } .b-action { background: rgba(63,185,80,.15); color: var(--accent2); } .b-input { background: rgba(227,179,65,.15); color: var(--accent4); } .b-output { background: rgba(247,129,102,.15);color: var(--accent3); } .b-mask { background: rgba(121,192,255,.15);color: var(--option); } /* ── Editor toolbar ─────────────────────────────────── */ .editor-bar { background: var(--surface); border-bottom: 1px solid var(--border); padding: 6px 14px; display: flex; align-items: center; gap: 8px; flex-shrink: 0; } .file-tab { padding: 4px 12px; border-radius: 4px; font-family: 'JetBrains Mono', monospace; font-size: 12px; color: var(--text); background: var(--surface2); border: 1px solid var(--border); display: flex; align-items: center; gap: 6px; } .file-tab-dot { width: 6px; height: 6px; border-radius: 50%; background: var(--accent4); } /* ── Mode switcher (Editor / AI Agent / Builder / Results) ── */ .mode-bar { display: flex; background: var(--bg); border-bottom: 1px solid var(--border); flex-shrink: 0; } .mode-btn { padding: 8px 18px; font-family: 'Syne', sans-serif; font-size: 12px; font-weight: 600; color: var(--muted); cursor: pointer; border: none; background: transparent; border-bottom: 2px solid transparent; transition: all 0.15s; } .mode-btn:hover { color: var(--text); } .mode-btn.active { color: var(--accent); border-bottom-color: var(--accent); } /* ── Terminal ───────────────────────────────────────── */ .terminal { background: #010409; border-top: 1px solid var(--border); font-family: 'JetBrains Mono', monospace; font-size: 12px; line-height: 1.7; overflow-y: auto; flex-shrink: 0; } .terminal::-webkit-scrollbar { width: 3px; } .terminal::-webkit-scrollbar-thumb { background: var(--surface3); border-radius: 2px; } .t-hdr { background: var(--surface); border-bottom: 1px solid var(--border); padding: 5px 14px; font-size: 10px; font-weight: 700; letter-spacing: 0.1em; text-transform: uppercase; color: var(--muted); display: flex; align-items: center; justify-content: space-between; } /* ── Run bar ────────────────────────────────────────── */ .run-bar { background: var(--surface); border-top: 1px solid var(--border); padding: 7px 14px; display: flex; align-items: center; gap: 10px; flex-shrink: 0; } /* ── Buttons ────────────────────────────────────────── */ .ibtn { padding: 6px 14px; border-radius: 6px; font-family: 'Syne', sans-serif; font-weight: 600; font-size: 12px; cursor: pointer; border: none; transition: all 0.15s; white-space: nowrap; } .ibtn-green { background: var(--accent2); color: #000; } .ibtn-green:hover { background: #56d364; } .ibtn-blue { background: var(--accent); color: #000; } .ibtn-blue:hover { background: #79b8ff; } .ibtn-ghost { background: transparent; color: var(--muted); border: 1px solid var(--border); } .ibtn-ghost:hover { border-color: var(--accent); color: var(--accent); } .ibtn-sm { padding: 4px 10px; font-size: 11px; } /* ── Upload zone ────────────────────────────────────── */ .upload-zone { margin: 10px; border: 2px dashed var(--border); border-radius: 8px; padding: 14px; text-align: center; cursor: pointer; transition: all 0.2s; } .upload-zone:hover { border-color: var(--accent); background: rgba(88,166,255,.05); } /* ── File items ─────────────────────────────────────── */ .fitem { display: flex; align-items: center; gap: 8px; padding: 6px 8px; border-radius: 6px; background: var(--surface2); margin: 4px 10px; font-size: 12px; } .fitem-name { font-family: 'JetBrains Mono', monospace; font-size: 11px; color: var(--text); flex: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } .fitem-meta { font-size: 10px; color: var(--muted); } .fitem-type { font-size: 9px; padding: 1px 5px; border-radius: 3px; background: rgba(88,166,255,.15); color: var(--accent); font-weight: 700; } /* ── Command detail ─────────────────────────────────── */ .detail-box { padding: 14px; overflow-y: auto; flex: 1; } .detail-box::-webkit-scrollbar { width: 3px; } .detail-box::-webkit-scrollbar-thumb { background: var(--surface3); border-radius: 2px; } .detail-cmd { font-family: 'JetBrains Mono', monospace; font-size: 18px; font-weight: 700; color: var(--keyword); } .detail-cat { font-size: 10px; color: var(--muted); text-transform: uppercase; letter-spacing: .1em; margin: 4px 0 10px; } .detail-desc { font-size: 12px; line-height: 1.6; color: var(--text); margin-bottom: 12px; } .sect-title { font-size: 9px; font-weight: 700; letter-spacing: .1em; text-transform: uppercase; color: var(--muted); padding-bottom: 4px; border-bottom: 1px solid var(--border); margin-bottom: 8px; } .syntax-box { background: var(--bg); border: 1px solid var(--border); border-radius: 6px; padding: 10px 12px; font-family: 'JetBrains Mono', monospace; font-size: 12px; line-height: 1.6; color: var(--text); position: relative; word-break: break-all; } .example-box { background: var(--bg); border: 1px solid var(--border); border-left: 3px solid var(--accent2); border-radius: 4px; padding: 8px 10px; font-family: 'JetBrains Mono', monospace; font-size: 11px; color: var(--text); line-height: 1.7; white-space: pre; overflow-x: auto; } .opt-row { display: flex; gap: 8px; margin-bottom: 6px; align-items: flex-start; } .opt-key { font-family: 'JetBrains Mono', monospace; font-size: 11px; color: var(--option); min-width: 80px; flex-shrink: 0; } .opt-val { font-size: 11px; color: var(--muted); line-height: 1.4; } .insert-btn { width: 100%; margin-top: 6px; padding: 6px; border-radius: 5px; background: rgba(88,166,255,.1); border: 1px solid rgba(88,166,255,.3); color: var(--accent); font-family: 'Syne', sans-serif; font-weight: 600; font-size: 11px; cursor: pointer; transition: all 0.15s; text-align: center; } .insert-btn:hover { background: rgba(88,166,255,.2); } /* ── Chat ───────────────────────────────────────────── */ .chat-area { flex: 1; overflow-y: auto; padding: 14px; display: flex; flex-direction: column; gap: 14px; } .chat-area::-webkit-scrollbar { width: 3px; } .chat-area::-webkit-scrollbar-thumb { background: var(--surface3); border-radius: 2px; } .chat-msg { display: flex; gap: 10px; } .chat-msg.user { flex-direction: row-reverse; } .avatar { width: 28px; height: 28px; border-radius: 50%; display: flex; align-items: center; justify-content: center; font-size: 12px; font-weight: 700; flex-shrink: 0; } .avatar-user { background: var(--accent); color: #000; } .avatar-ai { background: var(--accent2); color: #000; } .bubble { max-width: 85%; padding: 8px 12px; border-radius: 8px; font-size: 13px; line-height: 1.5; } .bubble-user { background: rgba(88,166,255,.15); border: 1px solid rgba(88,166,255,.3); color: var(--text); } .bubble-ai { background: var(--surface2); border: 1px solid var(--border); color: var(--text); } .chat-input-bar { background: var(--surface); border-top: 1px solid var(--border); padding: 10px 14px; display: flex; gap: 8px; flex-shrink: 0; } .chat-input { flex: 1; background: var(--bg); border: 1px solid var(--border); border-radius: 8px; padding: 8px 12px; color: var(--text); font-family: 'Syne', sans-serif; font-size: 13px; outline: none; resize: none; } .chat-input:focus { border-color: var(--accent); } /* ── Quick prompt chips ─────────────────────────────── */ .chip-row { display: flex; flex-wrap: wrap; gap: 6px; padding: 8px 14px; border-bottom: 1px solid var(--border); } .chip { padding: 4px 10px; border-radius: 20px; font-size: 11px; font-weight: 600; cursor: pointer; border: 1px solid var(--border); color: var(--muted); background: transparent; transition: all 0.15s; font-family: 'Syne', sans-serif; } .chip:hover { border-color: var(--accent); color: var(--accent); } /* ── Tool call accordion ────────────────────────────── */ .tool-call { margin-top: 6px; background: var(--bg); border: 1px solid var(--border); border-radius: 6px; overflow: hidden; font-size: 11px; } .tool-call-hdr { padding: 5px 10px; background: var(--surface3); color: var(--muted); font-family: 'JetBrains Mono', monospace; cursor: pointer; display: flex; align-items: center; gap: 6px; } .tool-call-body { padding: 8px 10px; font-family: 'JetBrains Mono', monospace; font-size: 11px; line-height: 1.6; color: var(--text); max-height: 200px; overflow-y: auto; white-space: pre-wrap; } .t-success { color: var(--accent2); } .t-warn { color: var(--accent4); } .t-error { color: var(--accent3); } .t-info { color: var(--muted); } .t-data { color: var(--accent); } .t-prompt { color: var(--accent2); } /* ── Streamlit widget overrides ─────────────────────── */ div[data-testid="stTextInput"] label, div[data-testid="stTextArea"] label, div[data-testid="stSelectbox"] label, div[data-testid="stCheckbox"] label, .stRadio label { color: var(--muted) !important; font-size: 11px !important; font-family: 'Syne', sans-serif !important; } div[data-testid="stTextInput"] input, div[data-testid="stNumberInput"] input { background: var(--bg) !important; border: 1px solid var(--border) !important; border-radius: 6px !important; color: var(--text) !important; font-family: 'JetBrains Mono', monospace !important; font-size: 13px !important; } div[data-testid="stTextInput"] input:focus, div[data-testid="stNumberInput"] input:focus { border-color: var(--accent) !important; box-shadow: none !important; } div[data-testid="stTextArea"] textarea { background: var(--bg) !important; border: 1px solid var(--border) !important; color: var(--text) !important; font-family: 'JetBrains Mono', monospace !important; font-size: 13px !important; line-height: 22px !important; } .stButton > button { background: var(--surface2) !important; border: 1px solid var(--border) !important; color: var(--text) !important; font-family: 'Syne', sans-serif !important; font-weight: 600 !important; border-radius: 6px !important; transition: all 0.15s !important; } .stButton > button:hover { border-color: var(--accent) !important; color: var(--accent) !important; } .stButton > button[kind="primary"] { background: var(--accent2) !important; border-color: var(--accent2) !important; color: #000 !important; } .stButton > button[kind="primary"]:hover { background: #56d364 !important; } div[data-testid="stSelectbox"] > div > div { background: var(--bg) !important; border: 1px solid var(--border) !important; color: var(--text) !important; font-family: 'JetBrains Mono', monospace !important; } .stExpander { background: var(--surface2) !important; border: 1px solid var(--border) !important; border-radius: 6px !important; } .stExpander header { color: var(--text) !important; font-family: 'JetBrains Mono', monospace !important; font-size: 13px !important; } .stExpander div[data-testid="stExpanderDetails"] { background: var(--bg) !important; } .stAlert { border-radius: 6px !important; font-family: 'Syne', sans-serif !important; } div[data-testid="stChatMessage"] { background: var(--surface2) !important; border: 1px solid var(--border) !important; border-radius: 8px !important; } div[data-testid="stChatInputContainer"] textarea { background: var(--bg) !important; border: 1px solid var(--border) !important; color: var(--text) !important; font-family: 'Syne', sans-serif !important; } [data-testid="column"] { padding: 0 !important; } /* ── Dataframe ──────────────────────────────────────── */ .stDataFrame { border: 1px solid var(--border) !important; border-radius: 6px !important; } /* ── Code block ─────────────────────────────────────── */ .stCode { background: var(--bg) !important; border: 1px solid var(--border) !important; } code { font-family: 'JetBrains Mono', monospace !important; color: var(--text) !important; } /* ── Plotly chart ───────────────────────────────────── */ .js-plotly-plot { border: 1px solid var(--border) !important; border-radius: 6px !important; } /* ── Scroll bars global ─────────────────────────────── */ ::-webkit-scrollbar { width: 4px; height: 4px; } ::-webkit-scrollbar-track { background: transparent; } ::-webkit-scrollbar-thumb { background: var(--surface3); border-radius: 2px; } /* ── Remove streamlit container gaps ────────────────── */ .element-container { margin: 0 !important; } div[data-testid="stVerticalBlock"] > div { gap: 0 !important; } """ _b64css = _b64.b64encode(_CSS.encode()).decode() st.markdown( '' f'', unsafe_allow_html=True, ) # ───────────────────────────────────────────────────────────────────────────── # SESSION STATE # ───────────────────────────────────────────────────────────────────────────── def _init(): defaults = { "parm_path": None, "traj_paths": [], "work_dir": tempfile.mkdtemp(prefix="cpptraj_"), "chat_history": [], "script": ( "# cpptraj analysis script\n" "parm topology.prmtop\n" "trajin trajectory.nc\n\n" "autoimage\n" "center !:WAT origin\n\n" "rmsd backbone @CA,C,N,O first out rmsd.dat\n\n" "go\n" ), "last_result": None, "api_key": os.environ.get("ANTHROPIC_API_KEY", ""), "runner": None, "agent": None, "center_mode": "editor", # editor | agent | builder | results "sel_cmd": None, # selected command key in left panel "cmd_filter": "all", "doc_search": "", } for k, v in defaults.items(): if k not in st.session_state: st.session_state[k] = v _init() # ───────────────────────────────────────────────────────────────────────────── # RESOURCE HELPERS # ───────────────────────────────────────────────────────────────────────────── @st.cache_resource def get_kb() -> CPPTrajKnowledgeBase: return CPPTrajKnowledgeBase() def get_runner() -> CPPTrajRunner: if st.session_state.runner is None: st.session_state.runner = CPPTrajRunner(work_dir=st.session_state.work_dir) return st.session_state.runner def get_agent() -> TrajectoryAgent: if st.session_state.agent is None: st.session_state.agent = TrajectoryAgent( runner=get_runner(), kb=get_kb(), api_key=st.session_state.api_key, ) return st.session_state.agent # ───────────────────────────────────────────────────────────────────────────── # HELPER: PLOT # ───────────────────────────────────────────────────────────────────────────── def _col_names(fname: str, ncols: int) -> list: maps = { "rmsd": ["Frame"] + ["RMSD_Å"] * (ncols - 1), "rmsf": ["Residue"] + ["RMSF_Å"] * (ncols - 1), "rg": ["Frame", "Rg_Å", "Rg_max_Å"], "radgyr": ["Frame", "Rg_Å", "Rg_max_Å"], "hbond": ["Frame", "N_HBonds"], "distance": ["Frame", "Dist_Å"], "angle": ["Frame", "Angle_°"], "dihedral": ["Frame", "Dihedral_°"], "msd": ["Time_ps", "MSD_Ų", "Dx", "Dy", "Dz"], "density": ["Pos_Å", "Density"], "surf": ["Frame", "SASA_Ų"], "sasa": ["Frame", "SASA_Ų"], "cluster": ["Frame", "Cluster_ID"], "pca_proj": ["Frame"] + [f"PC{i}" for i in range(1, ncols)], "watershell": ["Frame", "Shell1", "Shell2"], "nativecontacts": ["Frame", "Q_native"], } for k, cols in maps.items(): if k in fname: r = list(cols[:ncols]) while len(r) < ncols: r.append(f"col{len(r)}") return r return [f"col{i}" for i in range(ncols)] def plot_file(fp: Path, content: str, key_pfx: str = ""): from io import StringIO rows = [l for l in content.splitlines() if l.strip() and not l.strip().startswith(("#","@","$","%"))] if not rows: st.info("File is empty or comment-only.") return try: df = pd.read_csv(StringIO("\n".join(rows)), sep=r"\s+", header=None, on_bad_lines="skip") if df.empty or df.shape[1] < 2: st.info("Need ≥2 numeric columns to plot.") return names = _col_names(fp.stem.lower(), df.shape[1]) df.columns = names[:df.shape[1]] df = df.apply(pd.to_numeric, errors="coerce").dropna() if df.empty: return x_col = df.columns[0] y_cols = list(df.columns[1:]) ptype = st.radio("Plot type", ["Line","Scatter","Histogram","Box"], horizontal=True, key=f"pt_{key_pfx}_{fp.name}") sel_y = st.multiselect("Y-axis", y_cols, default=y_cols[:min(3,len(y_cols))], key=f"py_{key_pfx}_{fp.name}") if not sel_y: return if ptype == "Line": fig = px.line(df, x=x_col, y=sel_y, title=fp.stem, template="plotly_dark") elif ptype == "Scatter": cols = sel_y if len(sel_y) >= 2 else [x_col] + sel_y fig = px.scatter(df, x=cols[0], y=cols[1], title=fp.stem, template="plotly_dark", opacity=0.6) elif ptype == "Histogram": fig = px.histogram(df, x=sel_y[0], nbins=60, title=fp.stem, template="plotly_dark") else: fig = px.box(df, y=sel_y, title=fp.stem, template="plotly_dark") fig.update_layout( height=360, paper_bgcolor="#0d1117", plot_bgcolor="#0d1117", font_color="#e6edf3", xaxis=dict(gridcolor="#30363d"), yaxis=dict(gridcolor="#30363d"), ) st.plotly_chart(fig, use_container_width=True) with st.expander("Statistics"): st.dataframe(df[sel_y].describe(), use_container_width=True) except Exception as e: st.caption(f"Could not auto-plot: {e}") # ───────────────────────────────────────────────────────────────────────────── # HEADER # ───────────────────────────────────────────────────────────────────────────── runner = get_runner() cpptraj_ok = runner.is_cpptraj_available() parm_ok = st.session_state.parm_path is not None traj_ok = len(st.session_state.traj_paths) > 0 st.markdown(f"""