from typing import Optional

from langchain_core.pydantic_v1 import BaseModel, Field

class Person(BaseModel):
    """Information about a person."""

    # ^ Doc-string for the entity Person.
    # This doc-string is sent to the LLM as the description of the schema Person,
    # and it can help to improve extraction results.

    # Note that:
    # 1. Each field is an `optional` -- this allows the model to decline to extract it!
    # 2. Each field has a `description` -- this description is used by the LLM.
    # Having a good description can help improve extraction results.
    name: Optional[str] = Field(..., description="The name of the person")
    hair_color: Optional[str] = Field(
        ..., description="The color of the peron's hair if known"
    )
    height_in_meters: Optional[str] = Field(
        ..., description="Height measured in meters"
    )

from typing import List, Optional

from langchain_core.pydantic_v1 import BaseModel, Field


class Person(BaseModel):
    """Information about a person."""

    # ^ Doc-string for the entity Person.
    # This doc-string is sent to the LLM as the description of the schema Person,
    # and it can help to improve extraction results.

    # Note that:
    # 1. Each field is an `optional` -- this allows the model to decline to extract it!
    # 2. Each field has a `description` -- this description is used by the LLM.
    # Having a good description can help improve extraction results.
    name: Optional[str] = Field(..., description="The name of the person")
    hair_color: Optional[str] = Field(
        ..., description="The color of the peron's hair if known"
    )
    height_in_meters: Optional[str] = Field(
        ..., description="Height measured in meters"
    )


class Data(BaseModel):
    """Extracted data about people."""

    # Creates a model so that we can extract multiple entities.
    people: List[Person]

class Book(BaseModel):
    book_name: Optional[str] = Field(..., description="书的名称")
    writer: Optional[str] = Field( ..., description="书的作者")
    content:Optional[str] = Field( ..., description="书大概内容介绍")
    date :Optional[str] = Field( ..., description="预计出版时间")


class Data(BaseModel):
    books: List[Book]

examples = [
        (
            """
            《美棠来信：我们一家人》
作者：饶平如

出品方/出版社：小阅读Random

预计出版时间：2023年6月

（书封待定，图为作者饶平如）

继《平如美棠》后，《美棠来信》更加细致地呈现出家书中一个中国普通家庭的记忆。1973年至1979年，饶平如下放安徽，毛美棠留在上海照顾家庭。不久，家中年长的孩子们也响应知识青年上山下乡，去安徽、江西等地下乡。分散在各地的家庭成员，唯一连接他们亲情的，是一封封往来两地的家书。

本书收录了饶平如在1973年到1979年之间收到的来自妻子美棠和孩子的近两百封家书。在那个通讯不便的年代，在这些家书里，他们互相汇报生活近况，通报生活上遇到的困难，给对方出谋划策，嘘寒问暖，这些家书支撑他们度过了两地分散的艰难时期，同时也是那个年代一个普通中国家庭生活的真实侧写。
            """,
            Book(book_name='美棠来信：我们一家人', writer='饶平如', content="""继《平如美棠》后，《美棠来信》更加细致地呈现出家书中一个中国普通家庭的记忆。1973年至1979年，饶平如下放安徽，毛美棠留在上海照顾家庭。不久，家中年长的孩子们也响应知识青年上山下乡，去安徽、江西等地下乡。分散在各地的家庭成员，唯一连接他们亲情的，是一封封往来两地的家书。

本书收录了饶平如在1973年到1979年之间收到的来自妻子美棠和孩子的近两百封家书。在那个通讯不便的年代，在这些家书里，他们互相汇报生活近况，通报生活上遇到的困难，给对方出谋划策，嘘寒问暖，这些家书支撑他们度过了两地分散的艰难时期，同时也是那个年代一个普通中国家庭生活的真实侧写""",date='2023年6月'),
        ),

    ]

{'id': '932e3864-db6b-4888-a1ba-566609d6530d', 'type': 'function', 'function': {'name': 'Book', 'arguments': '{"book_name": "美棠来信：我们一家人", "writer": "饶平如", "content": "继《平如美棠》后，《美棠来信》更加细致..", "date": "2023年6月"}'}}

messages = []

