from dotenv import load_dotenv, find_dotenv
import os
from langchain.chat_models import init_chat_model

_ = load_dotenv(find_dotenv())
model_name = os.environ['MODEL_NAME']
api_key = os.environ['API_KEY']
temperature = 0.0

chat = init_chat_model(model=model_name, api_key=api_key, temperature=temperature)

template_string = """Translate the text \
that is delimited by triple backticks \
into a style that is {style}. \
text: ```{text}```
"""

from langchain_core.prompts import ChatPromptTemplate

prompt_template = ChatPromptTemplate.from_template(template_string)

customer_style = """American English \
in a calm and respectful tone
"""

customer_email = """
Arrr, I be fuming that me blender lid \
flew off and splattered me kitchen walls \
with smoothie! And to make matters worse, \
the warranty don't cover the cost of \
cleaning up me kitchen. I need yer help \
right now, matey!
"""

customer_messages = prompt_template.format_messages(
                    style=customer_style,
                    text=customer_email,)

print(customer_messages[0])

customer_response = chat.invoke(customer_messages)

print(customer_response.content)

from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel, Field


class ReviewOutput(BaseModel):
    gift: bool = Field(
        description="Was the item purchased as a gift for someone else? Answer True if yes,False if not or unknown."
    )

    delivery_days: int = Field(
        description="How many days did it take for the product to arrive? If this information is not found,output -1."
    )

    price_value: list[str] = Field(
        description="Extract any sentences about the value or price, and output them as a comma separated Python list."
    )

output_parser = JsonOutputParser(pydantic_object=ReviewOutput)

format_instructions = output_parser.get_format_instructions()

customer_review = """\
This leaf blower is pretty amazing.  It has four settings:\
candle blower, gentle breeze, windy city, and tornado. \
It arrived in two days, just in time for my wife's \
anniversary present. \
I think my wife liked it so much she was speechless. \
So far I've been the only one using it, and I've been \
using it every other morning to clear the leaves on our lawn. \
It's slightly more expensive than the other leaf blowers \
out there, but I think it's worth it for the extra features.
"""

review_template = """\
For the following text, extract the following information:

gift: Was the item purchased as a gift for someone else? \
Answer True if yes, False if not or unknown.

delivery_days: How many days did it take for the product\
to arrive? If this information is not found, output -1.

price_value: Extract any sentences about the value or price,\
and output them as a comma separated Python list.

text: {text}

{format_instructions}
"""

prompt = ChatPromptTemplate.from_template(template=review_template)

messages = prompt.format_messages(text=customer_review, 
                                format_instructions=format_instructions)

response = chat.invoke(messages)

prompt = ChatPromptTemplate.from_template(template=review_template).partial(format_instructions=format_instructions)

chain = prompt | chat | output_parser

response = chain.invoke({"text": customer_review})

print(response)

from langchain.chat_models import init_chat_model
from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver
from dotenv import load_dotenv, find_dotenv
import os

_ = load_dotenv(find_dotenv())
model_name = os.environ['MODEL_NAME']
api_key = os.environ['API_KEY']
temperature = 0.0

llm = init_chat_model(model=model_name, api_key=api_key, temperature=temperature)

memory = InMemorySaver()

conversation = create_agent(
    model=llm,
    tools=[],
    checkpointer=memory,
)

config = {"configurable": {"thread_id": "user_001"}}

conversation.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "Hi, my name is Andrew"
            }
        ]
    },
    config=config
)

conversation.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "What is 1+1?"
            }
        ]
    },
    config=config
)

conversation.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "What is my name?"
            }
        ]
    },
    config=config
)

# 查看历史对话
state = conversation.get_state(config)
for message in state.values["messages"]:
    print(message)

from langchain_core.messages import HumanMessage, AIMessage

history = [
    HumanMessage(content="Hi"),
    AIMessage(content="What's up"),
    HumanMessage(content="Not much, just hanging"),
    AIMessage(content="Cool"),
]

conversation.update_state(
    config=config,
    values={"messages": history}
)

# 查看历史对话
state = conversation.get_state(config)
for message in state.values["messages"]:
    print(message)

from langchain_core.messages import HumanMessage, AIMessage, trim_messages, RemoveMessage
from langchain.agents.middleware import before_model
from langgraph.checkpoint.memory import InMemorySaver

# 控制模型保留消息能力
trimmer = trim_messages(
    max_tokens=2,
    token_counter=len,
    strategy="last",
    start_on="human",
    include_system=False,
)

