Update prompt_engineer/call_llm.py

prompt_engineer/call_llm.py

@@ -1,15 +1,12 @@
-import re
 from openai import OpenAI, OpenAIError
 from anthropic import Anthropic, AnthropicError
-import requests
-import json
+import zai
+from zai import ZhipuAiClient
 
 import streamlit as st
 import pandas as pd
-import numpy as np
-from config import MODEL_CONFIGS
-from typing import IO, List, Dict
-from zai import ZhipuAiClient
+from typing import List, Dict
+
 
 class LLMClient:
     def __init__(self, model_configs: dict, api_keys: dict, model: str):
@@ -34,77 +31,104 @@ class LLMClient:
         )
 
         try:
-            if model_name == "GPT-4o" or model_name == "GPT-5" or model_name == "DeepSeek" or model_name == "通义千问" or model_name == "Claude" or model_name == "豆包":
+            # 获取 API 类型
+            api_type = config.get("api_type", "openai")
+            
+            # 根据 API 类型选择不同的调用方式
+            if api_type == "openai":
+                # OpenAI 兼容的 API（包括 OpenAI、DeepSeek、通义千问、豆包等）
                 try:
                     client = OpenAI(
                         api_key=api_key,
-                        base_url=config["api_base"]
+                        base_url=config.get("api_base", "https://api.openai.com/v1")
                     )
                     
-                    # 使用新的 API 调用方式
-                    resp = client.chat.completions.create(
-                        model=config["model_name"],
+                    response = client.chat.completions.create(
+                        model=config.get("model_name", "gpt-4o"),
                         messages=[
                             {"role": "system", "content": system_msg},
                             {"role": "user", "content": prompt},
                         ],
-                        stream = False
+                        stream=False
                     )
-                    return resp.choices[0].message.content
+                    return response.choices[0].message.content
                 
                 except OpenAIError as e:
-                    # 这里可以捕获所有OpenAI SDK定义的错误
-                    st.error(f"API调用失败: {str(e)}")
-                    # 记录日志或提示用户
-                    return "调用失败，请检查密钥或网络"
+                    msg = f"API 调用失败：{str(e)}"
+                    st.error(msg)
+                    return f"调用失败，请检查密钥或网络（{msg}）"
                 except Exception as e:
-                    # 捕获其他非预期的异常，如网络问题
-                    st.error(f"发生未知错误: {str(e)}")
-                    return "发生未知错误"
-
-            elif model_name == "智谱AI":
-                client = ZhipuAiClient(api_key=api_key)
-                response = client.chat.completions.create(
-                    model=config["model_name"],
-                    messages=[{"role": "system", "content": "你是一个专业的数据分析助手。"},
-                        {"role": "user", "content": prompt}],
-                    thinking={
-                        "type":"enabled"
-                    }
-                )
-                if response:
-                    print(response.choices[0].message)
-                    desc = response.choices[0].message.content if hasattr(response.choices[0].message, "content") else str(response.choices[0].message)
-                    return desc.replace("<|begin_of_box|>", "").replace("<|end_of_box|>", "").strip()
+                    msg = f"发生未知错误：{str(e)}"
+                    st.error(msg)
+                    return msg
+            
+            elif api_type == "claude":
+                # Claude 使用 Anthropic SDK
+                try:
+                    client = Anthropic(api_key=api_key)
+                    
+                    response = client.messages.create(
+                        model=config.get("model_name", "claude-3-5-sonnet-latest"),
+                        max_tokens=4096,
+                        system=system_msg,
+                        messages=[
+                            {"role": "user", "content": prompt}
+                        ]
+                    )
+                    
+                    return response.content[0].text
+                
+                except AnthropicError as e:
+                    msg = f"Claude API 调用失败：{str(e)}"
+                    st.error(msg)
+                    return f"调用失败，请检查密钥或网络（{msg}）"
+                except Exception as e:
+                    msg = f"发生未知错误：{str(e)}"
+                    st.error(msg)
+                    return msg
 
-                st.error(f"智谱调用失败：{response.text}")
-                return "调用失败，请检查密钥或网络"
+            elif api_type == "zhipu":
+                # 智谱 AI 使用自己的客户端
+                # 参考 https://github.com/zai-org/z-ai-sdk-python
+                try:
+                    client = ZhipuAiClient(
+                        api_key=api_key,
+                        base_url=config.get("api_base", "https://open.bigmodel.cn/api/paas/v4")
+                    )
+                    
+                    response = client.chat.completions.create(
+                        model=config.get("model_name", "glm-4v-plus-0111"),
+                        messages=[
+                            {"role": "system", "content": system_msg},
+                            {"role": "user", "content": prompt}
+                        ],
+                        thinking={"type": "enabled"}
+                    )
+                    desc = response.choices[0].message.content if hasattr(
+                        response.choices[0].message, "content"
+                    ) else str(response.choices[0].message)
+                    return desc.replace("<|begin_of_box|>", "").replace("<|end_of_box|>", "").strip()
 
-                
-            # elif model_name == "DeepSeek":
-            #     client = OpenAI(
-            #         api_key=api_key,
-            #         base_url=config["api_base"])
-
-            #     resp = client.chat.completions.create(
-            #         model=config["model_name"],
-            #         messages=[
-            #             {"role": "system", "content": system_msg},
-            #             {"role": "user", "content": prompt},
-            #         ],
-            #         stream=False
-            #     )
-            #     if resp:
-            #         return resp.choices[0].message.content
-            #     st.error(f"DeepSeek调用失败：{resp.text}")
-            #     return "调用失败，请检查密钥或网络"
+                except zai.core.APIStatusError as e:
+                    msg = f"智谱 API 状态错误：{str(e)}"
+                    st.error(msg)
+                    return f"调用失败，请检查密钥或网络（{msg}）"
+                except zai.core.APITimeoutError as e:
+                    msg = f"智谱 API 请求超时：{str(e)}"
+                    st.error(msg)
+                    return f"调用失败，请检查密钥或网络（{msg}）"
+                except Exception as e:
+                    msg = f"发生未知错误：{str(e)}"
+                    st.error(msg)
+                    return msg
 
             else:
-                return f"暂不支持模型：{model_name}"
+                return f"不支持的 API 类型：{api_type}"
 
         except Exception as e:
-            st.error(f"{model_name} 调用异常：{e}")
-            return "大模型调用失败，请检查 API 密钥或网络连接"
+            msg = f"{model_name} 调用异常：{e}"
+            st.error(msg)
+            return f"大模型调用失败，请检查 API 密钥或网络连接（{msg}）"
 
     
     def add_memory(self, entry: Dict[str, str]) -> None:
@@ -124,13 +148,11 @@ class LLMClient:
 
     def add_df(self, input_df) -> None:
 
- 
-        
         self.df = input_df
         
 
     def load_df(self) -> pd.DataFrame:
-        
+
         return self.df
     
 
@@ -141,4 +163,4 @@ class LLMClient:
 
     def has_df(self) -> bool:
 
-        return self.df == None
+        return self.df is None