for text, tool_call in examples:
    messages.extend(
        tool_example_to_messages({"input": text, "tool_calls": [tool_call]})
    )

  prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "您是一个专业的提取信息专家。 "
            "只从文本中提取相关信息。"
            "如果你不知道请求提取的属性值，请将该属性值返回为null。",
        ),
        # ↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓
        MessagesPlaceholder("examples"),  # <-- EXAMPLES!
        # ↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑
        ("human", "{text}"),
    ]
)  
  
# Azure Openai
llm = AzureChatOpenAI(
    openai_api_version="2024-02-15-preview",
    azure_deployment=os.getenv('DEPLOYMENT_NAME_GPT3_4K_JP'),
    temperature=0,
)

runnable = prompt | llm.with_structured_output(
    schema=Data,
    method="function_calling",
    include_raw=False,
)

    for _ in range(5):
        text = """
《浮生余情》（暂名）

作者：格非

出品方/出版社：译林出版社

预计出版时间：2023年8月

（书封待定，图为作者格非）

《浮生余情》是格非时隔四年推出的最新重磅长篇小说，叙写了1980年代至今四十余年的漫长时间里，四个彼此关联的人物的命运流转。

四个故事的主人公分别从江浙、北京、甘肃和天津四个地方来到北京的春台路21号——位于后厂村的中关村软件园。他们均供职于同一家现代物联网企业，彼此在生活中多有交集。四个故事由亲情、爱情的本能情感走向对自我与他者关系建构的哲学思考，进而关注现代人当下的存在方式与情感命题。

《大地上的家乡》

作者：刘亮程

出品方/出版社：译林出版社

预计出版时间：2023年6月

（书封待定，图为作者刘亮程）
...
        """
        print(runnable.invoke({"text": text, "examples": messages}))
        pass

books=[Book(book_name='浮生余情', writer='格非', content='《浮生余情》是格非时隔四年推出的最新重磅长篇小说，叙写了1980年代至今四十余年的漫长时间里，四个彼此关联的人物的命运流转。', date='2023年8月'), Book(book_name='大地上的家乡', writer='刘亮程', content='1998年，刘亮程站在乌鲁木齐的夕阳中，回望自己的家乡黄沙梁，写就《一个人的村庄》。此后，他在城市结婚、生子、写作、生活。2013年，刘亮程入住新疆木垒书院菜籽沟村落，重返晴耕雨读的田园生活，仿佛又回到早年的鸡鸣狗吠、虫鸣鸟语、风声落叶中，进入写作《一个人的村庄》时的状态。', date='2023年6月'), Book(book_name='俗世奇人新增本', writer='冯骥才', content='“俗世奇人”系列是当代文化大家冯骥才先生的代表作，也是当代文学的重要收获之一。天津卫本是水陆码头，居民五方杂处，冯骥才随想随记，每人一篇，冠之总名《俗世奇人》耳。\n\n2022年末，冯先生又创作了18篇“俗世奇人”系列新作，包括《欢喜》《小尊王五》《田大头》《谢二虎》等精彩作品，还专程为书中人物绘制了20余幅精美人物插画。18篇新作延续了之前一贯的写作手法、传奇风格和创作水准，艺术性层面更为灵动丰富，令人不忍释卷。', date='2023年1月'), Book(book_name='四', writer='杨本芬', content='杨本芬奶奶的第四本书，收入四个中短篇，写四个不同的人。对妈妈的回忆，对哥哥的眷恋，以及，哪怕是一位农妇，一个捡垃圾的老太太，依然拥有的，身为女性的，骄傲。也许，这本书可以叫《我的思念，她的骄傲》。', date='2023年7月'), Book(book_name='木星时刻', writer='李静睿', content='李静睿的最新短篇小说集，收录具有科幻风格的《木星时刻》、充满粗粝现实感的《温榆河》等8部短篇小说。\n\n李静睿在《木星时刻》里构建了一个依法由AI治理的未来社会，人们在其中享有健康的饮食和生活方式、绝对安全的生存环境，以及完美的AI管家，唯独没有自由。于是，作为“最后一代上学还能逃课的人类”的主人公夫妻二人，决心踏上一场“肖申克”式的逃亡之旅……', date='2023年5月')]

import os

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from typing import List, Optional

from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI, AzureChatOpenAI
import uuid
from typing import Dict, List, TypedDict

