from langchain_core.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate

# Defining the system prompt (how the AI should act)
system_prompt = SystemMessagePromptTemplate.from_template(
    "You are an AI assistant that helps generate article titles."
)

# the user prompt is provided by the user, in this case however the only dynamic
# input is the article
user_prompt = HumanMessagePromptTemplate.from_template(
    """You are tasked with creating a name for a article.
The article is here for you to examine {article}

The name should be based of the context of the article.

Be creative, but make sure the names are clear, catchy,

and relevant to the theme of the article.

Only output the article name, no other explanation or
text can be provided.""",
    input_variables=["article"]
)
# 这里的 `input_variables` 是一个列表，用来明确声明这个提示模板（Prompt Template）中包含了哪些需要被动态填充的变量

# 使用format方法来对提示词模版中的变量赋值
user_prompt.format(article="TEST STRING")

from langchain_core.prompts import ChatPromptTemplate

# 使用ChatPromptTemplate来组合系统提示词和用户提示词
first_prompt = ChatPromptTemplate.from_messages([system_prompt, user_prompt])

# 同样可以给组合提示词模版中的变量赋值
first_prompt.format(article="TEST STRING")

chain_one = (
    {"article": lambda x: x["article"]} # 
    | first_prompt
    | creative_llm  # <-- 语言模型输出了一个对象
    | {"article_title": lambda x: x.content} # <-- 这一步处理该对象
)

# 传入定义好的article变量
article_title_msg = chain_one.invoke({"article": article})

from pydantic import BaseModel, Field

class Paragraph(BaseModel):

    original_paragraph: str = Field(description="The original paragraph")

    edited_paragraph: str = Field(description="The improved edited paragraph")

    feedback: str = Field(description=(

        "Constructive feedback on the original paragraph"

    ))

structured_llm = creative_llm.with_structured_output(Paragraph)

chain_three = (
    {"article": lambda x: x["article"]}
    | third_prompt
    | structured_llm
    | {
        "original_paragraph": lambda x: x.original_paragraph,
        "edited_paragraph": lambda x: x.edited_paragraph,
        "feedback": lambda x: x.feedback
    }
)

out = chain_three.invoke({"article": article})

import os

# must enter API key
os.environ["LANGCHAIN_API_KEY"]

# below should not be changed
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"

# you can change this as preferred
os.environ["LANGCHAIN_PROJECT"] = "aurelioai-langchain-course-langsmith-openai"

from langsmith import traceable
import random
import time


@traceable
def generate_random_number():
    return random.randint(0, 100)

@traceable
def generate_string_delay(input_str: str):
    number = random.randint(1, 5)
    time.sleep(number)
    return f"{input_str} ({number})"

@traceable
def random_error():
    number = random.randint(0, 1)
    if number == 0:
        raise ValueError("Random error")
    else:
        return "No error"

from langsmith import traceable

@traceable(name="Chitchat Maker")
def error_generation_function(question: str):
    delay = random.randint(0, 3)
    time.sleep(delay)
    number = random.randint(0, 1)
    if number == 0:
        raise ValueError("Random error")
    else:
        return "I'm great how are you?"

[(Rules) For Our Prompt]
Answer the question based on the context below,                 
if you cannot answer the question using the                     
provided information answer with "I don't know"                 

[Context AI has]
Context: Aurelio AI is an AI development studio                 
focused on the fields of Natural Language Processing (NLP)     
and information retrieval using modern tooling                   
such as Large Language Models (LLMs),                           
vector databases, and LangChain.                                

[User Question]
Question: Does Aurelio AI do anything related to LangChain?        

[AI Answer]
Answer:                                                          

prompt = """
Answer the user's query based on the context below.
If you cannot answer the question using the
provided information answer with "I don't know".

Context: {context}
"""

from langchain.prompts import (
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate
)

