File size: 2,299 Bytes
e1a8ab5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
from typing import Callable, Any, Dict, get_type_hints, Optional
from dataclasses import dataclass
import inspect
from typing import _GenericAlias
@dataclass
class Tool:
name: str
description: str
func: Callable[..., str]
parameters: Dict[str, Dict[str, str]]
def __call__(self, *args, **kwargs) -> str:
return self.func(*args, **kwargs)
def parse_docstring_params(docstring: str) -> Dict[str, str]:
"""Extract parameter descriptions from docstring."""
if not docstring:
return {}
params = {}
lines = docstring.split('\n')
in_params = False
current_param = None
for line in lines:
line = line.strip()
if line.startswith('Parameters:'):
in_params = True
elif in_params:
if line.startswith('-') or line.startswith('*'):
current_param = line.lstrip('- *').split(':')[0].strip()
params[current_param] = line.lstrip('- *').split(':')[1].strip()
elif current_param and line:
params[current_param] += ' ' + line.strip()
elif not line:
in_params = False
return params
def get_type_description(type_hint: Any) -> str:
"""Get a human-readable description of a type hint."""
if isinstance(type_hint, _GenericAlias):
if type_hint._name == 'Literal':
return f"one of {type_hint.__args__}"
return type_hint.__name__
def tool(name: str = None):
def decorator(func: Callable[..., str]) -> Tool:
tool_name = name or func.__name__
description = inspect.getdoc(func) or "No description available"
type_hints = get_type_hints(func)
param_docs = parse_docstring_params(description)
sig = inspect.signature(func)
params = {}
for param_name, param in sig.parameters.items():
params[param_name] = {
"type": get_type_description(type_hints.get(param_name, Any)),
"description": param_docs.get(param_name, "No description available")
}
return Tool(
name=tool_name,
description=description.split('\n\n')[0],
func=func,
parameters=params
)
return decorator |