# 模型调用前裁剪历史
@before_model
def trim_history(state, runtime):
    trimmed = trimmer.invoke(state["messages"])
    keep_ids = {m.id for m in trimmed}
    message = {"messages": [RemoveMessage(id=m.id) for m in state["messages"] if m.id not in keep_ids]}
    return message

memory = InMemorySaver()

conversation = create_agent(
    model=llm,
    tools=[],
    checkpointer=memory,
    middleware=[trim_history],
)

config = {"configurable": {"thread_id": "user_002"}}

conversation.update_state(config, {"messages": [
    HumanMessage("Hi"),
    AIMessage("What's up"),
]})

conversation.update_state(config, {"messages": [
    HumanMessage("Not much, just hanging"),
    AIMessage("Cool"),
]})

# 查看注入对话
print([m.content for m in conversation.get_state(config).values["messages"]])

conversation.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "Hi, my name is Andrew"
            }
        ]
    },
    config=config
)

conversation.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "What is 1+1?"
            }
        ]
    },
    config=config
)

conversation.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "What is my name?"
            }
        ]
    },
    config=config
)

# 查看历史对话
state = conversation.get_state(config)
for message in state.values["messages"]:
    print(message)

from langchain_core.messages import HumanMessage, AIMessage, trim_messages, RemoveMessage
from langchain.agents.middleware import before_model
from langgraph.checkpoint.memory import InMemorySaver

# 控制模型保留消息能力
trimmer = trim_messages(
    max_tokens=30,
    token_counter="approximate",
    strategy="last",
    start_on="human",
    include_system=False,
)

# 模型调用前裁剪历史
@before_model
def trim_history(state, runtime):
    trimmed = trimmer.invoke(state["messages"])
    keep_ids = {m.id for m in trimmed}
    message = {"messages": [RemoveMessage(id=m.id) for m in state["messages"] if m.id not in keep_ids]}
    return message

memory = InMemorySaver()

conversation = create_agent(
    model=llm,
    tools=[],
    checkpointer=memory,
    middleware=[trim_history],
)

config = {"configurable": {"thread_id": "user_003"}}

conversation.update_state(config, {"messages": [
    HumanMessage("Hi"),
    AIMessage("What's up"),
]})

conversation.update_state(config, {"messages": [
    HumanMessage("Not much, just hanging"),
    AIMessage("Cool"),
]})

conversation.update_state(config, {"messages": [
    HumanMessage("What is on the schedule today?"),
    AIMessage(schedule),
]})

# 查看注入对话
print([m.content for m in conversation.get_state(config).values["messages"]])

conversation.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "Hi, my name is Andrew"
            }
        ]
    },
    config=config
)

conversation.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "What is 1+1?"
            }
        ]
    },
    config=config
)

conversation.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "What is my name?"
            }
        ]
    },
    config=config
)

# 查看历史对话
state = conversation.get_state(config)
for message in state.values["messages"]:
    print(message)

from langchain.agents import create_agent
from langchain.agents.middleware import SummarizationMiddleware
from langgraph.checkpoint.memory import InMemorySaver

summary_history = SummarizationMiddleware(
            model=llm,
            trigger=("tokens", 100),
            keep=("messages", 2),
            )

conversation = create_agent(
    model=llm,
    tools=[],
    checkpointer=InMemorySaver(),
    middleware=[summary_history,],
)

config = {"configurable": {"thread_id": "user_004"}}

conversation.update_state(config, {"messages": [
    HumanMessage("Hi"),
    AIMessage("What's up"),
]})

conversation.update_state(config, {"messages": [
    HumanMessage("Not much, just hanging"),
    AIMessage("Cool"),
]})

schedule = "There is a meeting at 8am with your product team. \
You will need your powerpoint presentation prepared. \
9am-12pm have time to work on your LangChain \
project which will go quickly because Langchain is such a powerful tool. \
At Noon, lunch at the italian resturant with a customer who is driving \
from over an hour away to meet you to understand the latest in AI. \
Be sure to bring your laptop to show the latest LLM demo."

conversation.update_state(config, {"messages": [
    HumanMessage("What is on the schedule today?"),
    AIMessage(schedule),
]})

# 查看注入消息
print([m.content for m in conversation.get_state(config).values["messages"]])

conversation.invoke(
    {"messages": [HumanMessage("What would be a good demo to show?")]},
    config,
)