prompt_template = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(prompt),
    HumanMessagePromptTemplate.from_template("{query}"),
])
# 可以使用prompt_template.messages来获取提示词模版的内容
# 可以使用prompt_template.input_variables来获取提示词模版中的变量
# prompt_template的完整内容如下：
'''
input_variables=['context', 'query'] input_types={} partial_variables={} messages=[SystemMessagePromptTemplate(prompt=PromptTemplate(input_variables=['context'], input_types={}, partial_variables={}, template='\nAnswer the user\'s query based on the context below.\nIf you cannot answer the question using the\nprovided information answer with "I don\'t know".\n\nContext: {context}\n'), additional_kwargs={}), HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['query'], input_types={}, partial_variables={}, template='{query}'), additional_kwargs={})] 
'''
pipeline = (
    {
        "query": lambda x: x["query"],
        "context": lambda x: x["context"]
    }
    | prompt_template
    | llm
)

context = """Aurelio AI is an AI company developing tooling for AI
engineers. Their focus is on language AI with the team having strong
expertise in building AI agents and a strong background in
information retrieval.

The company is behind several open source frameworks, most notably
Semantic Router and Semantic Chunkers. They also have an AI
Platform providing engineers with tooling to help them build with
AI. Finally, the team also provides development services to other
organizations to help them bring their AI tech to market.

Aurelio AI became LangChain Experts in September 2024 after a long
track record of delivering AI solutions built with the LangChain
ecosystem."""

query = "what does Aurelio AI do?"

pipeline.invoke({"query": query, "context": context})

new_system_prompt = """
Answer the user's query based on the context below.
If you cannot answer the question using the
provided information answer with "I don't know".
  
Always answer in markdown format. When doing so please
provide headers, short summaries, follow with bullet
points, then conclude.

Context: {context}
"""

example_prompt = ChatPromptTemplate.from_messages([
    ("human", "{input}"),
    ("ai", "{output}"),
])

examples = [
    {
        "input": "Can you explain gravity?",
        "output": (
            "## Gravity\n\n"
            "Gravity is one of the fundamental forces in the universe.\n\n"
            "### Discovery\n\n"
            "* Gravity was first discovered by Sir Isaac Newton in the late 17th century.\n"
            "* It was said that Newton theorized about gravity after seeing an apple fall from a tree.\n\n"
            "### In General Relativity\n\n"
            "* Gravity is described as the curvature of spacetime.\n"
            "* The more massive an object is, the more it curves spacetime.\n"
            "* This curvature is what causes objects to fall towards each other.\n\n"
            "### Gravitons\n\n"
            "* Gravitons are hypothetical particles that mediate the force of gravity.\n"
            "* They have not yet been detected.\n\n"
            "**To conclude**, Gravity is a fascinating topic and has been studied extensively since the time of Newton.\n\n"
        )
    },
    {
        "input": "What is the capital of France?",
        "output": (
            "## France\n\n"
            "The capital of France is Paris.\n\n"
            "### Origins\n\n"
            "* The name Paris comes from the Latin word \"Parisini\" which referred to a Celtic people living in the area.\n"
            "* The Romans named the city Lutetia, which means \"the place where the river turns\".\n"
            "* The city was renamed Paris in the 3rd century BC by the Celtic-speaking Parisii tribe.\n\n"
            "**To conclude**, Paris is highly regarded as one of the most beautiful cities in the world and is one of the world's greatest cultural and economic centres.\n\n"
        )
    }
]

from langchain_core.prompts import FewShotChatMessagePromptTemplate

few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples,
)

prompt_template = ChatPromptTemplate.from_messages([
    ("system", new_system_prompt),
    few_shot_prompt,
    ("user", "{query}"),
])

pipeline = prompt_template | llm

context = """Aurelio AI is an AI company developing tooling for AI
engineers. Their focus is on language AI with the team having strong
expertise in building AI agents and a strong background in
information retrieval.

The company is behind several open source frameworks, most notably
Semantic Router and Semantic Chunkers. They also have an AI
Platform providing engineers with tooling to help them build with
AI. Finally, the team also provides development services to other
organizations to help them bring their AI tech to market.

Aurelio AI became LangChain Experts in September 2024 after a long
track record of delivering AI solutions built with the LangChain
ecosystem."""

query = "what does Aurelio AI do?"

