# 创建一个llm类型的model并开始询问
from langchain_openai import OpenAI

llm = OpenAI(openai_api_key="...")
llm.invoke("1 + 1 = ?")

# 创建一个chat model类型的model并开始对话
from langchain_openai import ChatOpenAI

chat = ChatOpenAI(openai_api_key="...")

from langchain_core.messages import HumanMessage, SystemMessage
messages = [
    SystemMessage(content="You're good at math"),
    HumanMessage(content="1 + 1 = ?"),
]

chat.invoke(messages)

from langchain.prompts import PromptTemplate

# ====================== PromptTemplate ======================

# Prompt中有多个变量
prompt_template = PromptTemplate.from_template(
    "Tell me a {adjective} joke about {content}."
)
# 输出: Tell me a funny joke about math.
print(prompt_template.format(adjective="funny", content="math"))

# Prompt中没有任何变量
prompt_template = PromptTemplate.from_template("Tell me a joke")

# 输出: Tell me a joke
prompt = prompt_template.format()
print(prompt)

# 调用llm模型输出: Why couldn't the bicycle stand up by itself? Because it was two-tired.
from langchain_openai import OpenAI
llm = OpenAI()
print(llm.invoke(prompt))

# ====================== ChatPromptTemplate ======================

from langchain_core.prompts import ChatPromptTemplate

# 第1种
chat_template = ChatPromptTemplate.from_messages(
    [
        ("system", "You are good at math. Your name is {name}."),
        ("human", "please help me, my question is: {user_input}"),
    ]
)

# 输出: [SystemMessage(content='You are good at math. Your name is GaoSi.'), HumanMessage(content='please help me, my question is: 1 + 1 = ?')]
messages = chat_template.format_messages(name="GaoSi", user_input="1 + 1 = ?")
print(messages)

# 第2种
from langchain.prompts import HumanMessagePromptTemplate
from langchain_core.messages import SystemMessage
chat_template = ChatPromptTemplate.from_messages(
    [
        SystemMessage(
            content=(
                "You are good at math."
                "Your name is GaoSi."
            )
        ),
        HumanMessagePromptTemplate.from_template("please help me, my question is: {user_input}"),
    ]
)

# 输出: [SystemMessage(content='You are good at math.Your name is GaoSi.'), HumanMessage(content='please help me, my question is: 1 + 1 = ?')]
messages = chat_template.format_messages(user_input="1 + 1 = ?")
print(messages)

# 调用chat模型输出: content='1 + 1 = 2'
from langchain_openai import ChatOpenAI
chat = ChatOpenAI()
print(chat.invoke(messages))

# 创建Model
from langchain_openai import ChatOpenAI
model = ChatOpenAI()

# 创建output_parser(输出)
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.pydantic_v1 import BaseModel, Field
class MathProblem(BaseModel):
    question: str = Field(description="the question")
    answer: str = Field(description="the answer of question")
    steps: str = Field(description="the resolve steps of question")

output_parser = JsonOutputParser(pydantic_object=MathProblem)

# 创建prompt(输入)
from langchain.prompts import PromptTemplate
prompt = PromptTemplate(
    template="You are good at math, please answer the user query.\n{format_instructions}\n{query}\n",
    input_variables=["query"],
    partial_variables={"format_instructions": output_parser.get_format_instructions()},
)

# 创建Chain并链式调用
chain = prompt | model | output_parser
print(chain.invoke({"query": "1+1=?"}))

# 使用流式输出
for s in chain.stream({"query": "1+1=?"}):
    print(s)

from langchain.chains import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain_openai import OpenAI

template = """You are a chatbot having a conversation with a human.

{chat_history}
Human: {human_input}
Chatbot:"""

prompt = PromptTemplate(
    input_variables=["chat_history", "human_input"], template=template
)
memory = ConversationBufferMemory(memory_key="chat_history")