state = conversation.get_state(config)
for message in state.values["messages"]:
    print(message)

from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from dotenv import load_dotenv, find_dotenv
import os

_ = load_dotenv(find_dotenv())
model_name = os.environ['MODEL_NAME']
api_key = os.environ['API_KEY']
temperature = 0.9

llm = init_chat_model(model=model_name, api_key=api_key, temperature=temperature)

prompt = ChatPromptTemplate.from_template(
    "What is the best name to describe \
    a company that makes {product}?"
)

chain = prompt | llm

product = "Queen Size Sheet Set"
chain.invoke({"product": product})

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

# prompt template 1
first_prompt = ChatPromptTemplate.from_template(
    "What is the best name to describe a company that makes {product}?"
)

# Chain 1
chain_one = first_prompt | llm | StrOutputParser()

# prompt template 2
second_prompt = ChatPromptTemplate.from_template(
    "Write a 20 words description for the following company: {company_name}"
)
# chain 2
chain_two = second_prompt | llm | StrOutputParser()

overall_simple_chain = {"company_name": chain_one} | chain_two

product = "Queen Size Sheet Set"
overall_simple_chain.invoke({"product": product})

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

# prompt template 1: translate to english
first_prompt = ChatPromptTemplate.from_template(
    "Translate the following review to english:"
    "\n\n{Review}"
)
# chain 1: input= Review and output= English_Review
chain_one = first_prompt | llm | StrOutputParser()

second_prompt = ChatPromptTemplate.from_template(
    "Can you summarize the following review in 1 sentence:"
    "\n\n{English_Review}"
)
# chain 2: input= English_Review and output= summary
chain_two = second_prompt | llm | StrOutputParser()


# prompt template 3: translate to english
third_prompt = ChatPromptTemplate.from_template(
    "What language is the following review:\n\n{Review}"
)
# chain 3: input= Review and output= language
chain_three = third_prompt | llm | StrOutputParser()

# prompt template 4: follow up message
fourth_prompt = ChatPromptTemplate.from_template(
    "Write a follow up response to the following "
    "summary in the specified language:"
    "\n\nSummary: {summary}\n\nLanguage: {language}"
)
# chain 4: input= summary, language and output= followup_message
chain_four = fourth_prompt | llm | StrOutputParser()

overall_chain = (RunnablePassthrough.assign(English_Review=chain_one)
                | RunnablePassthrough.assign(summary=chain_two)
                | RunnablePassthrough.assign(language=chain_three)
                | RunnablePassthrough.assign(followup_message=chain_four)
                 )

review = df.Review[5]
overall_chain.invoke({"Review": review})

physics_template = """You are a very smart physics professor. \
You are great at answering questions about physics in a concise\
and easy to understand manner. \
When you don't know the answer to a question you admit\
that you don't know.

Here is a question:
{input}"""


math_template = """You are a very good mathematician. \
You are great at answering math questions. \
You are so good because you are able to break down \
hard problems into their component parts, 
answer the component parts, and then put them together\
to answer the broader question.

Here is a question:
{input}"""

history_template = """You are a very good historian. \
You have an excellent knowledge of and understanding of people,\
events and contexts from a range of historical periods. \
You have the ability to think, reflect, debate, discuss and \
evaluate the past. You have a respect for historical evidence\
and the ability to make use of it to support your explanations \
and judgements.

Here is a question:
{input}"""


computerscience_template = """ You are a successful computer scientist.\
You have a passion for creativity, collaboration,\
forward-thinking, confidence, strong problem-solving capabilities,\
understanding of theories and algorithms, and excellent communication \
skills. You are great at answering coding questions. \
You are so good because you know how to solve a problem by \
describing the solution in imperative steps \
that a machine can easily interpret and you know how to \
choose a solution that has a good balance between \
time complexity and space complexity. 

Here is a question:
{input}"""

prompt_infos = [
    {
        "name": "physics", 
        "description": "Good for answering questions about physics", 
        "prompt_template": physics_template
    },
    {
        "name": "math", 
        "description": "Good for answering math questions", 
        "prompt_template": math_template
    },
    {
        "name": "History", 
        "description": "Good for answering history questions", 
        "prompt_template": history_template
    },
    {
        "name": "computer science", 
        "description": "Good for answering computer science questions", 
        "prompt_template": computerscience_template
    }
]

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableBranch, RunnableLambda
from pydantic import BaseModel, Field
from typing import Literal

