from openai import OpenAI
client = OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_API_URL)

agent_memory = {"human": "", "agent": ""}
def core_memory_save(section: str, memory: str):
    agent_memory[section] += "\n"
    agent_memory[section] += memory.strip()
    print("current memory: ", agent_memory)
core_memory_save_desc = """Save important information about you, the agent or the human you are chatting with."""
core_memory_save_propertis = {
    "section": {
        "type": "string",
        "enum": ["human", "agent"],
        "description": "Must be either 'human'(to save information about the human) \
            or 'agent'(to save information about yourself)."
    },
    "memory": {
        "type": "string",
        "description": "memory to save in the section."
    }
}
# 这里提供了一个 core_memory_save 的 tool
core_memory_save_metadata = {
    "type": "function",
    "function": {
        "name": "core_memory_save",
        "description": core_memory_save_desc,
        "parameters": {
            "type": "object",
            "properties": core_memory_save_propertis,
            "required": ["section", "memory"]
        }
    }
}

system_prompt = """you are a chatbot.
You have a section of your context called [MEMORY]
that contains information relevant to your conversation
you can use this information to answer questions.
"""
messages=[
    {"role": "system", "content": system_prompt},
    {"role": "system", "content": "[MEMORY]\n" + json.dumps(agent_memory)},
    {"role": "user", "content": "My name is QuantumForge"},
]

chat_completion = client.chat.completions.create(
    model=OPENAI_MODEL_NAME,
    messages=messages,
    tools=[core_memory_save_metadata],
    tool_choice="auto",
)
response = chat_completion.choices[0]
print(response)

