|
| 1 | +# Function calling example using pydantic models. |
| 2 | + |
| 3 | +import json |
| 4 | +from enum import Enum |
| 5 | +from typing import Union, Optional |
| 6 | + |
| 7 | +import requests |
| 8 | +from pydantic import BaseModel, Field |
| 9 | + |
| 10 | +import importlib |
| 11 | +from pydantic_models_to_grammar import generate_gbnf_grammar_and_documentation |
| 12 | + |
| 13 | +# Function to get completion on the llama.cpp server with grammar. |
| 14 | +def create_completion(prompt, grammar): |
| 15 | + headers = {"Content-Type": "application/json"} |
| 16 | + data = {"prompt": prompt, "grammar": grammar} |
| 17 | + |
| 18 | + response = requests.post("http://127.0.0.1:8080/completion", headers=headers, json=data) |
| 19 | + data = response.json() |
| 20 | + |
| 21 | + print(data["content"]) |
| 22 | + return data["content"] |
| 23 | + |
| 24 | + |
| 25 | +# A function for the agent to send a message to the user. |
| 26 | +class SendMessageToUser(BaseModel): |
| 27 | + """ |
| 28 | + Send a message to the User. |
| 29 | + """ |
| 30 | + chain_of_thought: str = Field(..., description="Your chain of thought while sending the message.") |
| 31 | + message: str = Field(..., description="Message you want to send to the user.") |
| 32 | + |
| 33 | + def run(self): |
| 34 | + print(self.message) |
| 35 | + |
| 36 | + |
| 37 | +# Enum for the calculator function. |
| 38 | +class MathOperation(Enum): |
| 39 | + ADD = "add" |
| 40 | + SUBTRACT = "subtract" |
| 41 | + MULTIPLY = "multiply" |
| 42 | + DIVIDE = "divide" |
| 43 | + |
| 44 | + |
| 45 | +# Very simple calculator tool for the agent. |
| 46 | +class Calculator(BaseModel): |
| 47 | + """ |
| 48 | + Perform a math operation on two numbers. |
| 49 | + """ |
| 50 | + number_one: Union[int, float] = Field(..., description="First number.") |
| 51 | + operation: MathOperation = Field(..., description="Math operation to perform.") |
| 52 | + number_two: Union[int, float] = Field(..., description="Second number.") |
| 53 | + |
| 54 | + def run(self): |
| 55 | + if self.operation == MathOperation.ADD: |
| 56 | + return self.number_one + self.number_two |
| 57 | + elif self.operation == MathOperation.SUBTRACT: |
| 58 | + return self.number_one - self.number_two |
| 59 | + elif self.operation == MathOperation.MULTIPLY: |
| 60 | + return self.number_one * self.number_two |
| 61 | + elif self.operation == MathOperation.DIVIDE: |
| 62 | + return self.number_one / self.number_two |
| 63 | + else: |
| 64 | + raise ValueError("Unknown operation.") |
| 65 | + |
| 66 | + |
| 67 | +# Here the grammar gets generated by passing the available function models to generate_gbnf_grammar_and_documentation function. This also generates a documentation usable by the LLM. |
| 68 | +# pydantic_model_list is the list of pydanitc models |
| 69 | +# outer_object_name is an optional name for an outer object around the actual model object. Like a "function" object with "function_parameters" which contains the actual model object. If None, no outer object will be generated |
| 70 | +# outer_object_content is the name of outer object content. |
| 71 | +# model_prefix is the optional prefix for models in the documentation. (Default="Output Model") |
| 72 | +# fields_prefix is the prefix for the model fields in the documentation. (Default="Output Fields") |
| 73 | +gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation( |
| 74 | + pydantic_model_list=[SendMessageToUser, Calculator], outer_object_name="function", |
| 75 | + outer_object_content="function_parameters", model_prefix="Function", fields_prefix="Parameters") |
| 76 | + |
| 77 | +print(gbnf_grammar) |
| 78 | +print(documentation) |
| 79 | + |
| 80 | +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" + documentation |
| 81 | + |
| 82 | +user_message = "What is 42 * 42?" |
| 83 | +prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant" |
| 84 | + |
| 85 | +text = create_completion(prompt=prompt, grammar=gbnf_grammar) |
| 86 | +# This should output something like this: |
| 87 | +# { |
| 88 | +# "function": "calculator", |
| 89 | +# "function_parameters": { |
| 90 | +# "number_one": 42, |
| 91 | +# "operation": "multiply", |
| 92 | +# "number_two": 42 |
| 93 | +# } |
| 94 | +# } |
| 95 | +function_dictionary = json.loads(text) |
| 96 | +if function_dictionary["function"] == "calculator": |
| 97 | + function_parameters = {**function_dictionary["function_parameters"]} |
| 98 | + |
| 99 | + print(Calculator(**function_parameters).run()) |
| 100 | + # This should output: 1764 |
| 101 | + |
| 102 | + |
| 103 | +# A example structured output based on pydantic models. The LLM will create an entry for a Book database out of an unstructured text. |
| 104 | +class Category(Enum): |
| 105 | + """ |
| 106 | + The category of the book. |
| 107 | + """ |
| 108 | + Fiction = "Fiction" |
| 109 | + NonFiction = "Non-Fiction" |
| 110 | + |
| 111 | + |
| 112 | +class Book(BaseModel): |
| 113 | + """ |
| 114 | + Represents an entry about a book. |
| 115 | + """ |
| 116 | + title: str = Field(..., description="Title of the book.") |
| 117 | + author: str = Field(..., description="Author of the book.") |
| 118 | + published_year: Optional[int] = Field(..., description="Publishing year of the book.") |
| 119 | + keywords: list[str] = Field(..., description="A list of keywords.") |
| 120 | + category: Category = Field(..., description="Category of the book.") |
| 121 | + summary: str = Field(..., description="Summary of the book.") |
| 122 | + |
| 123 | + |
| 124 | +# We need no additional parameters other than our list of pydantic models. |
| 125 | +gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation([Book]) |
| 126 | + |
| 127 | +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" + documentation |
| 128 | + |
| 129 | +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.""" |
| 130 | +prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant" |
| 131 | + |
| 132 | +text = create_completion(prompt=prompt, grammar=gbnf_grammar) |
| 133 | + |
| 134 | +json_data = json.loads(text) |
| 135 | + |
| 136 | +print(Book(**json_data)) |
0 commit comments