destination_chains = {
    p["name"]: ChatPromptTemplate.from_template(p["prompt_template"]) | llm | StrOutputParser()
    for p in prompt_infos
}
default_chain = ChatPromptTemplate.from_template("{input}") | llm | StrOutputParser()

class Route(BaseModel):
    """Route a user query to the most relevant expert."""
    destination: Literal["physics", "math", "History", "computer science", "DEFAULT"] = Field(
        description="The expert best suited for the question, or DEFAULT if none fits."
    )
    next_inputs: str = Field(
        description="A potentially revised version of the original input."
    )

destinations_str = "\n".join(f"{p['name']}: {p['description']}" for p in prompt_infos)

router_prompt = ChatPromptTemplate.from_messages([
    ("system",
     "Given a raw text input to a language model, select the prompt best suited for it.\n"
     "You may also revise the original input if that leads to a better response.\n\n"
     "<< CANDIDATE PROMPTS >>\n{destinations}"),
    ("human", "{input}"),
])

router_chain = (
    router_prompt.partial(destinations=destinations_str)
    | llm.with_structured_output(Route)
)

overall_chain = (
    router_chain
    | {
        "destination": RunnableLambda(lambda x: x.destination),
        "input": RunnableLambda(lambda x: x.next_inputs)
    }
    | RunnableBranch(
        (lambda s: s["destination"] == "physics", destination_chains["physics"]),
        (lambda s: s["destination"] == "math", destination_chains["math"]),
        (lambda s: s["destination"] == "History", destination_chains["History"]),
        (lambda s: s["destination"] == "computer science", destination_chains["computer science"]),
        default_chain
    )
)

overall_chain.invoke({"input": "What is black body radiation?"})

from langchain_huggingface import HuggingFaceEmbeddings
# 根据你的Embedding模型进行选择，这里使用huggingface的做快速验证
# from langchain_openai import OpenAIEmbeddings
from langchain.chat_models import init_chat_model
from langchain_community.document_loaders import CSVLoader
from langchain_community.vectorstores import DocArrayInMemorySearch
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from IPython.display import display, Markdown

from dotenv import load_dotenv, find_dotenv
import os
from langchain.chat_models import init_chat_model

_ = load_dotenv(find_dotenv())
model_name = os.environ['MODEL_NAME']
api_key = os.environ['API_KEY']
temperature = 0.0

llm = init_chat_model(model=model_name, api_key=api_key, temperature=temperature)

file = 'OutdoorClothingCatalog_1000.csv'
loader = CSVLoader(file_path=file,encoding="utf-8")

docs = loader.load()

embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": "cuda"},
    encode_kwargs={"normalize_embeddings": True},
)
# embeddings = OpenAIEmbeddings()

vectorstore = DocArrayInMemorySearch.from_documents(docs, embeddings)
retriever = vectorstore.as_retriever()

prompt = ChatPromptTemplate.from_template(
    "基于以下上下文回答问题：\n\n{context}\n\n问题：{question}"
)

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

query = "Please list all your shirts with sun protection in a table in markdown and summarize each one."
response = rag_chain.invoke(query)
display(Markdown(response))

embed = embeddings.embed_query("Hi my name is Harrison")
print(len(embed))
print(embed[:5])

db = DocArrayInMemorySearch.from_documents(docs, embeddings)
query = "Please suggest a shirt with sunblocking"
docs = db.similarity_search(query)
retriever = db.as_retriever()

qdocs = "".join([docs[i].page_content for i in range(len(docs))])
response = llm.invoke(f"{qdocs} Question: Please list all your shirts with sun protection in a table in markdown and summarize each one.")
display(Markdown(response.content))

from dotenv import load_dotenv, find_dotenv
import os
from langchain.chat_models import init_chat_model
from langchain_huggingface import HuggingFaceEmbeddings
# 根据你的Embedding模型进行选择，这里使用huggingface的做快速验证
# from langchain_openai import OpenAIEmbeddings
from langchain_community.document_loaders import CSVLoader
from langchain_community.vectorstores import DocArrayInMemorySearch
from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate

_ = load_dotenv(find_dotenv())
model_name = os.environ['MODEL_NAME']
api_key = os.environ['API_KEY']
temperature = 0.0

llm = init_chat_model(model=model_name, api_key=api_key, temperature=temperature)

