|  | 
|  | 1 | +#!/usr/bin/env python3 | 
|  | 2 | +# function calling using llama-cli | 
|  | 3 | + | 
|  | 4 | +import subprocess | 
|  | 5 | +import sys | 
|  | 6 | +import select | 
|  | 7 | +import os | 
|  | 8 | +import re | 
|  | 9 | + | 
|  | 10 | +import json | 
|  | 11 | + | 
|  | 12 | +import functions | 
|  | 13 | +from function_tool import get_function_tool_json, generate_schema_from_functions | 
|  | 14 | + | 
|  | 15 | +function_name_list = [ name for name in dir(functions) if not name.startswith('_') ] | 
|  | 16 | +function_lookup = { name: getattr(functions, name) for name in function_name_list } | 
|  | 17 | +tools = [ get_function_tool_json(f) for (n, f) in function_lookup.items() ] | 
|  | 18 | +function_schema = generate_schema_from_functions(tools) | 
|  | 19 | + | 
|  | 20 | +prompt = """<|start_header_id|>system<|end_header_id|> | 
|  | 21 | +
 | 
|  | 22 | +You are capable of executing available function(s) if required. | 
|  | 23 | +Execute function(s) as needed. | 
|  | 24 | +The function calls are not shown in the conversation and should be called covertly to answer questions. | 
|  | 25 | +Ask for the required input to:recipient==all | 
|  | 26 | +Use JSON for function arguments. | 
|  | 27 | +Respond in this format: | 
|  | 28 | +>>>${recipient} | 
|  | 29 | +${content} | 
|  | 30 | +Available functions: | 
|  | 31 | +""" + function_schema + """<|eot_id|><|start_header_id|>system<|end_header_id|> | 
|  | 32 | +
 | 
|  | 33 | +When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 60.0 seconds. The drive at '/mnt/data' can be used to save and persist user files.<|eot_id|><|start_header_id|>user<|end_header_id|> | 
|  | 34 | +""" | 
|  | 35 | + | 
|  | 36 | +def main(): | 
|  | 37 | +    import argparse | 
|  | 38 | + | 
|  | 39 | +    parser = argparse.ArgumentParser(epilog='For more options: llama-cli --help') | 
|  | 40 | +    parser.add_argument('--display-prompt', action=argparse.BooleanOptionalAction, default=False) | 
|  | 41 | +    parser.add_argument('--special', action=argparse.BooleanOptionalAction, default=False) | 
|  | 42 | +    parser.add_argument('--reverse-prompt', type=str, default='<|start_header_id|>user<|end_header_id|>\n') | 
|  | 43 | +    parser.add_argument('--ctx-size', type=int, default=1024) | 
|  | 44 | +    args, other_args = parser.parse_known_args() | 
|  | 45 | + | 
|  | 46 | +    if args.display_prompt: print(prompt) | 
|  | 47 | + | 
|  | 48 | +    command = [ './llama-cli', '-i', '-p', prompt, '--reverse-prompt', args.reverse_prompt, '--escape', '--special', '--no-display-prompt', '--log-disable', '--simple-io', '--ctx-size',  str(args.ctx_size), *other_args] | 
|  | 49 | + | 
|  | 50 | +    process = subprocess.Popen( | 
|  | 51 | +        command, | 
|  | 52 | +        stdin=subprocess.PIPE, | 
|  | 53 | +        stdout=subprocess.PIPE, | 
|  | 54 | +        stderr=subprocess.PIPE, | 
|  | 55 | +        text=True, | 
|  | 56 | +    ) | 
|  | 57 | +    if process.stdout is not None: os.set_blocking(process.stdout.fileno(), False) | 
|  | 58 | + | 
|  | 59 | +    try: | 
|  | 60 | +        run_loop(process, args) | 
|  | 61 | +    except KeyboardInterrupt: | 
|  | 62 | +        print("\nInterrupted by user.") | 
|  | 63 | +    finally: | 
|  | 64 | +        process.terminate() | 
|  | 65 | +        process.wait() | 
|  | 66 | + | 
|  | 67 | +def run_loop(process, args): | 
|  | 68 | +    pbuffer = '' | 
|  | 69 | +    skip_output_until_result = False | 
|  | 70 | +    while True: | 
|  | 71 | +        readable, _, _ = select.select([process.stdout, process.stderr, sys.stdin], [], []) | 
|  | 72 | + | 
|  | 73 | +        for stream in readable: | 
|  | 74 | +            if stream == process.stdout: | 
|  | 75 | +                pdata = process.stdout.read() | 
|  | 76 | +                if not pdata: continue | 
|  | 77 | +                pbuffer += pdata | 
|  | 78 | + | 
|  | 79 | +                if(match := re.search(r'>>>([^\n]*)\n(.*)<\|eot_id\|>', pbuffer, re.S)): | 
|  | 80 | +                    if not args.special: | 
|  | 81 | +                        pdata = pdata[:match.pos] | 
|  | 82 | +                    pbuffer = '' | 
|  | 83 | +                    skip_output_until_result = False | 
|  | 84 | + | 
|  | 85 | +                    tool_name = match.group(1) | 
|  | 86 | +                    tool_args = match.group(2) | 
|  | 87 | + | 
|  | 88 | +                    if tool_name == 'python': | 
|  | 89 | +                        result = functions._run_python(tool_args); | 
|  | 90 | +                    else: | 
|  | 91 | +                        try: | 
|  | 92 | +                            tool_args = json.loads(tool_args) | 
|  | 93 | +                            result = function_lookup[tool_name](**tool_args) | 
|  | 94 | +                        except ValueError as e: | 
|  | 95 | +                            result = {'error': 'unknown'} | 
|  | 96 | + | 
|  | 97 | +                    result = json.dumps(result) + '<|eot_id|><|start_header_id|>assistant<|end_header_id|>' | 
|  | 98 | +                    process.stdin.write(result + '\n') | 
|  | 99 | +                    process.stdin.flush() | 
|  | 100 | +                    if(args.special): pdata += '\n' + result | 
|  | 101 | +                elif (n := pdata.find('>>>')) >= 0: | 
|  | 102 | +                    if not args.special: | 
|  | 103 | +                        pdata = pdata[:n] | 
|  | 104 | +                        skip_output_until_result = True | 
|  | 105 | +                elif skip_output_until_result: | 
|  | 106 | +                    pdata = '' | 
|  | 107 | + | 
|  | 108 | +                if not args.special: | 
|  | 109 | +                    pdata = re.sub(r'<\|[^\|>]*\|>', '', pdata) | 
|  | 110 | +                sys.stdout.write(pdata) | 
|  | 111 | +                sys.stdout.flush() | 
|  | 112 | + | 
|  | 113 | +            elif stream == sys.stdin: | 
|  | 114 | +                user_input = sys.stdin.readline() | 
|  | 115 | +                if user_input: | 
|  | 116 | +                    user_input = user_input.rstrip() | 
|  | 117 | +                    process.stdin.write(user_input + '<|eot_id|><|start_header_id|>assistant<|end_header_id|>' + '\n') | 
|  | 118 | +                    process.stdin.flush() | 
|  | 119 | + | 
|  | 120 | +if __name__ == '__main__': | 
|  | 121 | +    main() | 
|  | 122 | + | 
0 commit comments