Task2/LLM.py

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"
        )
        # print("##1")
        completion = client.chat.completions.create(
            model = model_name,
            messages=[
                {
                    "role": "user",
                    "content": content
                }
            ]
        )
        # print("##2")
        return completion.choices[0].message
        
    elif model_series == 'llama':
        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':

        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

        
    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)
        print(response.json())
        content = response.json()['choices'][0]['message']
        return content
    
    else:
        return "不支持的模型系列"