from dotenv import load_dotenv
load_dotenv()
import os
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
llm = ChatOpenAI()  # 初始化LLM用于推理
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")  # 初始化嵌入模型

from langchain_community.document_loaders import WebBaseLoader
urls = [
    "https://lilianweng.github.io/posts/2023-06-23-agent/",
    "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/"
]
docs = [WebBaseLoader(url).load() for url in urls]
docs_list = [item for sublist in docs for item in sublist]

from langchain_text_splitters import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=100,
    chunk_overlap=25
)
doc_splits = text_splitter.split_documents(docs_list)

from langchain_community.vectorstores import Chroma
vectorstore = Chroma.from_documents(
    documents=doc_splits,
    collection_name="rag-chrome",
    embedding=embeddings,
    persist_directory="./chroma_db"
)

from langchain.tools.retriever import create_retriever_tool
retriever = vectorstore.as_retriever()
retriever_tool = create_retriever_tool(
    retriever,
    "retriever_blog_post",
    """
    搜索并返回Lilian Weng關于LLM代理、提示工程和LLM對抗性攻擊的博客文章信息。
    僅當查詢明確與LangChain博客數據相關時使用'retriever_tool'。
    對于所有其他查詢，直接響應而不使用任何工具。
    對于'hi'、'hello'或'how are you'等簡單查詢，提供常規響應。
    """
)
tools = [retriever_tool]
from langgraph.prebuilt import ToolNode
retriever_node = ToolNode(tools)
llm_with_tool = llm.bind_tools(tools)

from typing import Annotated, Sequence, TypedDict
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]

def LLM_Decision_Maker(state: AgentState):
    print("----CALL LLM_DECISION_MAKER----")
    message = state["messages"]
    last_message = message[-1]
    question = last_message.content
    response = llm_with_tool.invoke(question)
    return {"messages": [response]}

from pydantic import BaseModel, Field
from typing import Literal
from langchain_core.prompts import PromptTemplate

class Grade(BaseModel):
    binary_score: str = Field(descriptinotallow="相關性分數'yes'或'no'")

def grade_documents(state: AgentState) -> Literal["Output Generator", "Query Rewriter"]:
    print("----CALLING GRADE FOR CHECKING RELEVANCY----")
    llm_with_structure_op = llm.with_structured_output(Grade)
    prompt = PromptTemplate(
        template="""你是一個評分員，決定文檔是否與用戶的問題相關。
                    這是文檔：{context}
                    這是用戶的問題：{question}
                    如果文檔談論或包含與用戶問題相關的信息，標記為相關。
                    給出'yes'或'no'答案，顯示文檔是否與問題相關。""",
        input_variables=["context", "question"]
    )
    chain = prompt | llm_with_structure_op
    message = state['messages']
    last_message = message[-1]
    question = message[0].content
    docs = last_message.content
    scored_result = chain.invoke({"question": question, "context": docs})
    score = scored_result.binary_score
    if score == "yes":
        print("----DECISION: DOCS ARE RELEVANT----")
        return "generator"
    else:
        print("----DECISION: DOCS ARE NOT RELEVANT----")
        return "rewriter"

from langchain import hub

def generate(state: AgentState):
    print("----RAG OUTPUT GENERATE----")
    message = state["messages"]
    question = message[0].content
    last_message = message[-1]
    docs = last_message.content
    prompt = hub.pull("rlm/rag-prompt")  # 從LangChain Hub獲取RAG提示模板
    rag_chain = prompt | llm
    response = rag_chain.invoke({"context": docs, "question": question})
    print(f"this is my response: {response}")
    return {"messages": [response]}

from langchain.schema import HumanMessage

def rewrite(state: AgentState):
    print("----TRANSFORM QUERY----")
    message = state["messages"]
    question = message[0].content
    input = [
        HumanMessage(cnotallow=f"""分析輸入并推理潛在的語義意圖或含義。
                    這是初始問題：{question}
                    制定一個改進的問題：""")
    ]
    response = llm.invoke(input)
    return {"messages": [response]}

from langgraph.graph import END, StateGraph, START
from langgraph.prebuilt import tools_condition

workflow = StateGraph(AgentState)
workflow.add_node("LLM Decision Maker", LLM_Decision_Maker)
workflow.add_node("Vector Retriever", retriever_node)
workflow.add_node("Output Generator", generate)
workflow.add_node("Query Rewriter", rewrite)

workflow.add_edge(START, "LLM Decision Maker")
workflow.add_conditional_edges(
    "LLM Decision Maker",
    tools_condition,
    {
        "tools": "Vector Retriever",
        END: END
    }
)
workflow.add_conditional_edges(
    "Vector Retriever",
    grade_documents,
    {
        "generator": "Output Generator",
        "rewriter": "Query Rewriter"
    }
)
workflow.add_edge("Output Generator", END)
workflow.add_edge("Query Rewriter", "LLM Decision Maker")

app = workflow.compile()
from IPython.display import Image, display
display(Image(app.get_graph().draw_mermaid_png()))

import warnings
warnings.filterwarnings("ignore")
app.invoke({
    "messages": [
        "what is LLM Powered Autonomous Agents explain the planning and reflection and prompt engineering explain me in terms of agents and langchain?"
    ]
})