Choice(finish_reason='tool_calls', index=0, logprobs=None, message=ChatCompletionMessage(content='\n\n', refusal=None, role='assistant', annotations=None, audio=None, function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_ortft9h7z3gh7l9r6kzt67w6', function=Function(arguments='{"section": "human", "memory": "QuantumForge"}', name='core_memory_save'), type='function')], reasoning_content='Okay, the user mentioned their name is QuantumForge. I need to save that into the core memory under the human section. Let me call the core_memory_save function. The parameters should be section "human" and memory "QuantumForge". Make sure the JSON is correctly formatted.\n'))
{'section': 'human', 'memory': 'QuantumForge'}
current memory:  {'human': 'QuantumForge'}

pip install arize-phoenix arize-phoenix-otel

pip install openinference-semantic-conventions opentelemetry-api opentelemetry-sdk
pip install openinference-instrumentation-openai

def start_main_span(messages):
    print("Starting main span with messages:", messages)

    with tracer.start_as_current_span(
        "AgentRun", openinference_span_kind="agent"
    ) as span:
        span.set_input(value=messages)
        ret = run_agent(messages)
        print("Main span completed with return value:", ret)
        span.set_output(value=ret)
        span.set_status(StatusCode.OK)
        return ret

# ...
while True:
        # Router Span
        print("Starting router call span")
<strong>        with tracer.start_as_current_span(</strong>
<strong>            "router_call", openinference_span_kind="chain",</strong>
<strong>        ) as span:</strong>
<strong>            span.set_input(value=messages)</strong>

            response = client.chat.completions.create(
                model=MODEL,
                messages=messages,
                tools=tools,
            )
        # ...

# 从 Phoenix 中加载 Tools Router 的 Spans
query = SpanQuery().where(
    "span_kind == 'LLM'",
).select(
    question="input.value",
    tool_call="llm.tools"
)
tool_calls_df = px.Client().query_spans(query,
                                        project_name=PROJECT_NAME,
                                        timeout=None)
tool_calls_df = tool_calls_df.dropna(subset=["tool_call"])

# 使用模型来判断 Tools 选择的正确性
with suppress_tracing():
    tool_call_eval = llm_classify(
        dataframe = tool_calls_df,
        template = TOOL_CALLING_PROMPT_TEMPLATE.template[0].template.replace(
            "{tool_definitions}",                                                                 json.dumps(tools).replace("{", '"').replace("}", '"')
        ),
        rails = ['correct', 'incorrect'],
        model=OpenAIModel(model="gpt-4o"),
        provide_explanation=True
    )

tool_call_eval['score'] = tool_call_eval.apply(
    lambda x: 1 if x['label']=='correct' else 0, axis=1
)
# 上报 evaluation 结果
px.Client().log_evaluations(
    SpanEvaluations(eval_name="Tool Calling Eval", dataframe=tool_call_eval),
)

convergence_questions = [
    "What was the average quantity sold per transaction?",
    # ...
]
convergence_df = pd.DataFrame({
    "question": convergence_questions
})
now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
dataset = px_client.upload_dataset(dataframe=convergence_df,
                                   dataset_name=f"convergence_questions-{now}",
                                   input_keys=["question"])

def format_message_steps(messages):
    steps = []
    for message in messages:
        if type(message) is ChatCompletionMessage:
            message = message.to_dict()
        role = message.get("role")
        if role == "user":
            steps.append(f"User: {message.get('content')}")
        elif role == "system":
            steps.append("System: Provided context")
        elif role == "assistant":
            if message.get("tool_calls"):
                for tool_call in message["tool_calls"]:
                    tool_name = tool_call["function"]["name"]
                    steps.append(f"Assistant: Called tool '{tool_name}'")
            else:
                steps.append(f"Assistant: {message.get('content')}")
        elif role == "tool":
            steps.append(f"Tool response: {message.get('content')}")
    return "\n".join(steps)

def run_agent_and_track_path(example: Example) -> str:
    messages = [{"role": "user", "content": example.input.get("question")}]
    # 重复运行 agent
    ret = run_agent(messages)
    return {"path_length": len(messages), "messages": format_message_steps(messages)}

experiment = run_experiment(dataset,
                            run_agent_and_track_path,
                            experiment_name="Convergence Eval",
                            experiment_description="Evaluating the convergence of the agent")
outputs = experiment.as_dataframe()["output"].to_dict().values()

optimal_path_length = min(output.get('path_length') for output in outputs \
    if output and output.get('path_length') is not None)
print(f"The optimal path length is {optimal_path_length}")

@create_evaluator(name="Convergence Eval", kind="CODE")
def evaluate_path_length(output: str) -> float:
    if output and output.get("path_length"):
        return optimal_path_length/float(output.get("path_length"))
    else:
        return 0
experiment = evaluate_experiment(experiment,
                            evaluators=[evaluate_path_length])
print(experiment.as_dataframe())

from phoenix.experiments import run_experiment

experiment = run_experiment(
    # 测试的输入问题和 Expected 问题集合
    dataset,
    # Agent 主逻辑代码
    run_agent_task,
    # evaluators 输入是 agent 的 input/output message
    # 输出是 bool 或者 float，用来评估效果
    evaluators=[function_calling_eval,
                evaluate_sql_result,
                evaluate_clarity,
                evaluate_entity_correctness,
                code_is_runnable],
    experiment_name="Overall Experiment",
    experiment_description="Evaluating the overall experiment")

MODEL_NAME = "../models/Qwen/Qwen2.5-3B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# 入参 prompt，本地推理生成 response 文本
def query_model(prompt: str) -> str:
    messages=[{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=10240
    )
    generated_ids = [
        output_ids[len(input_ids):]
        for input_ids, output_ids in zip(
            model_inputs.input_ids, generated_ids
        )
    ]

    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    response = client.get_ans(messages)
    print("model output: ", response)
    return response

# SQL 查询语句升成的 prompt 模板
SQL_GENERATION_PROMPT = """
Generate an SQL query based on a prompt. Do not reply with anything besides the SQL query.
The prompt is: {prompt}

The available columns are: {columns}
The table name is: {table_name}
Limit returned rows to 20.
"""
def generate_sql_query(prompt: str, columns: list, table_name: str) -> str:
    """Generate an SQL query based on a prompt"""
    formatted_prompt = SQL_GENERATION_PROMPT.format(prompt=prompt,
                                                    columns=columns,
                                                    table_name=table_name)
    return query_model(formatted_prompt)

TRANSACTION_DATA_FILE_PATH = './Store_Sales_Price_Elasticity_Promotions_Data.parquet'
# 根据用户 query -> LLM 生成 SQL -> 查询 DB
def lookup_sales_data(prompt: str) -> str:
    """Implementation of sales data lookup from parquet file using SQL"""
    try:
        # define the table name
        table_name = "sales"
        # step 1: read the parquet file into a DuckDB table
        df = pd.read_parquet(TRANSACTION_DATA_FILE_PATH)
        duckdb.sql(f"CREATE TABLE IF NOT EXISTS {table_name} AS SELECT * FROM df")
        # step 2: generate the SQL code
        sql_query = generate_sql_query(prompt, df.columns, table_name)
        # clean the response to make sure it only includes the SQL code
        sql_query = sql_query.strip()
        sql_query = sql_query.replace("```sql", "").replace("```", "")
        # step 3: execute the SQL query
        result = duckdb.sql(sql_query).df()

        return result.to_string()
    except Exception as e:
        return f"Error accessing data: {str(e)}"

# 数据分析的 Prompt 模板
DATA_ANALYSIS_PROMPT = """
Analyze the following data: {data}
Your job is to answer the following question: {prompt}
"""
def analyze_sales_data(prompt: str, data: str) -> str:
    """Implementation of AI-powered sales data analysis"""
    formatted_prompt = DATA_ANALYSIS_PROMPT.format(data=data, prompt=prompt)
    analysis = query_model(formatted_prompt)

    return analysis if analysis else "No analysis could be generated"

# 生成 Chart 配置的 Prompt 模板
CHART_CONFIGURATION_PROMPT = """
Generate a chart configuration based on this data: {data}
The goal is to show: {visualization_goal}
Return the chart configuration as a JSON object with the following keys:
- chart_type: Type of chart to generate
- x_axis: Name of the x-axis column
- y_axis: Name of the y-axis column
- title: Title of the chart
Only return the JSON object, no other text.
"""
class VisualizationConfig(BaseModel):
    chart_type: str = Field(..., description="Type of chart to generate")
    x_axis: str = Field(..., description="Name of the x-axis column")
    y_axis: str = Field(..., description="Name of the y-axis column")
    title: str = Field(..., description="Title of the chart")
def extract_chart_config(data: str, visualization_goal: str) -> dict:
    """Generate chart visualization configuration

    Args:
        data: String containing the data to visualize
        visualization_goal: Description of what the visualization should show

    Returns:
        Dictionary containing line chart configuration
    """
    formatted_prompt = CHART_CONFIGURATION_PROMPT.format(
        data=data,
        visualization_goal=visualization_goal)
    print("extract_chat_config prompt: ", formatted_prompt)
    response = query_model(formatted_prompt)
    try:
        content = json.loads(response)
        return {
            "chart_type": content.chart_type,
            "x_axis": content.x_axis,
            "y_axis": content.y_axis,
            "title": content.title,
            "data": data
        }
    except Exception:
        return {
            "chart_type": "line",
            "x_axis": "date",
            "y_axis": "value",
            "title": visualization_goal,
            "data": data
        }

# 生成画图 Python 代码的 Prompt 模板
CREATE_CHART_PROMPT = """
Write python code to create a chart based on the following configuration.
Only return the code, no other text.
config: {config}
"""
def create_chart(config: dict) -> str:
    """Create a chart based on the configuration"""
    formatted_prompt = CREATE_CHART_PROMPT.format(config=config)

    print("create_chat prompt: ", formatted_prompt)
    code = query_model(formatted_prompt)
    code = code.replace("```python", "").replace("```", "")
    code = code.strip()

    return code

def generate_visualization(data: str, visualization_goal: str) -> str:
    """Generate a visualization based on the data and goal"""
    config = extract_chart_config(data, visualization_goal)
    code = create_chart(config)
    return code

tools = [
    {
        "type": "function",
        "function": {
            "name": "lookup_sales_data",
            "description": "Look up data from Store Sales Price Elasticity Promotions dataset",
            "parameters": {
                "type": "object",
                "properties": {
                    "prompt": {"type": "string", "description": "The unchanged prompt that the user provided."}
                },
                "required": ["prompt"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "analyze_sales_data",
            "description": "Analyze sales data to extract insights",
            "parameters": {
                "type": "object",
                "properties": {
                    "data": {"type": "string", "description": "The lookup_sales_data tool's output."},
                    "prompt": {"type": "string", "description": "The unchanged prompt that the user provided."}
                },
                "required": ["data", "prompt"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "generate_visualization",
            "description": "Generate Python code to create data visualizations",
            "parameters": {
                "type": "object",
                "properties": {
                    "data": {"type": "string", "description": "The lookup_sales_data tool's output."},
                    "visualization_goal": {"type": "string", "description": "The goal of the visualization."}
                },
                "required": ["data", "visualization_goal"]
            }
        }
    }
]

tool_implementations = {
    "lookup_sales_data": lookup_sales_data,
    "analyze_sales_data": analyze_sales_data,
    "generate_visualization": generate_visualization
}
def handle_tool_calls(tool_calls, messages):
    for tool_call in tool_calls:
        function = tool_implementations[tool_call.function.name]
        function_args = json.loads(tool_call.function.arguments)
        result = function(**function_args)
        messages.append({
            "role": "tool", "content": result, "tool_call_id": tool_call.id
        })

    return messages

SYSTEM_PROMPT = """
You are a helpful assistant that can answer questions about the Store Sales Price Elasticity Promotions dataset.
"""
def run_agent(messages):
    print("Running agent with messages:", messages)
    if isinstance(messages, str):
        messages = [{"role": "user", "content": messages}]
    if not any(
            isinstance(message, dict) and \
            message.get("role") == "system" for message in messages
        ):
            system_prompt = {"role": "system", "content": SYSTEM_PROMPT}
            messages.append(system_prompt)

    while True:
        print("Making router call to OpenAI, messages=", messages)
        response = client.client.chat.completions.create(
            model=MODEL,
            messages=messages,
            tools=tools,
        )
        messages.append(response.choices[0].message)
        tool_calls = response.choices[0].message.tool_calls
        print("Received response with tool calls:", bool(tool_calls))

        # if the model decides to call function(s), call handle_tool_calls
        if tool_calls:
            print("Processing tool calls")
            messages = handle_tool_calls(tool_calls, messages)
        else:
            print("No tool calls, returning final response")
            return response.choices[0].message.content

result = run_agent('Show me the code for graph of sales by store in Nov 2021, and tell me what trends you see.')
print(result)

from transformers import AutoTokenizer

# A list of colors in RGB for representing the tokens
colors = [
    '102;194;165', '252;141;98', '141;160;203',
    '231;138;195', '166;216;84', '255;217;47'
]

def show_tokens(sentence: str, tokenizer_name: str):
    """ Show the tokens each separated by a different color """

    # Load the tokenizer and tokenize the input
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    token_ids = tokenizer(sentence).input_ids

    # Extract vocabulary length
    print(f"Vocab length: {len(tokenizer)}")

    # Print a colored list of tokens
    for idx, t in enumerate(token_ids):
        print(
            f'\x1b[0;30;48;2;{colors[idx % len(colors)]}m' +
            tokenizer.decode(t) +
            '\x1b[0m',
            end=' '
        )

sentence = """
English and CAPITALIZATION
🎵 鸟
show_tokens False None elif == >= else: two tabs:"    " Three tabs: "       "
12.0*50=600
"""
show_tokens(sentence, "bert-base-cased")
show_tokens(sentence, "Xenova/gpt-4")
show_tokens(sentence, "Qwen/Qwen2-VL-7B-Instruct")

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    "../models/microsoft/Phi-3-mini-4k-instruct"
)
model = AutoModelForCausalLM.from_pretrained(
    "../models/microsoft/Phi-3-mini-4k-instruct",
    device_map="cpu",
    torch_dtype="auto",
    trust_remote_code=True,
)

# Create a pipeline
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    return_full_text=False,
    do_sample=False,
)

prompt = "Write an email apologizing to Sarah for the tragic gardening mishap. Explain how it happened. "
output = generator(prompt)
print(output[0]['generated_text'])

prompt = "The capital of France is"
# Tokenize the input prompt
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
# Get the output of the model before the lm_head
model_output = model.model(input_ids)
# Get the output of the lm_head
lm_head_output = model.lm_head(model_output[0])
token_id = lm_head_output[0,-1].argmax(-1)
tokenizer.decode(token_id)