|
import inspect
|
|
from typing import get_type_hints, get_origin, get_args
|
|
import pr.tools
|
|
|
|
|
|
def _type_to_json_schema(py_type):
|
|
"""Convert Python type to JSON Schema type."""
|
|
if py_type == str:
|
|
return {"type": "string"}
|
|
elif py_type == int:
|
|
return {"type": "integer"}
|
|
elif py_type == float:
|
|
return {"type": "number"}
|
|
elif py_type == bool:
|
|
return {"type": "boolean"}
|
|
elif get_origin(py_type) == list:
|
|
return {"type": "array", "items": _type_to_json_schema(get_args(py_type)[0])}
|
|
elif get_origin(py_type) == dict:
|
|
return {"type": "object"}
|
|
else:
|
|
# Default to string for unknown types
|
|
return {"type": "string"}
|
|
|
|
|
|
def _generate_tool_schema(func):
|
|
"""Generate JSON Schema for a tool function."""
|
|
sig = inspect.signature(func)
|
|
docstring = func.__doc__ or ""
|
|
|
|
# Extract description from docstring
|
|
description = docstring.strip().split("\n")[0] if docstring else ""
|
|
|
|
# Get type hints
|
|
type_hints = get_type_hints(func)
|
|
|
|
properties = {}
|
|
required = []
|
|
|
|
for param_name, param in sig.parameters.items():
|
|
if param_name in ["db_conn", "python_globals"]: # Skip internal parameters
|
|
continue
|
|
|
|
param_type = type_hints.get(param_name, str)
|
|
schema = _type_to_json_schema(param_type)
|
|
|
|
# Add description from docstring if available
|
|
param_doc = ""
|
|
if docstring:
|
|
lines = docstring.split("\n")
|
|
in_args = False
|
|
for line in lines:
|
|
line = line.strip()
|
|
if line.startswith("Args:") or line.startswith("Arguments:"):
|
|
in_args = True
|
|
continue
|
|
elif in_args and line.startswith(param_name + ":"):
|
|
param_doc = line.split(":", 1)[1].strip()
|
|
break
|
|
elif in_args and line == "":
|
|
continue
|
|
elif in_args and not line.startswith(" "):
|
|
break
|
|
|
|
if param_doc:
|
|
schema["description"] = param_doc
|
|
|
|
# Set default if available
|
|
if param.default != inspect.Parameter.empty:
|
|
schema["default"] = param.default
|
|
|
|
properties[param_name] = schema
|
|
|
|
# Required if no default
|
|
if param.default == inspect.Parameter.empty:
|
|
required.append(param_name)
|
|
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": func.__name__,
|
|
"description": description,
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": properties,
|
|
"required": required,
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def get_tools_definition():
|
|
"""Dynamically generate tool definitions from all tool functions."""
|
|
tools = []
|
|
|
|
# Get all functions from pr.tools modules
|
|
for name in dir(pr.tools):
|
|
if name.startswith("_"):
|
|
continue
|
|
|
|
obj = getattr(pr.tools, name)
|
|
if callable(obj) and hasattr(obj, "__module__") and obj.__module__.startswith("pr.tools."):
|
|
# Check if it's a tool function (has docstring and proper signature)
|
|
if obj.__doc__:
|
|
try:
|
|
schema = _generate_tool_schema(obj)
|
|
tools.append(schema)
|
|
except Exception as e:
|
|
print(f"Warning: Could not generate schema for {name}: {e}")
|
|
continue
|
|
|
|
return tools
|
|
|
|
|
|
def get_func_map(db_conn=None, python_globals=None):
|
|
"""Dynamically generate function map for tool execution."""
|
|
func_map = {}
|
|
|
|
# Include all functions from __all__ in pr.tools
|
|
for name in getattr(pr.tools, "__all__", []):
|
|
if name.startswith("_"):
|
|
continue
|
|
|
|
obj = getattr(pr.tools, name, None)
|
|
if callable(obj) and hasattr(obj, "__module__") and obj.__module__.startswith("pr.tools."):
|
|
sig = inspect.signature(obj)
|
|
params = list(sig.parameters.keys())
|
|
|
|
# Create wrapper based on parameters
|
|
if "db_conn" in params and "python_globals" in params:
|
|
func_map[name] = (
|
|
lambda func=obj, db_conn=db_conn, python_globals=python_globals, **kw: func(
|
|
**kw, db_conn=db_conn, python_globals=python_globals
|
|
)
|
|
)
|
|
elif "db_conn" in params:
|
|
func_map[name] = lambda func=obj, db_conn=db_conn, **kw: func(**kw, db_conn=db_conn)
|
|
elif "python_globals" in params:
|
|
func_map[name] = lambda func=obj, python_globals=python_globals, **kw: func(
|
|
**kw, python_globals=python_globals
|
|
)
|
|
else:
|
|
func_map[name] = lambda func=obj, **kw: func(**kw)
|
|
|
|
return func_map
|