umbc-scify/ml-intern-trace

1import json2from pathlib import Path3 4import streamlit as st5 6TRACE_PATH = Path(__file__).parent / "trace.jsonl"7 8 9@st.cache_data10def load_trace(path: str):11    records = []12    with open(path) as f:13        for line in f:14            line = line.strip()15            if line:16                records.append(json.loads(line))17    return records18 19 20def stringify(value):21    if isinstance(value, str):22        return value23    return json.dumps(value, indent=2, ensure_ascii=False)24 25 26def render_user(msg):27    content = msg.get("content")28    if isinstance(content, str):29        with st.chat_message("user"):30            st.markdown(content)31        return32 33    for block in content:34        btype = block.get("type")35        if btype == "tool_result":36            tool_id = block.get("tool_use_id", "")37            with st.chat_message("tool", avatar="🔧"):38                st.caption(f"tool_result · {tool_id}")39                body = block.get("content")40                with st.expander("Result", expanded=False):41                    if isinstance(body, list):42                        for b in body:43                            if isinstance(b, dict) and b.get("type") == "text":44                                st.code(b.get("text", ""), language="markdown")45                            else:46                                st.code(stringify(b), language="json")47                    else:48                        st.code(stringify(body), language="markdown")49        else:50            with st.chat_message("user"):51                st.code(stringify(block), language="json")52 53 54def render_assistant(msg):55    content = msg.get("content")56    model = msg.get("model", "")57    if isinstance(content, str):58        with st.chat_message("assistant"):59            if model:60                st.caption(model)61            st.markdown(content)62        return63 64    for block in content:65        btype = block.get("type")66        if btype == "text":67            with st.chat_message("assistant"):68                if model:69                    st.caption(model)70                st.markdown(block.get("text", ""))71        elif btype == "tool_use":72            name = block.get("name", "tool")73            tool_id = block.get("id", "")74            with st.chat_message("assistant", avatar="🛠️"):75                st.caption(f"{model} · tool_use · {name} · {tool_id}")76                with st.expander(f"{name} input", expanded=False):77                    st.code(stringify(block.get("input", {})), language="json")78        else:79            with st.chat_message("assistant"):80                st.code(stringify(block), language="json")81 82 83def main():84    st.set_page_config(page_title="Trace Viewer", layout="wide")85    st.title("Trace Viewer")86 87    records = load_trace(str(TRACE_PATH))88 89    sessions = sorted({r.get("sessionId", "") for r in records})90    with st.sidebar:91        st.header("Filters")92        session = st.selectbox("Session", ["(all)"] + sessions)93        show_user = st.checkbox("Show user messages", value=True)94        show_assistant = st.checkbox("Show assistant messages", value=True)95        show_tool_results = st.checkbox("Show tool_result blocks", value=True)96        show_tool_use = st.checkbox("Show tool_use blocks", value=True)97        show_timestamps = st.checkbox("Show timestamps", value=False)98        st.caption(f"{len(records)} records in {TRACE_PATH.name}")99 100    filtered = []101    for r in records:102        if session != "(all)" and r.get("sessionId") != session:103            continue104        if r["type"] == "user" and not show_user:105            continue106        if r["type"] == "assistant" and not show_assistant:107            continue108        msg = r.get("message", {})109        content = msg.get("content")110        if isinstance(content, list):111            kept = []112            for b in content:113                bt = b.get("type")114                if bt == "tool_result" and not show_tool_results:115                    continue116                if bt == "tool_use" and not show_tool_use:117                    continue118                kept.append(b)119            if not kept:120                continue121            r = {**r, "message": {**msg, "content": kept}}122        filtered.append(r)123 124    st.caption(f"Showing {len(filtered)} of {len(records)} records")125 126    for r in filtered:127        if show_timestamps:128            st.caption(f"{r.get('timestamp', '')} · {r.get('uuid', '')[:8]}")129        if r["type"] == "user":130            render_user(r.get("message", {}))131        elif r["type"] == "assistant":132            render_assistant(r.get("message", {}))133        else:134            st.code(stringify(r), language="json")135 136 137if __name__ == "__main__":138    main()139