345 lines
14 KiB
Python
345 lines
14 KiB
Python
# Usage:
|
||
#! ./server -m some-model.gguf &
|
||
#! pip install pydantic
|
||
#! python examples/json-schema-pydantic-example.py
|
||
#
|
||
# TODO:
|
||
# - https://github.com/NousResearch/Hermes-Function-Calling
|
||
#
|
||
# <|im_start|>system
|
||
# You are a function calling AI model. You are provided with function signatures within <tools></tools> XML tags
|
||
# You may call one or more functions to assist with the user query. Don't make assumptions about what values to plug into functions. Here are the available tools:
|
||
# <tools> {'type': 'function', 'function': {'name': 'get_stock_fundamentals',
|
||
# 'description': 'get_stock_fundamentals(symbol: str) -> dict - Get fundamental data for a given stock symbol using yfinance API.\n\n Args:\n symbol (str): The stock symbol.\n\n Returns:\n dict: A dictionary containing fundamental data.', 'parameters': {'type': 'object', 'properties': {'symbol': {'type': 'string'}}, 'required': ['symbol']}}}
|
||
# </tools> Use the following pydantic model json schema for each tool call you will make: {'title': 'FunctionCall', 'type': 'object', 'properties': {'arguments': {'title': 'Arguments', 'type': 'object'}, 'name': {'title': 'Name', 'type': 'string'}}, 'required': ['arguments', 'name']} For each function call return a json object with function name and arguments within <tool_call></tool_call> XML tags as follows:
|
||
# <tool_call>
|
||
# {'arguments': <args-dict>, 'name': <function-name>}
|
||
# </tool_call><|im_end|>
|
||
|
||
from dataclasses import dataclass
|
||
import subprocess
|
||
import sys
|
||
from pydantic import BaseModel, TypeAdapter
|
||
from annotated_types import MinLen
|
||
from typing import Annotated, Callable, List, Union, Literal, Optional, Type, get_args, get_origin
|
||
import json, requests
|
||
|
||
from examples.openai.api import ToolCallsTypeAdapter
|
||
|
||
def type_to_str(t):
|
||
origin = get_origin(t)
|
||
if origin is None:
|
||
return t.__name__
|
||
args = get_args(t)
|
||
return origin.__name__ + (
|
||
f'[{", ".join(type_to_str(a) for a in args)}]' if args else ''
|
||
)
|
||
|
||
def build_union_type_adapter(*types):
|
||
src = '\n'.join([
|
||
'from pydantic import TypeAdapter',
|
||
'from typing import Union',
|
||
f'_out = TypeAdapter(Union[{", ".join(type_to_str(t) for t in types)}])',
|
||
])
|
||
globs = {
|
||
**globals(),
|
||
**{t.__name__: t for t in types},
|
||
}
|
||
exec(src, globs)
|
||
return globs['_out']
|
||
|
||
class Thought(BaseModel):
|
||
thought: str
|
||
|
||
|
||
def build_tool_call_adapter2(final_output_type, *tools):
|
||
lines = [
|
||
'from pydantic import BaseModel, TypeAdapter',
|
||
'from typing import Literal, Union',
|
||
]
|
||
globs = {
|
||
**globals(),
|
||
**locals(),
|
||
final_output_type.__name__: final_output_type,
|
||
}
|
||
tool_calls = []
|
||
for fn in tools:
|
||
# TODO: escape fn.__doc__ and fn.__doc__ to avoid comment or metadata injection!
|
||
fn_name = fn.__name__
|
||
fn_doc = fn.__doc__.replace('"""', "'''") if fn.__doc__ else None
|
||
name = fn_name.replace('_', ' ').title().replace(' ', '')
|
||
lines += [
|
||
f'class {name}ToolArgs(BaseModel):',
|
||
*(f' {k}: {type_to_str(v)}' for k, v in fn.__annotations__.items() if k != 'return'),
|
||
f'class {name}ToolCall(BaseModel):',
|
||
*([f' """{fn_doc}"""'] if fn_doc else []),
|
||
f' name: Literal["{fn_name}"]',
|
||
f' arguments: {name}ToolArgs',
|
||
f'class {name}Tool(BaseModel):',
|
||
# *([f' """{fn_doc}"""'] if fn_doc else []),
|
||
f' id: str',
|
||
f' type: Literal["function"]',
|
||
f' function: {name}ToolCall',
|
||
f' def __call__(self) -> {type_to_str(fn.__annotations__.get("return"))}:',
|
||
f' return {fn_name}(**self.function.arguments.dict())',
|
||
]
|
||
tool_calls.append(f'{name}Tool')
|
||
|
||
lines += [
|
||
# 'class FinalResult(BaseModel):',
|
||
# f' result: {type_to_str(final_output_type)}',
|
||
# 'class Response(BaseModel):',
|
||
# f' """A response that starts with a thought about whether we need tools or not, the plan about tool usage (maybe a sequence of tool calls), and then either a final result (of type {final_output_type.__name__}) or a first tool call"""',
|
||
# f' original_goal: str',
|
||
# f' thought_process: str',
|
||
# # f' thought: str',
|
||
# f' next_step: Union[FinalResult, {", ".join(tool_calls)}]',
|
||
# f'response_adapter = TypeAdapter(Response)'
|
||
f'response_adapter = TypeAdapter(Union[{", ".join(tool_calls)}])',
|
||
]
|
||
|
||
exec('\n'.join(lines), globs)
|
||
return globs['response_adapter']
|
||
|
||
def create_completion2(*, response_model=None, max_tool_iterations=None, tools=[], endpoint="http://localhost:8080/v1/chat/completions", messages, **kwargs):
|
||
'''
|
||
Creates a chat completion using an OpenAI-compatible endpoint w/ JSON schema support
|
||
(llama.cpp server, llama-cpp-python, Anyscale / Together...)
|
||
|
||
The response_model param takes a type (+ supports Pydantic) and behaves just as w/ Instructor (see below)
|
||
'''
|
||
if response_model:
|
||
type_adapter = TypeAdapter(response_model)
|
||
schema = type_adapter.json_schema()
|
||
# messages = [{
|
||
# "role": "system",
|
||
# "content": f"Respond in JSON format with the following schema: {json.dumps(schema, indent=2)}"
|
||
# }] + messages
|
||
# print("Completion: ", json.dumps(messages, indent=2))
|
||
# print("SCHEMA: " + json.dumps(schema, indent=2))
|
||
response_format={"type": "json_object", "schema": schema }
|
||
|
||
tool_call_adapter = build_tool_call_adapter2(response_model, *tools)
|
||
tool_adapters = [(fn, TypeAdapter(fn)) for fn in tools]
|
||
tools_schemas = [{
|
||
"type": "function",
|
||
"function": {
|
||
"name": fn.__name__,
|
||
"description": fn.__doc__,
|
||
"parameters": ta.json_schema()
|
||
}
|
||
} for (fn, ta) in tool_adapters]
|
||
|
||
# messages = [{
|
||
# "role": "system",
|
||
# "content": '\n'.join([
|
||
# # "You are a function calling AI model. You are provided with function signatures within <tools></tools> XML tags.",
|
||
# # "You may call one or more functions to assist with the user query. Don't make assumptions about what values to plug into functions. Here are the available tools:",
|
||
# # f'<tools>{json.dumps(tools_schemas)}</tools>',
|
||
# 'Before calling each tool, you think clearly and briefly about why and how you are using the tool.',
|
||
# f"Respond in JSON format with the following schema: {json.dumps(schema, indent=2)}" if schema else "",
|
||
# ])
|
||
# }] + messages
|
||
|
||
i = 0
|
||
while (max_tool_iterations is None or i < max_tool_iterations):
|
||
body=dict(
|
||
messages=messages,
|
||
response_format=response_format,
|
||
tools=tools_schemas,
|
||
**kwargs
|
||
)
|
||
# sys.stderr.write(f'# REQUEST: {json.dumps(body, indent=2)}\n')
|
||
response = requests.post(
|
||
endpoint,
|
||
headers={"Content-Type": "application/json"},
|
||
json=body,
|
||
)
|
||
if response.status_code != 200:
|
||
raise Exception(f"Request failed ({response.status_code}): {response.text}")
|
||
|
||
# sys.stderr.write(f"\n# RESPONSE:\n\n<<<{response.text}>>>\n\n")
|
||
data = response.json()
|
||
if 'error' in data:
|
||
raise Exception(data['error']['message'])
|
||
|
||
# sys.stderr.write(f"\n# RESPONSE DATA:\n\n{json.dumps(data, indent=2)}\n\n")
|
||
# print(json.dumps(data, indent=2))
|
||
choice = data["choices"][0]
|
||
|
||
content = choice["message"].get("content")
|
||
if choice.get("finish_reason") == "tool_calls":
|
||
# sys.stderr.write(f'\n# TOOL CALLS:\n{json.dumps(choice["message"]["tool_calls"], indent=2)}\n\n')
|
||
# tool_calls =ToolCallsTypeAdapter.validate_json(json.dumps(choice["tool_calls"]))
|
||
messages.append(choice["message"])
|
||
for tool_call in choice["message"]["tool_calls"]:
|
||
# id = tool_call.get("id")
|
||
# if id:
|
||
# del tool_call["id"]
|
||
|
||
if content:
|
||
print(f'💭 {content}')
|
||
|
||
tc = tool_call_adapter.validate_json(json.dumps(tool_call))
|
||
|
||
pretty_call = f'{tc.function.name}({", ".join(f"{k}={v}" for k, v in tc.function.arguments.model_dump().items())})'
|
||
sys.stdout.write(f'⚙️ {pretty_call}')
|
||
result = tc()
|
||
sys.stdout.write(f" -> {result}\n")
|
||
messages.append({
|
||
"tool_call_id": tc.id,
|
||
"role": "tool",
|
||
"name": tc.function.name,
|
||
# "content": f'{result}',
|
||
"content": f'{pretty_call} = {result}',
|
||
})
|
||
else:
|
||
assert content
|
||
# print(content)
|
||
# print(json.dumps(json.loads(content), indent=2))
|
||
result = type_adapter.validate_json(content) if type_adapter else content
|
||
# if isinstance(result, Thought):
|
||
# print(f'💭 {result.thought}')
|
||
# messages.append({
|
||
# "role": "assistant",
|
||
# "content": json.dumps(result.model_dump(), indent=2),
|
||
# })
|
||
# else:
|
||
return result
|
||
|
||
i += 1
|
||
|
||
if max_tool_iterations is not None:
|
||
raise Exception(f"Failed to get a valid response after {max_tool_iterations} tool calls")
|
||
|
||
if __name__ == '__main__':
|
||
|
||
class QAPair(BaseModel):
|
||
question: str
|
||
concise_answer: str
|
||
justification: str
|
||
|
||
class PyramidalSummary(BaseModel):
|
||
title: str
|
||
summary: str
|
||
question_answers: Annotated[List[QAPair], MinLen(2)]
|
||
sub_sections: Optional[Annotated[List['PyramidalSummary'], MinLen(2)]]
|
||
|
||
# print("# Summary\n", create_completion(
|
||
# model="...",
|
||
# response_model=PyramidalSummary,
|
||
# messages=[{
|
||
# "role": "user",
|
||
# "content": f"""
|
||
# You are a highly efficient corporate document summarizer.
|
||
# Create a pyramidal summary of an imaginary internal document about our company processes
|
||
# (starting high-level, going down to each sub sections).
|
||
# Keep questions short, and answers even shorter (trivia / quizz style).
|
||
# """
|
||
# }]))
|
||
|
||
import math
|
||
|
||
def eval_python_expression(expr: str) -> float:
|
||
"""
|
||
Evaluate a Python expression reliably.
|
||
This can be used to compute complex nested mathematical expressions, or any python, really.
|
||
"""
|
||
print("# Evaluating expression: ", expr)
|
||
return "0.0"
|
||
|
||
def add(a: float, b: float) -> float:
|
||
"""
|
||
Add a and b reliably.
|
||
Don't use this tool to compute the square of a number (use multiply or pow instead)
|
||
"""
|
||
return a + b
|
||
|
||
# def say(something: str) -> str:
|
||
# """
|
||
# Just says something. Used to say each thought out loud
|
||
# """
|
||
# return subprocess.check_call(["say", something])
|
||
|
||
def multiply(a: float, b: float) -> float:
|
||
"""Multiply a with b reliably"""
|
||
return a * b
|
||
|
||
def divide(a: float, b: float) -> float:
|
||
"""Divide a by b reliably"""
|
||
return a / b
|
||
|
||
def pow(value: float, power: float) -> float:
|
||
"""
|
||
Raise a value to a power (exponent) reliably.
|
||
The square of x is pow(x, 2), its cube is pow(x, 3), etc.
|
||
"""
|
||
return math.pow(value, power)
|
||
|
||
result = create_completion2(
|
||
model="...",
|
||
response_model=str,
|
||
tools=[add, multiply, divide, pow], #, say],#, eval_python_expression],
|
||
# tools=[eval_python_expression],
|
||
temperature=0.0,
|
||
# repetition_penalty=1.0,
|
||
n_predict=1000,
|
||
top_k=1,
|
||
top_p=0.0,
|
||
# logit_bias={
|
||
# i: 10.0
|
||
# for i in range(1, 259)
|
||
# },
|
||
messages=[{
|
||
# "role": "system",
|
||
# "content": f"""
|
||
# You are a reliable assistant. You think step by step and think before using tools
|
||
# """
|
||
# }, {
|
||
"role": "user",
|
||
# "content": f"""
|
||
# What is 10 squared?
|
||
# """
|
||
"content": f"""
|
||
What is the sum of 2535 squared and 32222000403 then multiplied by one and a half. What's a third of the result?
|
||
|
||
Keep your goal in mind at every step.
|
||
"""
|
||
# Think step by step, start expressing the problem as an arithmetic expression
|
||
}])
|
||
|
||
# result = create_completion(
|
||
# model="...",
|
||
# response_model=float,
|
||
# tools=[add, multiply, divide, pow], #, say],#, eval_python_expression],
|
||
# temperature=0.0,
|
||
# # logit_bias={
|
||
# # i: 10.0
|
||
# # for i in range(1, 259)
|
||
# # },
|
||
# messages=[{
|
||
# "role": "user",
|
||
# # "content": f"""
|
||
# # What is 10 squared?
|
||
# # """
|
||
# "content": f"""
|
||
# What is the sum of 2535 squared and 32222000403 then multiplied by one and a half. What's a third of the result?
|
||
# """
|
||
# # Think step by step, start expressing the problem as an arithmetic expression
|
||
# }])
|
||
|
||
# 💭 First, I need to square the number 2535. For this, I will use the 'pow' tool.
|
||
# ⚙️ pow(args={'value': 2535.0, 'power': 2.0})-> 6426225.0
|
||
# 💭 Now that I have the square of 2535, I need to add it to 32222000403.0 and store the result.
|
||
# ⚙️ add(args={'a': 6426225.0, 'b': 32222000403.0})-> 32228426628.0
|
||
# 💭 Now that I have the sum of 2535 squared and 32222000403, I need to multiply it by 1.5.
|
||
# ⚙️ pow(args={'value': 32228426628.0, 'power': 1.5})-> 5785736571757004.0
|
||
# 💭 Now that I have the result of the sum multiplied by 1.5, I need to divide it by 3 to get a third of the result.
|
||
# ⚙️ divide(args={'a': 5785736571757004.0, 'b': 3.0})-> 1928578857252334.8
|
||
# 💭 I have now calculated a third of the result, which is 1928578857252334.8. I can now share this as the final answer.
|
||
# Result: 1928578857252334.8
|
||
|
||
expected_result = (2535 ** 2 + 32222000403) * 1.5 / 3.0
|
||
print("➡️", result)
|
||
assert math.fabs(result - expected_result) < 0.0001, f"Expected {expected_result}, got {result}"
|