# 导入相应的库
import os
import random
import time
import datetime

import openai

from tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential, # 为指数回退 (for exponential backoff)
)  
import wandb
from wandb.sdk.data_types.trace_tree import Trace

# 获取 OpenAI 的 API
import openai

# 您的OpenAI-API
openai.api_key = 'sk-XXX'

# 定义相应变量
PROJECT = "dlai_llm"
MODEL_NAME = "gpt-3.5-turbo"

# 登陆wandb
wandb.login(anonymous="allow")

# 初始化
run = wandb.init(project=PROJECT, job_type="generation")

@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def completion_with_backoff(**kwargs):
    "这个函数能避免速率限制 (rate limits)"
    
    return openai.ChatCompletion.create(**kwargs)

def generate_and_print(system_prompt, user_prompt, table, n=5):
    messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    start_time = time.time()
    responses = completion_with_backoff(
        model=MODEL_NAME,
        messages=messages,
        n = n,
        )
    elapsed_time = time.time() - start_time
    for response in responses.choices:
        generation = response.message.content
        print(generation)
    table.add_data(system_prompt,
                user_prompt,
                [response.message.content for response in responses.choices],
                elapsed_time,
                datetime.datetime.fromtimestamp(responses.created),
                responses.model,
                responses.usage.prompt_tokens,
                responses.usage.completion_tokens,
                responses.usage.total_tokens
                )

system_prompt = """You are a creative copywriter.
You're given a category of game asset, \
and your goal is to design a name of that asset.
The game is set in a fantasy world \
where everyone laughs and respects each other, 
while celebrating diversity."""

system_prompt_cn = """你是一个有创意的撰稿人。 \
根据给定的游戏资产类别， \
你的目标是为该资产设计一个名称。 \
游戏设定在一个幻想的世界里，每个人都很快乐且互相尊重，同时庆祝多样性。"""

# 定义包含要追踪的列名的 W&B 表格
columns = ["system_prompt", "user_prompt", "generations", "elapsed_time", "timestamp",\
            "model", "prompt_tokens", "completion_tokens", "total_tokens"]
table = wandb.Table(columns=columns)

user_prompt = "hero"
generate_and_print(system_prompt, user_prompt, table)

user_prompt_cn = "英雄"
generate_and_print(system_prompt_cn, user_prompt_cn, table)

user_prompt = "jewel"
generate_and_print(system_prompt, user_prompt, table)

user_prompt_cn = "珠宝"
generate_and_print(system_prompt_cn, user_prompt_cn, table)

wandb.log({"simple_generations": table})
run.finish()

worlds = [
    "a mystic medieval island inhabited by intelligent and funny frogs",
    "a modern castle sitting on top of a volcano in a faraway galaxy",
    "a digital world inhabited by friendly machine learning engineers"
]

worlds_cn = [
    "一个神秘的中世纪岛屿，居住着聪明有趣的青蛙",
    "一座坐落在遥远星系火山顶上的现代城堡",
    "一个由友好的机器学习工程师居住的数字世界"
]

# 定义配置
model_name = "gpt-3.5-turbo"
temperature = 0.7
system_message = """You are a creative copywriter. 
You're given a category of game asset and a fantasy world.
Your goal is to design a name of that asset.
Provide the resulting name only, no additional description.
Single name, max 3 words output, remember!"""
system_message_cn = """你是一个有创意的撰稿人。给你一个游戏资产类别和一个幻想世界。您的目标是为该资产设计一个名称。只提供结果名称，不提供其他描述。单一名称，最多3个字输出，记住!"""

