Update app.py

app.py

@@ -2,7 +2,7 @@ import sys, os
 import tempfile
 import streamlit as st
 
-from config import MODEL_CONFIGS
+from config import MODEL_CONFIGS, CUSTOM_MODEL_KEY
 from utils.save_secrets import *
 from prompt_engineer.sec1_call_llm import DataLoadingAgent
 from prompt_engineer.sec2_call_llm import DataPreprocessAgent
@@ -20,12 +20,7 @@ np.set_printoptions(edgeitems=250, threshold=501)
 
 sys.path.append(os.path.dirname(__file__))
 
-CACHE_FILE = os.path.join(tempfile.gettempdir(), "anystat_cache.pkl")
-CACHE_DIR = './cache'
-SECRETS_PATH = Path(".streamlit") / "secrets.toml"
 
-
-# 设置页面配置
 st.set_page_config(
     page_title="Autostat",
     page_icon="🤖",
@@ -37,21 +32,36 @@ def init_session_state():
 
     if 'selected_model' not in st.session_state:
         st.session_state.selected_model = "DeepSeek"
-    if "api_keys" not in st.session_state:
-        st.session_state.api_keys = load_local_api_keys()
+    
+    if 'model_configs_runtime' not in st.session_state:
+        # 运行时模型配置，包含预设和自定义模型
+        st.session_state.model_configs_runtime = MODEL_CONFIGS.copy()
+        # 加载用户配置（包括 API 密钥和自定义模型）
+        user_configs = load_local_model_configs()
+        for model_name, config in user_configs.items():
+            if model_name in MODEL_CONFIGS:
+                # 预设模型：只更新 API 密钥
+                st.session_state.model_configs_runtime[model_name]["api_key"] = config.get("api_key", "")
+            else:
+                # 自定义模型：添加完整配置
+                st.session_state.model_configs_runtime[model_name] = {
+                    "api_base": config.get("api_base", ""),
+                    "model_name": config.get("model_name", ""),
+                    "api_key": config.get("api_key", ""),
+                    "api_type": "openai",
+                    "is_preset": False,
+                }
+    
+    # 从 model_configs_runtime 提取 api_keys（用于传递给 Agent）
+    if 'api_keys' not in st.session_state:
+        st.session_state.api_keys = {
+            name: config.get("api_key", "")
+            for name, config in st.session_state.model_configs_runtime.items()
+        }
+    
     if 'auto_mode' not in st.session_state:
         st.session_state.auto_mode = False
 
-    if 'loading_start_time' not in st.session_state:
-        st.session_state.loading_start_time = None
-    if 'prep_start_time' not in st.session_state:
-        st.session_state.prep_start_time = None
-    if 'vis_start_time' not in st.session_state:
-        st.session_state.vis_start_time = None
-    if 'modeling_start_time' not in st.session_state:
-        st.session_state.modeling_start_time = None
-    if 'report_start_time' not in st.session_state:
-        st.session_state.report_start_time = None
     if 'preference_select' not in st.session_state:
         st.session_state.preference_select = None
     if 'additional_preference' not in st.session_state:
@@ -62,37 +72,37 @@ def init_session_state():
     if 'data_loading_agent' not in st.session_state:
         st.session_state.data_loading_agent = DataLoadingAgent(
             api_keys=st.session_state.api_keys,
-            model_configs=MODEL_CONFIGS,
+            model_configs=st.session_state.model_configs_runtime,
             model=st.session_state.selected_model
         )
     if 'data_preprocess_agent' not in st.session_state:
         st.session_state.data_preprocess_agent = DataPreprocessAgent(
             api_keys=st.session_state.api_keys,
-            model_configs=MODEL_CONFIGS,
+            model_configs=st.session_state.model_configs_runtime,
             model=st.session_state.selected_model
         )
     if 'visualization_agent' not in st.session_state:
         st.session_state.visualization_agent = VisualizationAgent(
             api_keys=st.session_state.api_keys,
-            model_configs=MODEL_CONFIGS,
+            model_configs=st.session_state.model_configs_runtime,
             model=st.session_state.selected_model
         )
     if 'modeling_coding_agent' not in st.session_state:
         st.session_state.modeling_coding_agent = ModelingCodingAgent(
             api_keys=st.session_state.api_keys,
-            model_configs=MODEL_CONFIGS,
+            model_configs=st.session_state.model_configs_runtime,
             model=st.session_state.selected_model
         )
     if 'report_agent' not in st.session_state:
         st.session_state.report_agent = ReportAgent(
             api_keys=st.session_state.api_keys,
-            model_configs=MODEL_CONFIGS,
+            model_configs=st.session_state.model_configs_runtime,
             model=st.session_state.selected_model
         )
     if 'planner_agent' not in st.session_state:
         st.session_state.planner_agent = PlannerAgent(
             api_keys=st.session_state.api_keys,
-            model_configs=MODEL_CONFIGS,
+            model_configs=st.session_state.model_configs_runtime,
             model=st.session_state.selected_model
         )
 
@@ -111,11 +121,21 @@ def run_app():
     init_session_state()
     with st.sidebar:
         st.subheader("选择大模型")
-        models = list(MODEL_CONFIGS.keys())
+        
+        # 获取所有可用的模型（预设模型 + OpenAI API 兼容模型）
+        models = list(MODEL_CONFIGS.keys()) + [CUSTOM_MODEL_KEY]
+        
+        # 确保选择的索引有效
+        try:
+            current_index = models.index(st.session_state.selected_model)
+        except ValueError:
+            current_index = 0
+            st.session_state.selected_model = models[0]
+        
         st.selectbox(
             "选择要使用的大模型",
             models,
-            index=models.index(st.session_state.selected_model),
+            index=current_index,
             key="model_selector",
             on_change=on_model_selector_change,
         )
@@ -123,50 +143,132 @@ def run_app():
         st.subheader("API 密钥设置")
         selected = st.session_state.selected_model
 
-        api_key_input = st.text_input(
-            f"{selected} API 密钥",
-            value=st.session_state.api_keys.get(selected, ""),
-            type="password",
-            key="api_key_input",
-        )
-        st.session_state.api_keys[selected] = api_key_input
+        # 判断是否为 OpenAI API 兼容模型
+        is_custom_model = (selected == CUSTOM_MODEL_KEY)
+        
+        if is_custom_model:
+            # 显示 OpenAI API 兼容模型的配置界面
+            existing_config = st.session_state.model_configs_runtime.get(CUSTOM_MODEL_KEY, {})
+            
+            base_url_input = st.text_input(
+                "Base URL",
+                value=existing_config.get("api_base", ""),
+                key="base_url_input",
+                placeholder="例如: https://api.siliconflow.cn/v1"
+            )
+            
+            model_name_input = st.text_input(
+                "模型 ID",
+                value=existing_config.get("model_name", ""),
+                key="model_name_input",
+                placeholder="例如: Qwen/Qwen3-8B"
+            )
+            
+            api_key_input = st.text_input(
+                "API 密钥",
+                value=st.session_state.api_keys.get(CUSTOM_MODEL_KEY, ""),
+                type="password",
+                key="api_key_input",
+            )
+            
+            if existing_config and existing_config.get("api_base"):
+                st.info(f"当前配置: {existing_config.get('model_name', 'N/A')}")
+            else:
+                st.info("配置 OpenAI API 兼容模型")
+            
+            if st.button("💾 保存配置", use_container_width=True, key="save_key"):
+                if not base_url_input or not model_name_input or not api_key_input:
+                    st.error("请填写所有必需字段")
+                else:
+                    # 保存到配置文件
+                    update_local_model_config(
+                        display_name=CUSTOM_MODEL_KEY,
+                        api_key=api_key_input,
+                        base_url=base_url_input,
+                        model_name=model_name_input
+                    )
+                    
+                    # 更新运行时配置
+                    st.session_state.model_configs_runtime[CUSTOM_MODEL_KEY] = {
+                        "api_base": base_url_input,
+                        "model_name": model_name_input,
+                        "api_key": api_key_input,  # 也保存 api_key
+                        "api_type": "openai",
+                        "is_preset": False,
+                    }
+                    # 同步到 api_keys
+                    st.session_state.api_keys[CUSTOM_MODEL_KEY] = api_key_input
+                    st.session_state.selected_model = CUSTOM_MODEL_KEY
+                    
+                    st.success("已保存配置")
+                    st.rerun()
+        else:
+            # 预设模型或已保存的自定义模型
+            api_key_input = st.text_input(
+                f"{selected} API 密钥",
+                value=st.session_state.api_keys.get(selected, ""),
+                type="password",
+                key="api_key_input",
+            )
+            
+            # 如果是自定义模型，显示其配置信息
+            if selected in st.session_state.model_configs_runtime:
+                config = st.session_state.model_configs_runtime[selected]
+                if not config.get("is_preset", False):
+                    st.caption(f"Base URL: {config.get('api_base', 'N/A')}")
+                    st.caption(f"Model: {config.get('model_name', 'N/A')}")
 
-        if st.button("💾 保存密钥", use_container_width=True, key="save_key"):
+            if st.button("💾 保存密钥", use_container_width=True, key="save_key"):
+                # 保存到配置文件
+                config = st.session_state.model_configs_runtime.get(selected, {})
+                if config.get("is_preset", False):
+                    # 预设模型，只保存 API key
+                    update_local_model_config(display_name=selected, api_key=api_key_input)
+                else:
+                    # 自定义模型，保存完整配置
+                    update_local_model_config(
+                        display_name=selected,
+                        api_key=api_key_input,
+                        base_url=config.get("api_base"),
+                        model_name=config.get("model_name")
+                    )
 
-            st.session_state.api_keys[selected] = api_key_input
-            st.success("已保存")
-            st.rerun()
+                # 同步更新运行时配置和 api_keys
+                st.session_state.model_configs_runtime[selected]["api_key"] = api_key_input
+                st.session_state.api_keys[selected] = api_key_input
+                st.success("已保存")
+                st.rerun()
 
         if st.button("🧹 清空数据", use_container_width=True, key="clear_data"):
 
             st.session_state.data_loading_agent = DataLoadingAgent(
                 api_keys=st.session_state.api_keys,
-                model_configs=MODEL_CONFIGS,
+                model_configs=st.session_state.model_configs_runtime,
                 model=st.session_state.selected_model
             )
             st.session_state.data_preprocess_agent = DataPreprocessAgent(
                 api_keys=st.session_state.api_keys,
-                model_configs=MODEL_CONFIGS,
+                model_configs=st.session_state.model_configs_runtime,
                 model=st.session_state.selected_model
             )
             st.session_state.visualization_agent = VisualizationAgent(
                 api_keys=st.session_state.api_keys,
-                model_configs=MODEL_CONFIGS,
+                model_configs=st.session_state.model_configs_runtime,
                 model=st.session_state.selected_model
             )
             st.session_state.modeling_coding_agent = ModelingCodingAgent(
                 api_keys=st.session_state.api_keys,
-                model_configs=MODEL_CONFIGS,
+                model_configs=st.session_state.model_configs_runtime,
                 model=st.session_state.selected_model
             )
             st.session_state.report_agent = ReportAgent(
                 api_keys=st.session_state.api_keys,
-                model_configs=MODEL_CONFIGS,
+                model_configs=st.session_state.model_configs_runtime,
                 model=st.session_state.selected_model
             )
             st.session_state.planner_agent = PlannerAgent(
                 api_keys=st.session_state.api_keys,
-                model_configs=MODEL_CONFIGS,
+                model_configs=st.session_state.model_configs_runtime,
                 model=st.session_state.selected_model
             )
             st.session_state.auto_mode = False
@@ -186,7 +288,7 @@ def run_app():
                     st.session_state.auto_mode = False
                     st.session_state.planner_agent = PlannerAgent(
                     api_keys=st.session_state.api_keys,
-                    model_configs=MODEL_CONFIGS,
+                    model_configs=st.session_state.model_configs_runtime,
                     model=st.session_state.selected_model
                     )
                     st.rerun()
@@ -231,5 +333,4 @@ def run_app():
     pg.run()
     
 if __name__ == "__main__":
-    run_app()
-
+    run_app()