out = pipeline.invoke({"query": query, "context": context}).content

# return_messages=True使得在输出记忆内容时会附上消息类型，比如HumanMessage
memory = ConversationBufferMemory(return_messages=True)

memory.chat_memory.add_user_message("Hi, my name is James")
memory.chat_memory.add_ai_message("Hey James, what's up? I'm an AI model called Zeta.")
memory.chat_memory.add_user_message("I'm researching the different types of conversational memory.")
memory.chat_memory.add_ai_message("That's interesting, what are some examples?")
memory.chat_memory.add_user_message("I've been looking at ConversationBufferMemory and ConversationBufferWindowMemory.")
memory.chat_memory.add_ai_message("That's interesting, what's the difference?")
memory.chat_memory.add_user_message("Buffer memory just stores the entire conversation, right?")
memory.chat_memory.add_ai_message("That makes sense, what about ConversationBufferWindowMemory?")
memory.chat_memory.add_user_message("Buffer window memory stores the last k messages, dropping the rest.")
memory.chat_memory.add_ai_message("Very cool!")

memory.load_memory_variables({}) # 输出记忆内容

from langchain.chains import ConversationChain

chain = ConversationChain(
    llm=llm,
    memory=memory,
    verbose=True
)

chain.invoke({"input": "what is my name again?"})

from langchain.prompts import (
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
    ChatPromptTemplate
)

system_prompt = "You are a helpful assistant called Zeta."

prompt_template = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(system_prompt),
    MessagesPlaceholder(variable_name="history"),
    # MessagesPlaceholder相当于一个占位符，它可以插入
    # 一系列完整的消息对象（如 `HumanMessage`, `AIMessage`）
    HumanMessagePromptTemplate.from_template("{query}"),
])

pipeline = prompt_template | llm

from langchain_core.chat_history import InMemoryChatMessageHistory

chat_map = {}

def get_chat_history(session_id: str) -> InMemoryChatMessageHistory:
    if session_id not in chat_map:
        # if session ID doesn't exist, create a new chat history
        chat_map[session_id] = InMemoryChatMessageHistory()
    return chat_map[session_id]
    
from langchain_core.runnables.history import RunnableWithMessageHistory

pipeline_with_history = RunnableWithMessageHistory(
    pipeline,
    get_session_history=get_chat_history,
    # 这样langchain就可以使用get_chat_history函数来导入记忆
    input_messages_key="query",
    history_messages_key="history"
    # 这里的"history"对应了上文中的MessagesPlaceholder，它会把记忆中的
    # 内容全部导入到prompt_template中MessagesPlaceholder所占的位置
)

pipeline_with_history.invoke(
    {"query": "Hi, my name is James"},
    config={"session_id": "id_123"}
    # config参数告诉一个RunnableWithMessageHistory对象该如何使用
    # get_session_hitory中的这个可调用对象，相当于get_chat_history("id_123")

)

from langchain_core.tools import tool

# 创建工具时一定要把工具功能写清楚，否则大模型可能不知道这个工具的功能
@tool
def add(x: float, y: float) -> float:
    """Add 'x' and 'y'."""
    return x + y

@tool
def multiply(x: float, y: float) -> float:
    """Multiply 'x' and 'y'."""
    return x * y

@tool
def exponentiate(x: float, y: float) -> float:
    """Raise 'x' to the power of 'y'."""
    return x ** y

@tool
def subtract(x: float, y: float) -> float:
    """Subtract 'x' from 'y'."""
    return y - x

from langchain.memory import ConversationBufferMemory

memory = ConversationBufferMemory(
    memory_key="chat_history",  # must align with MessagesPlaceholder variable_name
    return_messages=True  # to return Message objects
)

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

prompt = ChatPromptTemplate.from_messages([
    ("system", "you're a helpful assistant"),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])

from langchain.agents import create_tool_calling_agent

tools = [add, subtract, multiply, exponentiate]

agent = create_tool_calling_agent(
    llm=llm, tools=tools, prompt=prompt
)

from langchain.agents import AgentExecutor

agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    memory=memory,
    verbose=True
)