def run_creative_chain(query):
    # 部分1 - 链开始了
    start_time_ms = round(datetime.datetime.now().timestamp() * 1000)

    # 我们会从顶层的 root span 开始
    root_span = Trace(
          name="MyCreativeChain",
          kind="chain",
          start_time_ms=start_time_ms,
          metadata={"user": "student_1"},
          model_dict={"_kind": "CreativeChain"}
          )

    # 部分2 - 链随机选择一个虚幻世界
    time.sleep(3)
    world = random.choice(worlds)
    expanded_prompt = f'Game asset category: {query}; fantasy world description: {world}'
    tool_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)

    # 创建一个tool span
    tool_span = Trace(
          name="WorldPicker",
          kind="tool",
          status_code="success",
          start_time_ms=start_time_ms,
          end_time_ms=tool_end_time_ms,
          inputs={"input": query},
          outputs={"result": expanded_prompt},
          model_dict={"_kind": "tool", "num_worlds": len(worlds)}
          )

    # 将上面的 tool span 添加为 root span 的子
    root_span.add_child(tool_span)

    # 部分3 - 将工具 span 的输出传递至大语言模型链
    messages=[
      {"role": "system", "content": system_message},
      {"role": "user", "content": expanded_prompt}
    ]

    response = completion_with_backoff(model=model_name,
                                       messages=messages,
                                       max_tokens=12,
                                       temperature=temperature)   

    llm_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
    response_text = response["choices"][0]["message"]["content"]
    token_usage = response["usage"].to_dict()

    llm_span = Trace(
          name="OpenAI",
          kind="llm",
          status_code="success",
          metadata={"temperature":temperature,
                    "token_usage": token_usage, 
                    "model_name":model_name},
          start_time_ms=tool_end_time_ms,
          end_time_ms=llm_end_time_ms,
          inputs={"system_prompt":system_message, "query":expanded_prompt},
          outputs={"response": response_text},
          model_dict={"_kind": "Openai", "engine": response["model"], "model": response["object"]}
          )

    # 将大模型 span 添加为链 span的子
    root_span.add_child(llm_span)

    # 更新链 span 的输入和输出
    root_span.add_inputs_and_outputs(
          inputs={"query":query},
          outputs={"response": response_text})

    # u更新链 span 的结束时间
    root_span.end_time_ms = llm_end_time_ms


    # 部分4 - 通过记录 root span 来记录所有 spans 到 W&B 表
    root_span.log(name="creative_trace")
    print(f"Result: {response_text}")

def run_creative_chain_cn(query):
    # 部分1 - 链开始了
    start_time_ms = round(datetime.datetime.now().timestamp() * 1000)

    # 我们会从顶层的 root span 开始
    root_span = Trace(
          name="MyCreativeChain",
          kind="chain",
          start_time_ms=start_time_ms,
          metadata={"user": "student_1"},
          model_dict={"_kind": "CreativeChain"}
          )

    # 部分2 - 链随机选择一个虚幻世界
    time.sleep(3)
    world = random.choice(worlds_cn)
    expanded_prompt = f'Game asset category: {query}; fantasy world description: {world}'
    tool_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)

    # 创建一个tool span
    tool_span = Trace(
          name="WorldPicker",
          kind="tool",
          status_code="success",
          start_time_ms=start_time_ms,
          end_time_ms=tool_end_time_ms,
          inputs={"input": query},
          outputs={"result": expanded_prompt},
          model_dict={"_kind": "tool", "num_worlds": len(worlds)}
          )

    # 将上面的 tool span 添加为 root span 的子
    root_span.add_child(tool_span)

    # 部分3 - 将工具 span 的输出传递至大语言模型链
    messages=[
      {"role": "system", "content": system_message_cn},
      {"role": "user", "content": expanded_prompt}
    ]

    response = completion_with_backoff(model=model_name,
                                       messages=messages,
                                       max_tokens=12,
                                       temperature=temperature)   

    llm_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
    response_text = response["choices"][0]["message"]["content"]
    token_usage = response["usage"].to_dict()

    llm_span = Trace(
          name="OpenAI",
          kind="llm",
          status_code="success",
          metadata={"temperature":temperature,
                    "token_usage": token_usage, 
                    "model_name":model_name},
          start_time_ms=tool_end_time_ms,
          end_time_ms=llm_end_time_ms,
          inputs={"system_prompt":system_message, "query":expanded_prompt},
          outputs={"response": response_text},
          model_dict={"_kind": "Openai", "engine": response["model"], "model": response["object"]}
          )

    # 将大模型 span 添加为链 span的子
    root_span.add_child(llm_span)

    # 更新链 span 的输入和输出
    root_span.add_inputs_and_outputs(
          inputs={"query":query},
          outputs={"response": response_text})

    # u更新链 span 的结束时间
    root_span.end_time_ms = llm_end_time_ms


    # 部分4 - 通过记录 root span 来记录所有 spans 到 W&B 表
    root_span.log(name="creative_trace")
    print(f"Result: {response_text}")

# 开启新的 W&B 表
wandb.init(project=PROJECT, job_type="generation")

# 运行
run_creative_chain("hero")

# 运行
run_creative_chain_cn("英雄")

# 运行
run_creative_chain("jewel")

