from langchain_core.tools import tool

@tool
def search_documents(query: str) -> str:
    """벡터 데이터베이스에서 관련 문서를 검색합니다."""
    docs = retriever.invoke(query)
    return "\n\n".join(doc.page_content for doc in docs)

@tool
def web_search(query: str) -> str:
    """웹에서 최신 정보를 검색합니다."""
    from langchain_community.tools import TavilySearchResults
    search = TavilySearchResults(max_results=3)
    results = search.invoke(query)
    return "\n\n".join(r["content"] for r in results)

@tool
def calculator(expression: str) -> str:
    """수학 계산을 수행합니다. 예: '2 + 3 * 4'"""
    try:
        # 안전한 수식 평가
        import ast
        import operator
        ops = {
            ast.Add: operator.add, ast.Sub: operator.sub,
            ast.Mult: operator.mul, ast.Div: operator.truediv,
        }
        def _eval(node):
            if isinstance(node, ast.Expression):
                return _eval(node.body)
            elif isinstance(node, ast.Constant):
                return node.value
            elif isinstance(node, ast.BinOp) and type(node.op) in ops:
                return ops[type(node.op)](_eval(node.left), _eval(node.right))
            raise ValueError(f"지원하지 않는 연산")
        return str(_eval(ast.parse(expression, mode="eval")))
    except Exception as e:
        return f"계산 오류: {e}"

tools = [search_documents, web_search, calculator]

from typing import TypedDict, Annotated, Sequence
from langchain.chat_models import init_chat_model
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition

# 상태 정의
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]

# LLM에 도구 바인딩
llm = init_chat_model("gpt-4o-mini", temperature=0)
llm_with_tools = llm.bind_tools(tools)

SYSTEM_PROMPT = """당신은 질문에 답변하는 AI 어시스턴트입니다.
주어진 도구를 활용하여 정확한 정보를 수집한 뒤 답변하세요.
- search_documents: 내부 문서에서 정보를 찾을 때
- web_search: 최신 정보나 내부 문서에 없는 정보가 필요할 때
- calculator: 수학 계산이 필요할 때
충분한 정보를 수집한 후에 답변하세요."""

# 에이전트 노드
def agent(state: AgentState) -> AgentState:
    messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

# 그래프 구성
workflow = StateGraph(AgentState)

workflow.add_node("agent", agent)
workflow.add_node("tools", ToolNode(tools))

workflow.add_edge(START, "agent")
workflow.add_conditional_edges("agent", tools_condition)
workflow.add_edge("tools", "agent")

app = workflow.compile()

# 문서 검색 활용
result = app.invoke({
    "messages": [HumanMessage(content="Self-RAG의 Reflection Token에 대해 설명해줘")]
})
print(result["messages"][-1].content)

# 웹 검색 활용
result = app.invoke({
    "messages": [HumanMessage(content="2024년 주목받은 RAG 논문은?")]
})
print(result["messages"][-1].content)

# 복합 도구 활용 (검색 + 계산)
result = app.invoke({
    "messages": [HumanMessage(content="RAG 논문이 2023년 대비 2024년에 몇 배 증가했나?")]
})
print(result["messages"][-1].content)

# 최대 반복 횟수 제한
result = app.invoke(
    {"messages": [HumanMessage(content="질문")]},
    config={"recursion_limit": 10},
)