llm = OpenAI()
llm_chain = LLMChain(
    llm=llm,
    prompt=prompt,
    verbose=True,
    memory=memory,
)
# Hello there! How are you?
print(llm_chain.predict(human_input="Hi, my friend"))

from langchain.agents import AgentExecutor, Tool, ZeroShotAgent
from langchain.chains import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain_community.utilities import GoogleSearchAPIWrapper
from langchain_openai import OpenAI

# 定义Tool
# 需要定义环境变量 export GOOGLE_API_KEY="", 在网站上注册并生成API Key: https://serpapi.com/searches
search = GoogleSearchAPIWrapper()
tools = [
    Tool(
        name="Search",
        func=search.run,
        description="useful for when you need to answer questions about current events",
    )
]

# 定义Prompt
prefix = """Have a conversation with a human, answering the following questions as best you can. You have access to the following tools:"""
suffix = """Begin!"

{chat_history}
Question: {input}
{agent_scratchpad}"""

prompt = ZeroShotAgent.create_prompt(
    tools,
    prefix=prefix,
    suffix=suffix,
    input_variables=["input", "chat_history", "agent_scratchpad"],
)

# 定义Memory
memory = ConversationBufferMemory(memory_key="chat_history")

# 定义LLMChain
llm_chain = LLMChain(llm=OpenAI(temperature=0), prompt=prompt)

# 定义Agent
agent = ZeroShotAgent(llm_chain=llm_chain, tools=tools, verbose=True)
agent_chain = AgentExecutor.from_agent_and_tools(
    agent=agent, tools=tools, verbose=True, memory=memory
)
agent_chain.run(input="How many people live in canada?")

from langchain.tools import tool

@tool
def search(query: str) -> str:
    """Look up things online."""
    return "LangChain"

print(search)

from langchain_core.utils.function_calling import convert_to_openai_function
tools = [search]
functions = [convert_to_openai_function(t) for t in tools]

from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
model = ChatOpenAI(model="gpt-3.5-turbo")
message = model.invoke(
    [HumanMessage(content="hello, my friend ! i want to search dog")], functions=functions
)
print(message)

from langchain_community.document_loaders import TextLoader

# 加载markdown文档
loader = TextLoader("./index.md")
loader.load()

# 定义一个自定义的回调函数
from typing import Any, Dict
from langchain.callbacks.base import BaseCallbackHandler
class MyCustomHandler(BaseCallbackHandler):
    def on_chain_start(
            self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """Print out that we are entering a chain."""
        class_name = serialized.get("name", serialized.get("id", ["<unknown>"])[-1])
        print(f"\n\n\033[1m> [This is custom callback handler] Entering new {class_name} chain...\033[0m")

# 创建llm/prompt/handler
from langchain_openai import OpenAI
from langchain.prompts import PromptTemplate
llm = OpenAI()
prompt = PromptTemplate.from_template("1 + {number} = ")
handler = MyCustomHandler()

# 链式调用
from langchain.chains import LLMChain
chain = LLMChain(llm=llm, prompt=prompt, callbacks=[handler], verbose=True)
chain.invoke({"number": 2})

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

prompt = ChatPromptTemplate.from_template("tell me a short joke about {topic}")
model = ChatOpenAI()
output_parser = StrOutputParser()

chain = prompt | model | output_parser

print(chain.invoke({"topic": "math"}))

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

prompt = ChatPromptTemplate.from_template("tell me a short joke about {topic}")
model = ChatOpenAI()
output_parser = StrOutputParser()

chain = prompt | model | output_parser

# invoke: 普通输出
print(chain.invoke({"topic": "math"}))

# ainvoke: 异步输出
chain.ainvoke({"topic": "math"})

# stream: 流式输出
for chunk in chain.stream({"topic": "math"}):
    print(chunk, end="", flush=True)

# Batch: 批量输入
print(chain.batch([{"topic": "math"}, {"topic": "English"}]))