from langchain_core.tools import tool
from pydantic.v1 import BaseModel, Field

class PythonFileInput(BaseModel):
    # Описание параметров определения
    filename: str = Field(description="filename")
    source_code: str = Field(description="source code data")

class PytestFileName(BaseModel):
    # Описание параметров определения
    filename: str = Field(description="The name of the file to be executed")

@tool(args_schema=PythonFileInput)
def write_file(filename, source_code):
    """
    Generate python files based on input source code
    """
    with open(filename, "w") as f:
        f.write(source_code)


@tool(args_schema=PytestFileName)
def execute_test_file(filename):
    """
    Pass in the file name, execute the test case and return the execution result
    """
    import subprocess
    # Выполните команду pytest, используя модуль подпроцесса
    result = subprocess.run(['pytest', filename], capture_output=True, text=True)
    # Проверьте результаты выполнения pytest
    if result.returncode == 0:
        print("Тест пройден успешно!")
    else:
        print("Тестовый запуск не пройден:")
    print(result.stdout)
    return result.stdout

#  Уведомление: Необходимо добавить на основе исходного подсказки слово {code} переменная
# prompt = hub.pull("hwchase17/structured-chat-agent")
llm = ChatOpenAI()

agent1 = create_structured_chat_agent(llm, tools_all, prompt)

agent_executor = AgentExecutor(
    agent=agent1, tools=tools_all,
    verbose=True,
    return_intermediate_steps=True,
    handle_parsing_errors=True)

if __name__ == '__main__':
    agent_executor.invoke({"input": «Пожалуйста, сгенерируйте файлы на основе приведенного выше исходного кода», "code": """def test_demo(): return True"""})

if __name__ == '__main__':
    agent_executor.invoke({"input": """
               Пожалуйста, следуйте инструкциям ниже, чтобы завершить операцию, и я позволю вам завершить операцию. Вы не можете остановиться, пока не будут выполнены все шаги:
                1. Сначала сгенерируйте документ на основе приведенного выше исходного кода.
                2. На основе файла исходного кода, созданного на предыдущем шаге, выполните Выполнение. тестовых случае операции и вернуть окончательный результат выполнения
                """,
               "code": """def test_demo(): return True"""})

# load_case Возвращаемый результат — автоматизированный тестовый пример интерфейса.
chain = (
        RunnablePassthrough.assign(code=load_case) | agent1
)

agent_executor = AgentExecutor(
    agent=chain, tools=tools_all,
    verbose=True,
    return_intermediate_steps=True,
    handle_parsing_errors=True)

if __name__ == '__main__':
    agent_executor.invoke({"input": """
               Пожалуйста, следуйте инструкциям ниже, чтобы завершить операцию, и я позволю вам завершить операцию. Вы не можете остановиться, пока не будут выполнены все шаги:
                1. Сначала сгенерируйте документ на основе приведенного выше исходного кода.
                2. На основе файла исходного кода, созданного на предыдущем шаге, выполните Выполнение. тестовых случае операции и вернуть окончательный результат выполнения
                """})