from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langchain_core.pydantic_v1 import BaseModel, Field

if __name__ == '__main__':

    #
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "您是一个专业的提取信息专家。 "
                "只从文本中提取相关信息。"
                "如果你不知道请求提取的属性值，请将该属性值返回为null。",
            ),
            # ↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓
            MessagesPlaceholder("examples"),  # <-- EXAMPLES!
            # ↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑
            ("human", "{text}"),
        ]
    )


    class Book(BaseModel):
        book_name: Optional[str] = Field(..., description="书的名称")
        writer: Optional[str] = Field( ..., description="书的作者")
        content:Optional[str] = Field( ..., description="书大概内容介绍")
        date :Optional[str] = Field( ..., description="预计出版时间")


    class Data(BaseModel):
        books: List[Book]


    class Example(TypedDict):
        """A representation of an example consisting of text input and expected tool calls.

        For extraction, the tool calls are represented as instances of pydantic model.
        """

        input: str  # This is the example text
        tool_calls: List[BaseModel]  # Instances of pydantic model that should be extracted


    def tool_example_to_messages(example: Example) -> List[BaseMessage]:
        messages: List[BaseMessage] = [HumanMessage(content=example["input"])]
        openai_tool_calls = []
        for tool_call in example["tool_calls"]:
            openai_tool_calls.append(
                {
                    "id": str(uuid.uuid4()),
                    "type": "function",
                    "function": {
                        "name": tool_call.__class__.__name__,
                        "arguments": tool_call.json(),
                    },
                }
            )
        messages.append(
            AIMessage(content="", additional_kwargs={"tool_calls": openai_tool_calls})
        )
        tool_outputs = example.get("tool_outputs") or [
            "You have correctly called this tool."
        ] * len(openai_tool_calls)
        for output, tool_call in zip(tool_outputs, openai_tool_calls):
            messages.append(ToolMessage(content=output, tool_call_id=tool_call["id"]))
        return messages


    examples = [
        (
            """
            《美棠来信：我们一家人》
作者：饶平如

出品方/出版社：小阅读Random

预计出版时间：2023年6月

（书封待定，图为作者饶平如）

继《平如美棠》后，《美棠来信》更加细致地呈现出家书中一个中国普通家庭的记忆。1973年至1979年，饶平如下放安徽，毛美棠留在上海照顾家庭。不久，家中年长的孩子们也响应知识青年上山下乡，去安徽、江西等地下乡。分散在各地的家庭成员，唯一连接他们亲情的，是一封封往来两地的家书。

本书收录了饶平如在1973年到1979年之间收到的来自妻子美棠和孩子的近两百封家书。在那个通讯不便的年代，在这些家书里，他们互相汇报生活近况，通报生活上遇到的困难，给对方出谋划策，嘘寒问暖，这些家书支撑他们度过了两地分散的艰难时期，同时也是那个年代一个普通中国家庭生活的真实侧写。
            """,
            Book(book_name='美棠来信：我们一家人', writer='饶平如', content="""继《平如美棠》后，《美棠来信》更加细致地呈现出家书中一个中国普通家庭的记忆。1973年至1979年，饶平如下放安徽，毛美棠留在上海照顾家庭。不久，家中年长的孩子们也响应知识青年上山下乡，去安徽、江西等地下乡。分散在各地的家庭成员，唯一连接他们亲情的，是一封封往来两地的家书。

本书收录了饶平如在1973年到1979年之间收到的来自妻子美棠和孩子的近两百封家书。在那个通讯不便的年代，在这些家书里，他们互相汇报生活近况，通报生活上遇到的困难，给对方出谋划策，嘘寒问暖，这些家书支撑他们度过了两地分散的艰难时期，同时也是那个年代一个普通中国家庭生活的真实侧写""",date='2023年6月'),
        ),

    ]

    messages = []

    for text, tool_call in examples:
        messages.extend(
            tool_example_to_messages({"input": text, "tool_calls": [tool_call]})
        )

    # Azure Openai
    llm = AzureChatOpenAI(
        openai_api_version="2024-02-15-preview",
        azure_deployment=os.getenv('DEPLOYMENT_NAME_GPT3_4K_JP'),
        temperature=0,
    )

    runnable = prompt | llm.with_structured_output(
        schema=Data,
        method="function_calling",
        include_raw=False,
    )

    for _ in range(5):
        text = """
《浮生余情》（暂名）

作者：格非

出品方/出版社：译林出版社

预计出版时间：2023年8月

（书封待定，图为作者格非）

《浮生余情》是格非时隔四年推出的最新重磅长篇小说，叙写了1980年代至今四十余年的漫长时间里，四个彼此关联的人物的命运流转。

四个故事的主人公分别从江浙、北京、甘肃和天津四个地方来到北京的春台路21号——位于后厂村的中关村软件园。他们均供职于同一家现代物联网企业，彼此在生活中多有交集。四个故事由亲情、爱情的本能情感走向对自我与他者关系建构的哲学思考，进而关注现代人当下的存在方式与情感命题。

《大地上的家乡》

作者：刘亮程

出品方/出版社：译林出版社

预计出版时间：2023年6月

（书封待定，图为作者刘亮程）

1998年，刘亮程站在乌鲁木齐的夕阳中，回望自己的家乡黄沙梁，写就《一个人的村庄》。此后，他在城市结婚、生子、写作、生活。2013年，刘亮程入住新疆木垒书院菜籽沟村落，重返晴耕雨读的田园生活，仿佛又回到早年的鸡鸣狗吠、虫鸣鸟语、风声落叶中，进入写作《一个人的村庄》时的状态。

菜籽沟村堆满故事，早晨做梦的气味被一只狗闻见，在一棵大树下慢慢变老，散步于开满窗户的山坡，看从天坑外背土豆的人……这些飘在空中被人视若寻常的故事，均收在了新书《大地上的家乡》里。

《俗世奇人新增本》

作者: 冯骥才

出品方/出版社: 人民文学出版社

预计出版时间：2023年1月

“俗世奇人”系列是当代文化大家冯骥才先生的代表作，也是当代文学的重要收获之一。天津卫本是水陆码头，居民五方杂处，性格迥然相异，冯骥才随想随记，每人一篇，冠之总名《俗世奇人》耳。

2022年末，冯先生又创作了18篇“俗世奇人”系列新作，包括《欢喜》《小尊王五》《田大头》《谢二虎》等精彩作品，还专程为书中人物绘制了20余幅精美人物插画。18篇新作延续了之前一贯的写作手法、传奇风格和创作水准，艺术性层面更为灵动丰富，令人不忍释卷。

《四》（书名待定）

作者：杨本芬

出品方/出版社：乐府文化

预计出版时间：2023年7月

（书封待定，图为作者杨本芬）

杨本芬奶奶的第四本书，收入四个中短篇，写四个不同的人。对妈妈的回忆，对哥哥的眷恋，以及，哪怕是一位农妇，一个捡垃圾的老太太，依然拥有的，身为女性的，骄傲。也许，这本书可以叫《我的思念，她的骄傲》。

《木星时刻》

作者：李静睿

出品方/出版社：小阅读Random

预计出版时间：2023年5月

（书封待定，上图为作者李静睿）

李静睿的最新短篇小说集，收录具有科幻风格的《木星时刻》、充满粗粝现实感的《温榆河》等8部短篇小说。

李静睿在《木星时刻》里构建了一个依法由AI治理的未来社会，人们在其中享有健康的饮食和生活方式、绝对安全的生存环境，以及完美的AI管家，唯独没有自由。于是，作为“最后一代上学还能逃课的人类”的主人公夫妻二人，决心踏上一场“肖申克”式的逃亡之旅……
        """
        print(runnable.invoke({"text": text, "examples": messages}))
        pass

