from typing import Union, Optional, TypedDict, Annotated
from pydantic import BaseModel, Field
from dotenv import dotenv_values
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
import operator, requests, json
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage, AIMessage
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode


env_vars = dotenv_values('.env')
OPENAI_KEY = env_vars['OPENAI_API_KEY'] 
OPENAI_BASE_URL = env_vars['OPENAI_API_BASE'] 
SERPER_KEY = env_vars['SERPER_KEY'] 

llm = ChatOpenAI(model="gpt-4o-mini", api_key=OPENAI_KEY,base_url=OPENAI_BASE_URL)  


class SearchQuery(BaseModel):
    query: str = Field(description="Questions for networking queries")

class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]

@tool(args_schema = SearchQuery)
def fetch_real_time_info(query):
    """Get real-time Internet information"""
    url = "https://google.serper.dev/search"
    payload = json.dumps({
      "q": query,
      "num": 1,
    })
    headers = {
      'X-API-KEY': SERPER_KEY,
      'Content-Type': 'application/json'
    }
    response = requests.post(url, headers=headers, data=payload)
    data = json.loads(response.text)
    if 'organic' in data:
        return json.dumps(data['organic'],  ensure_ascii=False) 
    else:
        return json.dumps({"error": "No organic results found"},  ensure_ascii=False)  


def chat_with_model(state):
    """generate structured output"""
    messages = state['messages']
    response = llm.invoke(messages) 
    return {"messages": [response]}

# 判断是否要工具调用
def exists_function_calling(state: AgentState):
    result = state['messages'][-1]
    print(563, "exists_function_calling")
    return len(result.tool_calls) > 0

# 不调用工具
def final_answer(state):
    """generate natural language responses"""
    messages = state['messages'][-1]
    return {"messages": [messages]}

# 调用工具
def execute_function(state: AgentState):
    tool_calls = state['messages'][-1].tool_calls
    results = []
    tools = [fetch_real_time_info]
    tools = {t.name: t for t in tools}
    for t in tool_calls:
        if not t['name'] in tools:     
            result = "bad tool name, retry" 
        else:
            result = tools[t['name']].invoke(t['args'])
        results.append(ToolMessage(tool_call_id=t['id'], name=t['name'], content=str(result)))
    return {'messages': results}


# 请你基于现在得到的信息，进行总结，生成专业的回复
SYSTEM_PROMPT = """
Please summarize the information obtained so far and generate a professional response.
"""

# 拼接查找的信息后再最终生成结果
def natural_response(state):
    """generate final language responses"""
    messages = state['messages'][-1]
    messages = [SystemMessage(content=SYSTEM_PROMPT)] + [HumanMessage(content=messages.content)]
    response = llm.invoke(messages)
    return {"messages": [response]}


graph = StateGraph(AgentState)

graph.add_node("chat_with_model", chat_with_model)
graph.add_node("execute_function", execute_function)
graph.add_node("final_answer", final_answer)
graph.add_node("natural_response", natural_response)

# 设置图的启动节点
graph.set_entry_point("chat_with_model")
graph.add_conditional_edges(
    "chat_with_model",
    exists_function_calling,
    {True: "execute_function", False: "final_answer"}
    )
graph.add_edge("execute_function", "natural_response")
graph.set_finish_point("final_answer")
graph.set_finish_point("natural_response")
graph = graph.compile()


tools = [fetch_real_time_info]
llm = llm.bind_tools(tools)

messages = [HumanMessage(content="what is labubu")]  #测试
result = graph.invoke({"messages": messages})
res = result["messages"][-1].content
print(896, res)

整体流程如上所示