# 运行
run_creative_chain_cn("珠宝")

wandb.finish()

# 导入需要的库
from langchain.agents import AgentType, initialize_agent
from langchain.chat_models import ChatOpenAI
from langchain.tools import BaseTool

from typing import Optional

from langchain.callbacks.manager import (
    AsyncCallbackManagerForToolRun,
    CallbackManagerForToolRun,
)

# 初始化
wandb.init(project=PROJECT, job_type="generation")

# 这将设置跟踪，因此会自动记录这些跟踪
os.environ["LANGCHAIN_WANDB_TRACING"] = "true"

class WorldPickerTool(BaseTool):
    # 随机返回列表中的一种选择
    
    name = "pick_world"
    description = "pick a virtual game world for your character or item naming"
    worlds = [
                "a mystic medieval island inhabited by intelligent and funny frogs",
                "a modern anthill featuring a cyber-ant queen and her cyber-ant-workers",
                "a digital world inhabited by friendly machine learning engineers"
            ]

    def _run(
        self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> str:
        """使用 tool"""
        time.sleep(1)
        return random.choice(self.worlds)

    async def _arun(
        self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
    ) -> str:
        """不同时使用 tool"""
        raise NotImplementedError("pick_world does not support async")
        
class NameValidatorTool(BaseTool):
    # 检查 query 或 name 是否少于20个字符
    
    name = "validate_name"
    description = "validate if the name is properly generated"

    def _run(
        self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> str:
        """使用 tool"""
        time.sleep(1)
        if len(query) < 20:
            return f"This is a correct name: {query}"
        else:
            return f"This name is too long. It should be shorter than 20 characters."

    async def _arun(
        self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
    ) -> str:
        """不同时使用 tool"""
        raise NotImplementedError("validate_name does not support async")

class WorldPickerTool_cn(BaseTool):
    # 随机返回列表中的一种选择
    
    name = "pick_world"
    description = "为命名或角色选择一个虚拟的游戏世界"
    worlds = [
                "一个神秘的中世纪岛屿，居住着聪明有趣的青蛙",
                "一个现代蚁丘，有一个网络蚁女王和她的网络蚁工",
                "一个由友好的机器学习工程师居住的数字世界"
            ]

    def _run(
        self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> str:
        """使用 tool"""
        time.sleep(1)
        return random.choice(self.worlds)

    async def _arun(
        self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
    ) -> str:
        """不同时使用 tool"""
        raise NotImplementedError("pick_world 不支持异步")
        
class NameValidatorTool_cn(BaseTool):
    # 检查 query 或 name 是否少于20个字符
    
    name = "validate_name"
    description = "验证名称是否正确生成"

    def _run(
        self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> str:
        """使用 tool"""
        time.sleep(1)
        if len(query) < 20:
            return f"这是正确名字: {query}"
        else:
            return f"这个名字太长了，它应该少于20个字符"

    async def _arun(
        self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
    ) -> str:
        """不同时使用 tool"""
        raise NotImplementedError("validate_name 不支持异步")

# 实例化大模型
llm = ChatOpenAI(temperature=0.7, openai_api_key=openai.api_key)

# 生成可选择工具的list、生成代理
tools = [WorldPickerTool(), NameValidatorTool()]
agent = initialize_agent(
    tools, 
    llm, 
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    handle_parsing_errors=True,
    verbose=True
)

# 生成可选择的工具库变量、生成代理
tools_cn = [WorldPickerTool_cn(), NameValidatorTool_cn()]
agent_cn = initialize_agent(
    tools_cn, 
    llm, 
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    handle_parsing_errors=True,
    verbose=True
)

# 输入、运行
agent.run(
    "Find a virtual game world for me and imagine the name of a hero in that world"
)

# 输入、运行
agent_cn.run(
    "为我找一个虚拟的游戏世界，想象一下那个世界里英雄的名字"
)

# 输入、运行
agent.run(
    "Find a virtual game world for me and imagine the name of a jewel in that world"
)

# 输入、运行
agent_cn.run(
    "为我找一个虚拟的游戏世界，想象一下那个世界里一颗宝石的名字"
)

# 输入、运行
agent.run(
    "Find a virtual game world for me and imagine the name of food in that world."
)

# 输入、运行
agent_cn.run(
    "给我找一个虚拟的游戏世界，想象一下这个世界里食物的名字。"
)

wandb.finish()