import bilbystats as bs

def write_test(fun_name, instructions = '', model = 'anthropic/claude-sonnet-4.5', n_examples = 3, max_tokens = 1000):
    code = bs.extract_fun(fun_name)

    prompt = '''Write python code containing ''' + str(n_examples) + ''' examples for the provided function. 
    Separate each example with '\n#%%\n'. 
    DO NOT include the function itself in the code output. 
    The function parallelize belongs to the bilbystats package so can be imported from there in the code, so include "import bilbystats as bs" at the beggining of the code output.'''
    prompt += "Here's the function:\n"
    prompt += code
    prompt += "\nProvide only the Python code without explanations."

    code = bs.llm_api(prompt, instructions, model, max_tokens = max_tokens)
    code_clean = bs.process_python_string(code)
    return code_clean