file = 'OutdoorClothingCatalog_1000.csv'
loader = CSVLoader(file_path=file,encoding="utf-8")

docs = loader.load()

embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": "cuda"},
    encode_kwargs={"normalize_embeddings": True},
)
# embeddings = OpenAIEmbeddings()

vectorstore = DocArrayInMemorySearch.from_documents(docs, embeddings)
retriever = vectorstore.as_retriever()

system_prompt = (
    "Use the following pieces of context to answer the user's question. "
    "If you don't know the answer, just say that you don't know, "
    "don't try to make up an answer.\n"
    "----------------\n"
    "{context}"
)

prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}"),
])

question_answer_chain = create_stuff_documents_chain(
    llm,
    prompt,
    document_separator="<<<<>>>>>"
)
qa_chain = create_retrieval_chain(retriever, question_answer_chain)

examples = [
    {
        "query": "Do the Cozy Comfort Pullover Set have side pockets?",
        "answer": "Yes"
    },
    {
        "query": "What collection is the Ultra-Lofty 850 Stretch Down Hooded Jacket from?",
        "answer": "The DownTek collection"
    }
]

result = qa_chain.invoke({"input": examples[0]["query"]})
print(result["answer"])

from langchain_core.output_parsers import StrOutputParser

example_gen_prompt = ChatPromptTemplate.from_messages([
    ("system",
     "You are a teacher. Given a document, generate one question and answer "
     "based on the document. Output in JSON format with keys 'query' and 'answer'."),
    ("human", "{doc}")
])

example_gen_chain = example_gen_prompt | llm | StrOutputParser()

import json
new_examples = []
for doc in docs[:5]:
    raw = example_gen_chain.invoke({"doc": doc.page_content})
    parsed = json.loads(raw.strip("```json").strip("```"))
    new_examples.append(parsed)

examples += new_examples

eval_prompt = ChatPromptTemplate.from_messages([
    ("system",
     "You are a teacher grading a quiz. You are given a question, the ground truth "
     "answer, and a student's answer. Grade the student's answer as CORRECT or INCORRECT.\n\n"
     "Question: {query}\n"
     "Ground truth: {answer}\n"
     "Student answer: {result}"),
])

eval_chain = eval_prompt | llm | StrOutputParser()

predictions = [{"query": e["query"], "result": qa_chain.invoke({"input": e["query"]})["answer"]} for e in examples]

for i, (eg, pred) in enumerate(zip(examples, predictions)):
    grade = eval_chain.invoke({
        "query": eg["query"],
        "answer": eg["answer"],
        "result": pred["result"]
    })
    print(f"Example {i}: {grade}")

from dotenv import load_dotenv, find_dotenv
import os
from langchain.chat_models import init_chat_model

_ = load_dotenv(find_dotenv())
model_name = os.environ['MODEL_NAME']
api_key = os.environ['API_KEY']
temperature = 0.0

llm = init_chat_model(model=model_name, api_key=api_key, temperature=temperature)

from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper


# Wikipedia 工具
wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())

# DuckDuckGo 搜索工具
from langchain_community.tools import DuckDuckGoSearchRun
duckduckgo = DuckDuckGoSearchRun()

from langchain.agents import create_agent
from langchain_core.messages import HumanMessage

tools = [wikipedia, duckduckgo]

agent = create_agent(
    model=llm,
    tools=tools,
    system_prompt="You are a helpful assistant. Use tools when needed."
)

# 调用方式：传入 messages 列表
result = agent.invoke({
    "messages": [HumanMessage(content="What is the 25% of 300?")]
})
print(result["messages"][-1].content)

result = agent.invoke({
    "messages": [HumanMessage(content="What is the Apple Inc.?")]
})
print(result["messages"][-1].content)

from langchain_core.tools import tool
from datetime import date

@tool
def time(text: str) -> str:
    """Returns today's date. Use this for any questions related to
    knowing today's date. The input should always be an empty string,
    and this function will always return today's date - any date
    mathematics should occur outside this function."""
    return str(date.today())

tools_with_time = [wikipedia, duckduckgo, time]

agent_with_time = create_agent(
    model=llm,
    tools=tools_with_time,
    system_prompt="You are a helpful assistant with access to web search, Wikipedia, and date tools."
)

try:
    result = agent_with_time.invoke({
        "messages": [HumanMessage(content="What's the date today?")]
    })
    print(result["messages"][-1].content)
except Exception as e:
    print(f"exception: {e}")