Felipe97/llama-cpp-compiled

1#!/usr/bin/env python32 3"""Function calling example using pydantic models."""4 5from __future__ import annotations6 7import argparse8import datetime9import json10import logging11import textwrap12import sys13from enum import Enum14from typing import Optional, Union15 16import requests17from pydantic import BaseModel, Field18from pydantic_models_to_grammar import (add_run_method_to_dynamic_model, convert_dictionary_to_pydantic_model,19                                        create_dynamic_model_from_function, generate_gbnf_grammar_and_documentation)20 21 22def create_completion(host, prompt, gbnf_grammar):23    """Calls the /completion API on llama-server.24 25    See26    https://github.com/ggml-org/llama.cpp/tree/HEAD/tools/server#api-endpoints27    """28    print(f"  Request:\n    Grammar:\n{textwrap.indent(gbnf_grammar, '      ')}\n    Prompt:\n{textwrap.indent(prompt.rstrip(), '      ')}")29    headers = {"Content-Type": "application/json"}30    data = {"prompt": prompt, "grammar": gbnf_grammar}31    result = requests.post(f"http://{host}/completion", headers=headers, json=data).json()32    assert data.get("error") is None, data33    logging.info("Result: %s", result)34    content = result["content"]35    print(f"  Model: {result['model']}")36    print(f"  Result:\n{textwrap.indent(json.dumps(json.loads(content), indent=2), '    ')}")37    return content38 39 40# A function for the agent to send a message to the user.41class SendMessageToUser(BaseModel):42    """Send a message to the User."""43    chain_of_thought: str = Field(..., description="Your chain of thought while sending the message.")44    message: str = Field(..., description="Message you want to send to the user.")45 46    def run(self):47        print(f"SendMessageToUser: {self.message}")48 49 50def example_rce(host):51    """Minimal test case where the LLM call an arbitrary python function."""52    print("- example_rce")53    tools = [SendMessageToUser]54    gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation(55        pydantic_model_list=tools, outer_object_name="function",56        outer_object_content="function_parameters", model_prefix="Function", fields_prefix="Parameters")57    system_message = "You are an advanced AI, tasked to assist the user by calling functions in JSON format. The following are the available functions and their parameters and types:\n\n" + documentation58    user_message = "What is 42 * 42?"59    prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant"60    text = create_completion(host, prompt, gbnf_grammar)61    json_data = json.loads(text)62    tools_map = {tool.__name__:tool for tool in tools}63    # This finds "SendMessageToUser":64    tool = tools_map.get(json_data["function"])65    if not tool:66        print(f"Error: unknown tool {json_data['function']}")67        return 168    tool(**json_data["function_parameters"]).run()69    return 070 71 72# Enum for the calculator tool.73class MathOperation(Enum):74    ADD = "add"75    SUBTRACT = "subtract"76    MULTIPLY = "multiply"77    DIVIDE = "divide"78 79 80# Simple pydantic calculator tool for the agent that can add, subtract,81# multiply, and divide. Docstring and description of fields will be used in82# system prompt.83class Calculator(BaseModel):84    """Perform a math operation on two numbers."""85    number_one: Union[int, float] = Field(..., description="First number.")86    operation: MathOperation = Field(..., description="Math operation to perform.")87    number_two: Union[int, float] = Field(..., description="Second number.")88 89    def run(self):90        if self.operation == MathOperation.ADD:91            return self.number_one + self.number_two92        elif self.operation == MathOperation.SUBTRACT:93            return self.number_one - self.number_two94        elif self.operation == MathOperation.MULTIPLY:95            return self.number_one * self.number_two96        elif self.operation == MathOperation.DIVIDE:97            return self.number_one / self.number_two98        else:99            raise ValueError("Unknown operation.")100 101 102def example_calculator(host):103    """Have the LLM ask to get a calculation done.104 105    Here the grammar gets generated by passing the available function models to106    generate_gbnf_grammar_and_documentation function. This also generates a107    documentation usable by the LLM.108 109    pydantic_model_list is the list of pydantic models outer_object_name is an110    optional name for an outer object around the actual model object. Like a111    "function" object with "function_parameters" which contains the actual model112    object. If None, no outer object will be generated outer_object_content is113    the name of outer object content.114 115    model_prefix is the optional prefix for models in the documentation. (Default="Output Model")116    fields_prefix is the prefix for the model fields in the documentation. (Default="Output Fields")117    """118    print("- example_calculator")119    tools = [SendMessageToUser, Calculator]120    gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation(121        pydantic_model_list=tools, outer_object_name="function",122        outer_object_content="function_parameters", model_prefix="Function", fields_prefix="Parameters")123    system_message = "You are an advanced AI, tasked to assist the user by calling functions in JSON format. The following are the available functions and their parameters and types:\n\n" + documentation124    user_message1 = "What is 42 * 42?"125    prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message1}<|im_end|>\n<|im_start|>assistant"126    text = create_completion(host, prompt, gbnf_grammar)127    json_data = json.loads(text)128    expected = {129        "function": "Calculator",130        "function_parameters": {131            "number_one": 42,132            "operation": "multiply",133            "number_two": 42134        }135    }136    if json_data != expected:137        print("  Result is not as expected!")138    tools_map = {tool.__name__:tool for tool in tools}139    # This finds "Calculator":140    tool = tools_map.get(json_data["function"])141    if not tool:142        print(f"Error: unknown tool {json_data['function']}")143        return 1144    result = tool(**json_data["function_parameters"]).run()145    print(f"  Call {json_data['function']} gave result {result}")146    return 0147 148 149class Category(Enum):150    """The category of the book."""151    Fiction = "Fiction"152    NonFiction = "Non-Fiction"153 154 155class Book(BaseModel):156    """Represents an entry about a book."""157    title: str = Field(..., description="Title of the book.")158    author: str = Field(..., description="Author of the book.")159    published_year: Optional[int] = Field(..., description="Publishing year of the book.")160    keywords: list[str] = Field(..., description="A list of keywords.")161    category: Category = Field(..., description="Category of the book.")162    summary: str = Field(..., description="Summary of the book.")163 164 165def example_struct(host):166    """A example structured output based on pydantic models.167 168    The LLM will create an entry for a Book database out of an unstructured169    text. We need no additional parameters other than our list of pydantic170    models.171    """172    print("- example_struct")173    tools = [Book]174    gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation(pydantic_model_list=tools)175    system_message = "You are an advanced AI, tasked to create a dataset entry in JSON for a Book. The following is the expected output model:\n\n" + documentation176    text = """The Feynman Lectures on Physics is a physics textbook based on some lectures by Richard Feynman, a Nobel laureate who has sometimes been called "The Great Explainer". The lectures were presented before undergraduate students at the California Institute of Technology (Caltech), during 1961–1963. The book's co-authors are Feynman, Robert B. Leighton, and Matthew Sands."""177    prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant"178    text = create_completion(host, prompt, gbnf_grammar)179    json_data = json.loads(text)180    # In this case, there's no function nor function_parameters.181    # Here the result will vary based on the LLM used.182    keys = sorted(["title", "author", "published_year", "keywords", "category", "summary"])183    if keys != sorted(json_data.keys()):184        print(f"Unexpected result: {sorted(json_data.keys())}")185        return 1186    book = Book(**json_data)187    print(f"  As a Book object: %s" % book)188    return 0189 190 191def get_current_datetime(output_format: Optional[str] = None):192    """Get the current date and time in the given format.193 194    Args:195         output_format: formatting string for the date and time, defaults to '%Y-%m-%d %H:%M:%S'196    """197    return datetime.datetime.now().strftime(output_format or "%Y-%m-%d %H:%M:%S")198 199 200# Example function to get the weather.201def get_current_weather(location, unit):202    """Get the current weather in a given location"""203    if "London" in location:204        return json.dumps({"location": "London", "temperature": "42", "unit": unit.value})205    elif "New York" in location:206        return json.dumps({"location": "New York", "temperature": "24", "unit": unit.value})207    elif "North Pole" in location:208        return json.dumps({"location": "North Pole", "temperature": "-42", "unit": unit.value})209    return json.dumps({"location": location, "temperature": "unknown"})210 211 212def example_concurrent(host):213    """An example for parallel function calling with a Python function, a pydantic214    function model and an OpenAI like function definition.215    """216    print("- example_concurrent")217    # Function definition in OpenAI style.218    current_weather_tool = {219        "type": "function",220        "function": {221            "name": "get_current_weather",222            "description": "Get the current weather in a given location",223            "parameters": {224                "type": "object",225                "properties": {226                    "location": {227                        "type": "string",228                        "description": "The city and state, e.g. San Francisco, CA",229                    },230                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},231                },232                "required": ["location"],233            },234        },235    }236    # Convert OpenAI function definition into pydantic model.237    current_weather_tool_model = convert_dictionary_to_pydantic_model(current_weather_tool)238    # Add the actual function to a pydantic model.239    current_weather_tool_model = add_run_method_to_dynamic_model(current_weather_tool_model, get_current_weather)240 241    # Convert normal Python function to a pydantic model.242    current_datetime_model = create_dynamic_model_from_function(get_current_datetime)243 244    tools = [SendMessageToUser, Calculator, current_datetime_model, current_weather_tool_model]245    gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation(246        pydantic_model_list=tools, outer_object_name="function",247        outer_object_content="params", model_prefix="Function", fields_prefix="Parameters", list_of_outputs=True)248    system_message = "You are an advanced AI assistant. You are interacting with the user and with your environment by calling functions. You call functions by writing JSON objects, which represent specific function calls.\nBelow is a list of your available function calls:\n\n" + documentation249    text = """Get the date and time, get the current weather in celsius in London and solve the following calculation: 42 * 42"""250    prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant"251    text = create_completion(host, prompt, gbnf_grammar)252    json_data = json.loads(text)253    expected = [254      {255        "function": "get_current_datetime",256        "params": {257          "output_format": "%Y-%m-%d %H:%M:%S"258        }259      },260      {261        "function": "get_current_weather",262        "params": {263          "location": "London",264          "unit": "celsius"265        }266      },267      {268        "function": "Calculator",269        "params": {270          "number_one": 42,271          "operation": "multiply",272          "number_two": 42273        }274      }275    ]276    res = 0277    if json_data != expected:278        print("  Result is not as expected!")279        print("  This can happen on highly quantized models")280        res = 1281    tools_map = {tool.__name__:tool for tool in tools}282    for call in json_data:283      tool = tools_map.get(call["function"])284      if not tool:285          print(f"Error: unknown tool {call['function']}")286          return 1287      result = tool(**call["params"]).run()288      print(f"  Call {call['function']} returned {result}")289    # Should output something like this:290    #   Call get_current_datetime returned 2024-07-15 09:50:38291    #   Call get_current_weather returned {"location": "London", "temperature": "42", "unit": "celsius"}292    #   Call Calculator returned 1764293    return res294 295 296def main():297    parser = argparse.ArgumentParser(description=sys.modules[__name__].__doc__)298    parser.add_argument("--host", default="localhost:8080", help="llama.cpp server")299    parser.add_argument("-v", "--verbose", action="store_true", help="enables logging")300    args = parser.parse_args()301    logging.basicConfig(level=logging.INFO if args.verbose else logging.ERROR)302    ret = 0303    # Comment out below to only run the example you want.304    ret = ret or example_rce(args.host)305    ret = ret or example_calculator(args.host)306    ret = ret or example_struct(args.host)307    ret = ret or example_concurrent(args.host)308    return ret309 310 311if __name__ == "__main__":312    sys.exit(main())313