import json
import os

from langchain.chains.llm import LLMChain
from langchain_community.chat_models.baidu_qianfan_endpoint import QianfanChatEndpoint
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate

def parse_or_fix(text: str,chat):
    fixing_chain = (

            ChatPromptTemplate.from_template(
                "根据错误修改text中的文本内容:\n\n```text\n{input}\n```\n错误: {error}"
                "只返回修改后的text，千万不要解释，不要说明，不要添加```json"
            )

            | chat
            | StrOutputParser()
    )
    for _ in range(3):
        try:
            filtered_text = text.replace('```json\n', '').replace('```', '')
            return json.loads(filtered_text)
        except Exception as e:
            text = fixing_chain.invoke({"input": filtered_text, "error": e})
            print("百度修正结果： = [{}]".format(text))

    print('百度修正失败！')
    return None

if __name__ == '__main__':
    
    os.environ["QIANFAN_ACCESS_KEY"] = os.getenv('MY_QIANFAN_ACCESS_KEY')
    os.environ["QIANFAN_SECRET_KEY"] = os.getenv('MY_QIANFAN_SECRET_KEY')
    
    chat = QianfanChatEndpoint(model="ERNIE-Bot-4")

    examples = """
            <<曾国藩传>>
            作者:AAA
            出版时间:2000年11月3日
            内容:曾国藩（1811-1872年），湖南湘乡人，初名子城，字伯函，号涤生，是中国历史上最有影响的人物之一。
    他的人生，他的智慧，他的思想，深深地影响了几代中国人，以至他虽已去世一百余年，提起曾国藩，人们仍然津津乐道。
    有的评论者说：如果以人物断代的话，曾国藩是中国古代历史上的最后一人，近代历史上的第一人。
    这句话从某一角度，概括了曾国藩的个人作用和影响。
    曾国藩出生于晚清一个地主家庭，自幼勤奋好学，6岁入塾读书。8岁能读八股文、诵五经，14岁能读《周礼》《史记》文选，
    同年参加长沙的童子试，成绩列为优等。父麟书，有田产，不事耕种，醉心功名，然童试17次皆不第，父设馆授徒。
    曾国藩幼从父学。道光十三年（1833）入县学为秀才。翌年就读于长沙岳麓书院，同年中举人。
    此后赴京会试，一再落榜。十八年，始中第三十八名贡士，旋赴殿试，中三甲第四十二名，赐同进士出身。
    朝考选翰林院庶吉士。自此供职京师，结交穆彰阿、倭仁及唐鉴等。二十七年任四川乡试正考官，二十八年升侍读，
    后年升侍讲学士。二十七年授内阁学士，兼礼部侍郎衔。二十九年任礼部右侍郎，旋兼兵部右侍郎。
    三十年兼署工部右侍郎。咸丰二年（1852）兼署吏部左侍郎。
    后丁忧在湘乡老家，此时奉诏以礼部侍郎身份帮同湖南巡抚督办团练，创建湘军。
    """

    # 摘要提示词
    prompt_template = """
        你是一个文本分析专家,你抽取文本中的书籍信息。
        抽取的内容包括书名，作者，出版时间，内容评述。
        你需要按照下面的格式输出，输出格式为JSON：
        不要返回```json 和 多余内容。
        book_name：
        writer：
        time：
        content：
        comments:
        {text}
        你的结果:"""
    edit_prompt = PromptTemplate.from_template(prompt_template)
    edit_chain = LLMChain(llm=chat, prompt=edit_prompt)
    result = edit_chain.run(text=examples)
    print(result)

    # 字符串转JSON
    json_obj = parse_or_fix(result,chat)
    if json_obj:
        print('解析成功')
        print(json_obj)
    else:
        print('解析失败')

#Schema
class Person(BaseModel):
    age: Optional[int] = Field(None, description="The age of the person in years.")
    name: Optional[str] = Field(None, description="The name of the person.")
    nick_name: Optional[str] = Field(None, description="Alias, if any.")


runnable = RemoteRunnable("http://localhost:8000/extract_text/")

text = """
My name is Chester. i am 42 years old. My friend Jane is a year older than me.
"""

response = runnable.invoke({"text": text, "schema": Person.schema()})
print(response)

python -m scripts.run_migrations create 
All tables created successfully.

def get_postgres_url() -> URL:
    url = URL.create(
        drivername="postgresql",
        username="xxxx", # 修改
        password="xxxxx", # 修改
        host=os.environ.get("PG_HOST", "localhost"),
        database="langchain",
        port=5432,
    )
    return url
  
  def get_model() -> ChatOpenAI:
    """Get the model."""
    model = xxxx # 修改
    # return ChatOpenAI(model=MODEL_NAME, temperature=0)
    return model
  
  async def extract_from_content()
    # 修改
    # vectorstore = FAISS.from_texts(doc_contents, embedding=OpenAIEmbeddings())
    # retriever = vectorstore.as_retriever()

# 我的环境里已经设置相关 环境变量
python -m server.main