Update LLM.py

LLM.py

@@ -1,100 +1,87 @@
-from openai import OpenAI
-import os
-import pandas as pd
-import google.generativeai as genai
-import requests
-
-# 输入模型系列和内容
-def zero_shot(model_series, model_name, content):
-    # 读取txt文件
-    df = pd.read_csv('./LLM_APIs.txt')
-    # 由于列名包含制表符,需要先分割列名
-    df = pd.DataFrame([x.split('\t') for x in df.values.flatten()], columns=['name','series','API'])
-    # 获取对应的API密钥
-    api_key = df[df['name'] == model_series]['API'].values[0]
-    
-    # 根据不同模型系列调用不同API
-    if model_series == 'gpt':
-        client = OpenAI(
-            api_key=api_key,
-            base_url="https://api.bianxie.ai/v1"
-        )
-        completion = client.chat.completions.create(
-            model = model_name,
-            messages=[
-                {
-                    "role": "user",
-                    "content": content
-                }
-            ]
-        )
-        return completion.choices[0].message
-   
-
-    elif model_series == 'deepSeek' or model_series == 'qwen' or model_series == 'llama':
-
-        client = OpenAI(api_key=api_key,
-                         base_url="https://api.studio.nebius.ai/v1")
-
-        response = client.chat.completions.create(
-            model=model_name,
-            messages=[
-                {"role": "user", "content": content},
-            ],
-            stream=False
-        )
-        # client = OpenAI(api_key=api_key,
-        #                  base_url="https://api.deepseek.com")
-
-        # response = client.chat.completions.create(
-        #     model=model_name,
-        #     messages=[
-        #         {"role": "user", "content": content},
-        #     ],
-        #     stream=False
-        # )
-
-        return response.choices[0].message.content
-
-        
-    elif model_series == 'gemini':
-        # genai.configure(api_key=api_key)
-        # model = genai.GenerativeModel(model_name)
-        # response = model.generate_content(content)
-        # print(response.text)
-        # return response.text
-        api_key = api_key
-        url = 'https://api.bianxie.ai/v1/chat/completions'
-        headers = {
-            'Content-Type': 'application/json',
-            'Authorization': f'Bearer {api_key}'
-        }
-        data = {
-            'model': model_name,
-            'messages': [{'role': 'user', 'content': content}],
-        }
-        response = requests.post(url, headers=headers, json=data)
-        content = response.json()['choices'][0]['message']
-        return content
-        
-    elif model_series == 'claude':
-        api_key = api_key
-        url = 'https://api.bianxie.ai/v1/chat/completions'
-
-        headers = {
-            'Content-Type': 'application/json',
-            'Authorization': f'Bearer {api_key}'
-        }
-
-        data = {
-            'model': model_name,
-            'messages': [{'role': 'user', 'content': content}],
-        }
-
-        response = requests.post(url, headers=headers, json=data)
-        content = response.json()['choices'][0]['message']
-        return content
-    
-    else:
-        return "不支持的模型系列"
-
+from openai import OpenAI
+import os
+import pandas as pd
+import google.generativeai as genai
+import requests
+
+# 输入模型系列和内容
+def zero_shot(model_series, model_name, content):
+    # 根据模型系列获取对应的环境变量
+    if model_series in ['gpt', 'gemini', 'claude']:
+        api_key = os.environ.get("OPENAI_API_KEY")
+    elif model_series in ['deepSeek', 'qwen', 'llama']:
+        api_key = os.environ.get("DEEPSEEK_API_KEY")
+    else:
+        return "不支持的模型系列"
+    
+    if not api_key:
+        return f"未找到 {model_series} 对应的API密钥环境变量"
+    
+    # 根据不同模型系列调用不同API
+    if model_series == 'gpt':
+        client = OpenAI(
+            api_key=api_key,
+            base_url="https://api.bianxie.ai/v1"
+        )
+        completion = client.chat.completions.create(
+            model = model_name,
+            messages=[
+                {
+                    "role": "user",
+                    "content": content
+                }
+            ]
+        )
+        return completion.choices[0].message
+   
+
+    elif model_series == 'deepSeek' or model_series == 'qwen' or model_series == 'llama':
+
+        client = OpenAI(api_key=api_key,
+                         base_url="https://api.studio.nebius.ai/v1")
+
+        response = client.chat.completions.create(
+            model=model_name,
+            messages=[
+                {"role": "user", "content": content},
+            ],
+            stream=False
+        )
+
+        return response.choices[0].message.content
+
+        
+    elif model_series == 'gemini':
+        api_key = api_key
+        url = 'https://api.bianxie.ai/v1/chat/completions'
+        headers = {
+            'Content-Type': 'application/json',
+            'Authorization': f'Bearer {api_key}'
+        }
+        data = {
+            'model': model_name,
+            'messages': [{'role': 'user', 'content': content}],
+        }
+        response = requests.post(url, headers=headers, json=data)
+        content = response.json()['choices'][0]['message']
+        return content
+        
+    elif model_series == 'claude':
+        url = 'https://api.bianxie.ai/v1/chat/completions'
+
+        headers = {
+            'Content-Type': 'application/json',
+            'Authorization': f'Bearer {api_key}'
+        }
+
+        data = {
+            'model': model_name,
+            'messages': [{'role': 'user', 'content': content}],
+        }
+
+        response = requests.post(url, headers=headers, json=data)
+        content = response.json()['choices'][0]['message']
+        return content
+    
+    else:
+        return "不支持的模型系列"