from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores.faiss import FAISS
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain import hub
from plantuml import PlantUML

# Объявить модель
llm = ChatOpenAI()
# # 1. Прочтите файл.
loader = TextLoader("./Требования Document.md")
data = loader.load()
# 3. embedding
embeddings = OpenAIEmbeddings()
# # 4. векторное хранилище
vector = FAISS.from_documents(data, embeddings)
retriever = vector.as_retriever()

from langchain.tools.retriever import create_retriever_tool

retriever_tool = create_retriever_tool(
    retriever,
    "search_demand",
    "Найти документ с требованиями середина, в котором указаны требования",
)
@tool
def generate_png(uml_code, filename):
    """Введите plantumlкод, чтобы сгенерировать изображение и сохранить его в файл"""
    plantuml = PlantUML(url='https://plantuml.ceshiren.com/img/')
    image_bytes = plantuml.processes(uml_code)
    with open(f'{filename}.png', 'wb') as f:
        f.write(image_bytes)

tools = [retriever_tool, generate_png]
llm_with_tools = llm.bind_tools(tools)

prompt = hub.pull("hwchase17/openai-tools-agent")
agent = create_openai_tools_agent(llm, tools, prompt, )
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
result = agent_executor.invoke({
    "input": """Я инженер-испытатель. Мне нужно разобраться в информации о требованиях из приведенного выше документа с требованиями. Пожалуйста, помогите мне разобраться во всех требованиях."
             «Первый уровень интеллект-карты — это информация заголовка, начиная с версии 4.x документа с требованиями середина, который представляет функциональный модуль, а второй уровень — это тестовая точка функционального модуля»,
             «Пожалуйста, сначала выведите один plantuml Исходный код формата, формат исходного кода показан в коде.
            @startmindmap
            * root node
                * some first level node
                    * second level node
                    * another second level node
                * another first level node
            @endmindmap
             Затем выведите файл интеллект-карты в формате plantuml на основе информации исходного кода. Создайте файл изображения с именем файла Хогвартс добавляет любое случайное число """
})