import os
import json
from openai import OpenAI
from typing import Sequence

os.environ["OPENAI_BASE_URL"] = "http://192.168.0.1:8000/v1" # 替换为本地主机 IP
os.environ["OPENAI_API_KEY"] = "0"

def calculate_gpa(grades: Sequence[str], hours: Sequence[int]) -> float:
    grade_to_score = {"A": 4, "B": 3, "C": 2}
    total_score, total_hour = 0, 0
    for grade, hour in zip(grades, hours):
        total_score += grade_to_score[grade] * hour
        total_hour += hour
    return total_score / total_hour

tool_map = {
    "calculate_gpa": calculate_gpa
}

if __name__ == "__main__":
    client = OpenAI()
    tools = [
        {
            "type": "function",
            "function": {
                "name": "calculate_gpa",
                "description": "根据课程成绩和学时计算 GPA",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "grades": {"type": "array", "items": {"type": "string"}, "description": "课程成绩"},
                        "hours": {"type": "array", "items": {"type": "integer"}, "description": "课程学时"},
                    },
                    "required": ["grades", "hours"],
                },
            },
        }
    ]
    messages = []
    messages.append({"role": "user", "content": "我的成绩是 A, A, B, C，学时是 3, 4, 3, 2"})
    result = client.chat.completions.create(messages=messages, model="yi-agent-6b", tools=tools)
    tool_call = result.choices[0].message.tool_calls[0].function
    name, arguments = tool_call.name, json.loads(tool_call.arguments)
    messages.append({"role": "function", "content": json.dumps({"name": name, "argument": arguments}, ensure_ascii=False)})
    tool_result = tool_map[name](**arguments)
    messages.append({"role": "tool", "content": json.dumps({"gpa": tool_result}, ensure_ascii=False)})
    result = client.chat.completions.create(messages=messages, model="yi-agent-6b", tools=tools)
    print(result.choices[0].message.content)
    # 根据你的成绩和学时，你的平均绩点 (GPA) 为 3.4166666666666665。