Spaces:
Running
Running
import os | |
import gradio as gr | |
from transformers import ReactCodeAgent, HfEngine, Tool | |
from gradio_agentchatbot import ( | |
AgentChatbot, | |
stream_from_transformers_agent, | |
ChatMessage, | |
) | |
from dotenv import load_dotenv | |
from huggingface_hub import login | |
from transformers.agents.default_tools import ( | |
BASE_PYTHON_TOOLS, | |
LIST_SAFE_MODULES, | |
evaluate_python_code, | |
) | |
# to load SerpAPI key | |
load_dotenv() | |
login(os.getenv("HUGGINGFACEHUB_API_TOKEN")) | |
llm_engine = HfEngine(model="meta-llama/Meta-Llama-3-70B-Instruct") | |
authorized_imports = ["numpy"] | |
agent = ReactCodeAgent( | |
llm_engine=llm_engine, | |
tools=[], | |
additional_authorized_imports=authorized_imports, | |
max_iterations=10, | |
) | |
class FinalAnswerToolWithVerification(Tool): | |
name = "final_answer" | |
description = "Provides a final answer to the given problem" | |
inputs = { | |
"answer": {"type": "text", "description": "The final answer to the problem"} | |
} | |
output_type = "any" | |
def forward(self, answer): | |
if "def test" not in answer: | |
raise Exception( | |
"I can only accept from you a code snippet answer that defines test functions in python, anything else will not work. PLEASE PROVIDE ME A FULL CODE SNIPPET CONTAINING THE DEFINITION OF THE TESTS." | |
) | |
return answer | |
final_answer_tool = FinalAnswerToolWithVerification() | |
agent._toolbox.update_tool(final_answer_tool) | |
function = """import numpy as np | |
def moving_average(x, w): | |
return np.convolve(x, np.ones(w), 'valid') / w""" | |
task = """I will give you a basic function that I've created. | |
Now I want you to generate a set of unit tests for these functions, check that they run, and give them to me. | |
Please follow these steps in order: | |
1. Define and run the function given o you, so that it gets defined in your interpreter. | |
2. Generate one test function as a python blob, with assert statements | |
3. Run the test function in a code snippet and make sure the tests pass | |
4. Return to me the complete TEST function (not the original function) as a string code snippet. | |
--- | |
Example: | |
Here is your function: | |
```py | |
def get_even_numbers(numbers): | |
even_numbers = [] | |
for number in numbers: | |
if number % 2 == 0: | |
even_numbers.append(number) | |
return even_numbers | |
``` | |
Now generate test functions for me! | |
Thought: Let's re-define the given function and generate a test. | |
Code: | |
```py | |
def get_even_numbers(numbers): | |
even_numbers = [] | |
for number in numbers: | |
if number % 2 == 0: | |
even_numbers.append(number) | |
return even_numbers | |
def test_get_even_numbers(): | |
assert get_even_numbers([1, 2, 3, 4, 5]) == [2, 4] | |
print("No error found!") | |
test_get_even_numbers() | |
``` | |
Observation: "No error found!" | |
Thought: the interpreter ran tests with no error. So we can return the function IN A TEXT SNIPPET. | |
Code: | |
```py | |
fianl_answer_snippet = \"\"\" | |
def test_get_even_numbers(): | |
assert get_even_numbers([1, 2, 3, 4, 5]) == [2, 4] | |
print("No error found!") | |
\"\"\" | |
final_answer(final_answer_snippet) | |
``` | |
--- | |
Now proceed! | |
Here is your function: | |
```py | |
<<function>> | |
``` | |
Now generate test functions for me! | |
""" | |
def interact_with_agent(prompt): | |
full_prompt = task.replace("<<function>>", prompt) | |
messages = [] | |
messages.append(ChatMessage(role="user", content=full_prompt)) | |
yield messages | |
for msg in stream_from_transformers_agent(agent, full_prompt): | |
messages.append(msg) | |
yield messages | |
yield messages | |
with gr.Blocks(theme="soft") as demo: | |
gr.Markdown("""### Python test generator | |
Write your function in the left textbox, and the agent on the right will generate tests for you!""") | |
with gr.Row(): | |
with gr.Column(): | |
text_input = gr.Textbox( | |
lines=1, label="Your function to test", value=function | |
) | |
submit = gr.Button("Generate tests!") | |
with gr.Column(): | |
chatbot = AgentChatbot(label="Agent") | |
submit.click(interact_with_agent, [text_input], [chatbot]) | |
if __name__ == "__main__": | |
demo.launch() |