ailinnesse/Subreddit_Pulse

1import inspect
2import sys
3import typing
4from typing import get_origin, get_args, Literal, Union, Optional
5
6try:
7    from typing import TypedDict  # py3.8+
8except ImportError:
9    TypedDict = None
10
11def _is_typeddict(t):
12    try:
13        return isinstance(t, type) and TypedDict is not None and issubclass(t, TypedDict)
14    except TypeError:
15        return False
16
17def _is_dataclass(t):
18    try:
19        import dataclasses
20        return dataclasses.is_dataclass(t)
21    except Exception:
22        return False
23
24def _docstring_split_sections(doc: str):
25    """Very small parser to extract:
26       - short summary (first non-empty line)
27       - param descriptions from sections like 'Args:', 'Parameters:', ':param x:'.
28    """
29    if not doc:
30        return "", {}
31
32    lines = [l.rstrip() for l in doc.strip().splitlines()]
33    # Summary = first nonempty line
34    summary = next((l for l in lines if l.strip()), "")
35    params_desc = {}
36
37    # Gather “Args/Parameters/Arguments” blocks (Google/Numpy style)
38    markers = {"args:", "parameters:", "arguments:"}
39    i = 0
40    while i < len(lines):
41        line = lines[i].strip().lower()
42        if line in markers:
43            i += 1
44            while i < len(lines):
45                raw = lines[i]
46                if raw.strip() == "" or raw.startswith(" "):
47                    # keep reading indented or blank continuation lines
48                    # detect "name (type): desc" or "name: desc"
49                    stripped = raw.strip()
50                    if stripped:
51                        # Try common patterns
52                        if ":" in stripped:
53                            name, desc = stripped.split(":", 1)
54                            name = name.strip().split()[0].split("(")[0]
55                            params_desc.setdefault(name, desc.strip())
56                        else:
57                            # continuation line: append to last desc if any
58                            if params_desc:
59                                last = list(params_desc.keys())[-1]
60                                params_desc[last] += " " + stripped
61                    i += 1
62                else:
63                    break
64            continue
65        i += 1
66
67    # Sphinx-style ":param name: desc"
68    for l in lines:
69        ls = l.strip()
70        if ls.lower().startswith(":param "):
71            try:
72                rest = ls[len(":param "):]
73                name, desc = rest.split(":", 1)
74                name = name.strip().split()[0]
75                params_desc[name] = desc.strip()
76            except ValueError:
77                pass
78
79    return summary, params_desc
80
81
82def _json_type_for_python(t):
83    """Return a JSON Schema fragment for python/typing type t."""
84    origin = get_origin(t)
85    args = get_args(t)
86
87    # NoneType
88    if t is type(None):
89        return {"type": "null"}
90
91    # Builtins
92    if t is str:
93        return {"type": "string"}
94    if t is int:
95        return {"type": "integer"}
96    if t is float:
97        return {"type": "number"}
98    if t is bool:
99        return {"type": "boolean"}
100
101    # datetime-like
102    try:
103        import datetime as _dt
104        if t in (_dt.datetime,):
105            return {"type": "string", "format": "date-time"}
106        if t in (_dt.date,):
107            return {"type": "string", "format": "date"}
108        if t in (_dt.time,):
109            return {"type": "string", "format": "time"}
110        if t in (_dt.timedelta,):
111            # no standard JSON Schema, fallback to string
112            return {"type": "string", "description": "Duration (ISO 8601 or human-readable)."}
113    except Exception:
114        pass
115
116    # Enum
117    import enum
118    if isinstance(t, type) and issubclass(t, enum.Enum):
119        values = [e.value for e in t]
120        # infer primitive type of the enum values
121        if all(isinstance(v, str) for v in values):
122            return {"type": "string", "enum": values}
123        if all(isinstance(v, int) for v in values):
124            return {"type": "integer", "enum": values}
125        # mixed types
126        return {"enum": values}
127
128    # TypedDict
129    if _is_typeddict(t):
130        props = {}
131        required = []
132        # __annotations__ holds fields
133        ann = t.__annotations__
134        total = getattr(t, "__total__", True)
135        for k, v in ann.items():
136            props[k] = _json_type_for_python(v)
137            # In total=True, all are required unless Optional/Union[..., None]
138            if total:
139                if not _is_optional(v):
140                    required.append(k)
141            else:
142                # total=False => all optional
143                pass
144        schema = {"type": "object", "properties": props}
145        if required:
146            schema["required"] = required
147        return schema
148
149    # dataclass
150    if _is_dataclass(t):
151        import dataclasses
152        props = {}
153        required = []
154        for f in dataclasses.fields(t):
155            props[f.name] = _json_type_for_python(f.type)
156            has_default = f.default is not dataclasses.MISSING or f.default_factory is not dataclasses.MISSING
157            if not has_default and not _is_optional(f.type):
158                required.append(f.name)
159        schema = {"type": "object", "properties": props}
160        if required:
161            schema["required"] = required
162        return schema
163
164    # Literal
165    if origin is Literal:
166        vals = list(args)
167        # infer a base type if uniform
168        if all(isinstance(v, str) for v in vals):
169            return {"type": "string", "enum": vals}
170        if all(isinstance(v, int) for v in vals):
171            return {"type": "integer", "enum": vals}
172        if all(isinstance(v, (int, float)) for v in vals):
173            # number enum
174            return {"type": "number", "enum": vals}
175        return {"enum": vals}
176
177    # Optional[T] == Union[T, None]
178    if _is_optional(t):
179        # caller should handle required vs optional; here return the underlying schema
180        non_none = [a for a in args if a is not type(None)]
181        if len(non_none) == 1:
182            return _json_type_for_python(non_none[0])
183        # Optional of Union[…, None] falls through to anyOf
184        return {"anyOf": [_json_type_for_python(a) for a in non_none] + [{"type": "null"}]}
185
186    # Union
187    if origin is Union:
188        return {"anyOf": [_json_type_for_python(a) for a in args]}
189
190    # List/Tuple/Set
191    if origin in (list, tuple, set, typing.Sequence, typing.MutableSequence):
192        item_t = args[0] if args else typing.Any
193        return {"type": "array", "items": _json_type_for_python(item_t)}
194
195    # Dict / Mapping
196    if origin in (dict, typing.Mapping, typing.MutableMapping):
197        key_t, val_t = (args + (typing.Any, typing.Any))[:2]
198        # JSON keys must be strings; if key_t != str, we note it in description
199        schema = {"type": "object", "additionalProperties": _json_type_for_python(val_t)}
200        if key_t is not str:
201            schema["description"] = (schema.get("description", "") + " Keys will be stringified.").strip()
202        return schema
203
204    # Fallbacks
205    if t is typing.Any or t is None:
206        return {}  # unconstrained
207    # Unknown type => treat as string with note
208    return {"type": "string", "description": f"Serialized {getattr(t, '__name__', str(t))}."}
209
210
211def _is_optional(t):
212    origin = get_origin(t)
213    if origin is Union:
214        args = get_args(t)
215        return any(a is type(None) for a in args)
216    return False
217
218
219def function_to_tool(
220    func,
221    *,
222    name: str | None = None,
223    description: str | None = None,
224    param_overrides: dict | None = None,
225) -> dict:
226    """
227    Build an OpenAI-style 'tool' schema from a Python function.
228
229    - `name`: override function name.
230    - `description`: override function description (otherwise from docstring summary).
231    - `param_overrides`: dict of per-param overrides, e.g.
232        {
233          "city": {"description": "City name", "enum": ["Dubai", "Abu Dhabi"]},
234          "units": {"default": "metric"}  # note: default isn't used by schema; make param optional instead
235        }
236    """
237    sig = inspect.signature(func)
238    hints = typing.get_type_hints(func, include_extras=True)
239    doc = inspect.getdoc(func) or ""
240    summary, param_descs = _docstring_split_sections(doc)
241
242    tool_name = name or func.__name__
243    tool_desc = description or summary or f"Callable function `{tool_name}`."
244
245    properties = {}
246    required = []
247
248    for pname, param in sig.parameters.items():
249        if pname == "self":
250            continue
251
252        ann = hints.get(pname, typing.Any)
253        schema = _json_type_for_python(ann)
254
255        # base description from docstring, if any
256        if param_descs.get(pname):
257            schema["description"] = param_descs[pname]
258
259        # overrides
260        if param_overrides and pname in param_overrides:
261            schema.update(param_overrides[pname])
262
263        # required vs optional
264        is_required = (
265            param.default is inspect._empty
266            and not _is_optional(ann)
267        )
268        if is_required:
269            required.append(pname)
270
271        # ensure at least a basic type if none inferred
272        if not schema:
273            schema = {"type": "string"}
274
275        properties[pname] = schema
276
277    parameters = {
278        "type": "object",
279        "properties": properties,
280    }
281    if required:
282        parameters["required"] = required
283
284    return {
285        "type": "function",
286        "name": tool_name,
287        "description": tool_desc,
288        "parameters": parameters,
289    }
290