Skip to content

Commit e614270

Browse files
committed
initial commit
0 parents  commit e614270

File tree

6 files changed

+340
-0
lines changed

6 files changed

+340
-0
lines changed

.gitignore

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# Byte-compiled / optimized / DLL files
2+
__pycache__/
3+
*.py[cod]
4+
*$py.class
5+
6+
# C extensions
7+
*.so
8+
9+
# Distribution / packaging
10+
dist/
11+
build/
12+
*.egg-info/
13+
14+
# PyCharm
15+
.idea/
16+
17+
# VSCode
18+
.vscode/
19+
20+
# Jupyter Notebook
21+
.ipynb_checkpoints/
22+
23+
# Environment
24+
.env
25+
venv/
26+
env/

README.md

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
# Function schema
2+
3+
This is a small utility to generate JSON schemas for python functions.
4+
With power of type annotations, it is possible to generate a schema for a function without describing it twice.
5+
6+
At this moment, extracting schema from a function is only useful for [OpenAI API function-call](https://platform.openai.com/docs/guides/gpt/function-calling) feature.
7+
But it can be used for other purposes for example to generate documentation in the future.
8+
9+
## Installation
10+
11+
```sh
12+
pip install function-schema
13+
```
14+
15+
## Usage
16+
17+
```python
18+
from typing import Annotated, Optional
19+
import enum
20+
21+
def get_weather(
22+
city: Annotated[str, "The city to get the weather for"],
23+
unit: Annotated[
24+
Optional[str],
25+
"The unit to return the temperature in",
26+
enum.Enum("Unit", "celcius fahrenheit")
27+
] = "celcius",
28+
) -> str:
29+
"""Returns the weather for the given city."""
30+
return f"Weather for {city} is 20°C"
31+
```
32+
33+
Function description is taken from the docstring.
34+
Type hinting with `typing.Annotated` for annotate additional information about the parameters and return type.
35+
36+
- type can be `typing.Union`, `typing.Optional`. (`T | None` for python 3.10+)
37+
- string value of `Annotated` is used as a description
38+
- enum value of `Annotated` is used as an enum schema
39+
40+
```python
41+
import json
42+
from function_schema import get_function_schema
43+
44+
schema = get_function_schema(get_weather)
45+
print(json.dumps(schema, indent=2))
46+
```
47+
48+
Will output:
49+
50+
```json
51+
{
52+
"name": "get_weather",
53+
"description": "Returns the weather for the given city.",
54+
"parameters": {
55+
"type": "object",
56+
"properties": {
57+
"city": {
58+
"type": "string",
59+
"description": "The city to get the weather for"
60+
},
61+
"unit": {
62+
"type": "string",
63+
"description": "The unit to return the temperature in",
64+
"enum": [
65+
"celcius",
66+
"fahrenheit"
67+
],
68+
"default": "celcius"
69+
}
70+
},
71+
}
72+
"required": [
73+
"city"
74+
]
75+
}
76+
```
77+
78+
You can use this schema to make a function call in OpenAI API:
79+
```python
80+
import openai
81+
openai.api_key = "sk-..."
82+
83+
result = openai.ChatCompletion.create(
84+
model="gpt-3.5-turbo",
85+
messages=[
86+
{"role": "user", "content": "What's the weather like in Seoul?"}
87+
],
88+
functions=[
89+
get_function_schema(get_weather)
90+
],
91+
function_call="auto",
92+
)
93+
```
94+
95+
### CLI usage
96+
97+
```sh
98+
function_schema mymodule.py my_function
99+
```
100+
101+
## License
102+
MIT License

function_schema/__init__.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
"""
2+
A small utility to generate JSON schemas for python functions.
3+
"""
4+
from .core import get_function_schema
5+
6+
__version__ = "0.1.0"
7+
__all__ = (
8+
"get_function_schema",
9+
"__version__",
10+
)

function_schema/cli.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
import sys
2+
from importlib.util import module_from_spec, spec_from_file_location
3+
import inspect
4+
import json
5+
from core import get_function_schema
6+
7+
8+
def print_usage():
9+
print("Usage: function_schema <file_path> <function_name>")
10+
11+
12+
def main():
13+
# get cli args of file path
14+
try:
15+
file_path = sys.argv[1]
16+
spec = spec_from_file_location("_defendant", file_path)
17+
module = module_from_spec(spec)
18+
spec.loader.exec_module(module)
19+
members = inspect.getmembers(module)
20+
21+
for name, func in members:
22+
if name == sys.argv[2]:
23+
print(json.dumps(get_function_schema(func), indent=2))
24+
sys.exit(0)
25+
print(f"Function {sys.argv[2]} not found in {file_path}")
26+
except IndexError:
27+
print_usage()
28+
sys.exit(1)
29+
30+
if __name__ == "__main__":
31+
main()

function_schema/core.py

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
import enum
2+
import typing
3+
import inspect
4+
5+
def get_function_schema(
6+
func: typing.Annotated[typing.Callable, "The function to get the schema for"]
7+
) -> typing.Annotated[dict[str, typing.Any], "The JSON schema for the given function"]:
8+
"""
9+
Returns a JSON schema for the given function.
10+
11+
You can annotate your function parameters with the special Annotated type.
12+
Then get the schema for the function without writing the schema by hand.
13+
14+
Especially useful for OpenAI API function-call.
15+
16+
Example:
17+
>>> from typing import Annotated, Optional
18+
>>> import enum
19+
>>> def get_weather(
20+
... city: Annotated[str, "The city to get the weather for"],
21+
... unit: Annotated[
22+
... Optional[str],
23+
... "The unit to return the temperature in",
24+
... enum.Enum("Unit", "celcius fahrenheit")
25+
... ] = "celcius",
26+
... ) -> str:
27+
... \"\"\"Returns the weather for the given city.\"\"\"
28+
... return f"Hello {name}, you are {age} years old."
29+
>>> get_function_schema(get_weather)
30+
{
31+
"name": "get_weather",
32+
"description": "Returns the weather for the given city.",
33+
"parameters": {
34+
"type": "object",
35+
"properties": {
36+
"city": {
37+
"type": "string",
38+
"description": "The city to get the weather for"
39+
},
40+
"unit": {
41+
"type": "string",
42+
"description": "The unit to return the temperature in",
43+
"enum": ["celcius", "fahrenheit"],
44+
"default": "celcius"
45+
}
46+
},
47+
"required": ["city"]
48+
}
49+
}
50+
"""
51+
sig = inspect.signature(func)
52+
params = sig.parameters
53+
schema = {
54+
"type": "object",
55+
"properties": {},
56+
"required": [],
57+
}
58+
for name, param in params.items():
59+
param_args = typing.get_args(param.annotation)
60+
is_annotated = len(param_args) > 1
61+
62+
enum_ = None
63+
default_value = inspect._empty
64+
65+
if is_annotated:
66+
# first arg is type
67+
(T, _) = param_args
68+
69+
# find description in param_args tuple
70+
description = next(
71+
(arg for arg in param_args if isinstance(arg, str)),
72+
f"The {name} parameter",
73+
)
74+
75+
# find enum in param_args tuple
76+
enum_ = next(
77+
(arg for arg in param_args if isinstance(arg, enum.Enum)), None
78+
)
79+
else:
80+
T = param.annotation
81+
description = f"The {name} parameter"
82+
83+
# find default value for param
84+
if param.default is not inspect._empty:
85+
default_value = param.default
86+
87+
schema["properties"][name] = {
88+
"type": guess_type(T),
89+
"description": description, # type: ignore
90+
}
91+
92+
if enum_ is not None:
93+
schema["properties"][name]["enum"] = enum_.values
94+
95+
if default_value is not inspect._empty:
96+
schema["properties"][name]["default"] = default_value
97+
98+
if not isinstance(None, T):
99+
schema["required"].append(name)
100+
return {
101+
"name": func.__qualname__,
102+
"description": inspect.getdoc(func),
103+
"parameters": schema,
104+
}
105+
106+
107+
def guess_type(
108+
T: typing.Annotated[type, "The type to guess the JSON schema type for"]
109+
) -> typing.Annotated[
110+
typing.Union[str, list[str]], "str | list of str that representing JSON schema type"
111+
]:
112+
"""Guesses the JSON schema type for the given python type."""
113+
_types = []
114+
115+
# hacking around typing modules, `typing.Union` and `types.UnitonType`
116+
if isinstance(1, T):
117+
_types.append("integer")
118+
elif isinstance(1.1, T):
119+
_types.append("number")
120+
121+
if isinstance("", T):
122+
_types.append("string")
123+
if not isinstance(1, T) and isinstance(True, T):
124+
_types.append("boolean")
125+
if isinstance([], T):
126+
_types.append("array")
127+
if isinstance({}, T):
128+
return "object"
129+
130+
if len(_types) == 0:
131+
return "object"
132+
133+
if len(_types) == 1:
134+
return _types[0]
135+
136+
return _types

pyproject.toml

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
[project]
2+
name = "function-schema"
3+
version = "0.1.0"
4+
requires-python = ">= 3.9"
5+
description = "A small utility to generate JSON schemas for python functions."
6+
readme = "README.md"
7+
license = {text = "MIT License"}
8+
authors = [
9+
{name = "Changkyun Kim", email = "[email protected]"}
10+
]
11+
maintainers = [
12+
{name = "Changkyun Kim", email = "[email protected]"}
13+
]
14+
keywords = ["json-schema", "function", "documentation", "openai", "utility"]
15+
classifiers = [
16+
"Development Status :: 3 - Alpha",
17+
"Programming Language :: Python :: 3.9",
18+
]
19+
dynamic = ["version"]
20+
21+
[project.optional-depencencies]
22+
test = [
23+
"pytest",
24+
]
25+
26+
[project.urls]
27+
Homepage = "https://github.com/comfuture/function-schema"
28+
Repository = "https://github.com/comfuture/function-schema"
29+
30+
[project.scripts]
31+
function_schema = "function_schema.cli:main"
32+
33+
[build-system]
34+
requires = ["setuptools", "wheel"]
35+
build-backend = "setuptools.build_meta"

0 commit comments

Comments
 (0)