监控Streamlit中每段代码的rerun情况

把代码放到不同的函数中并用下面的装饰器包裹即可在log中输出rerun次数。

依赖：

• loguru

import streamlit as st
from functools import wraps
from typing import Any, Callable

from loguru._logger import Logger

def run_counter(_func: Callable | None = None, *, logger: Logger | None = None):
    """
    用法：
        @run_counter               # 无参
        @run_counter()             # 无参，显式括号
        @run_counter(logger=logger) # 带参
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            key = f"run_{func.__name__}"
            st.session_state.setdefault(key, 0)
            st.session_state[key] += 1
            if logger is not None:
                logger.debug(f"Rerun {func.__name__} ({st.session_state[key]})")
            return func(*args, **kwargs)

        return wrapper

    # 1. 无参调用：@run_counter 或 @run_counter()
    if _func is None:
        return decorator
    # 2. 带参调用：@run_counter(logger=logger)
    return decorator(_func)

使用方法：

当成正常装饰器使用即可，需要放在streamlit的cache装饰器内部

@st.cache_resource
@run_counter(logger=log)
def get_model():
    model = ChatOpenAI(
        base_url=st.secrets["ollama_base_url"],
        model="qwen3:0.6b",
        api_key=st.secrets["ollama_api_key"],
    )
    return model

@run_counter(logger=log)
def init_app():
    # 初始 session_state
    init_session_state("messages", [])
    init_session_state("context", [])
    init_session_state("chat", [])
    init_session_state("generating", False)

    st.set_page_config(layout="wide", initial_sidebar_state="expanded")