Spaces:
Paused
Paused
| # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import importlib.util | |
| import json | |
| import sys | |
| import types | |
| import docstring_parser | |
| def import_from_path(module_name, file_path): | |
| """Import a module given its name and file path.""" | |
| spec = importlib.util.spec_from_file_location(module_name, file_path) | |
| module = importlib.util.module_from_spec(spec) | |
| sys.modules[module_name] = module | |
| spec.loader.exec_module(module) | |
| return module | |
| class Routine: | |
| def __init__( | |
| self, | |
| name: str, | |
| description: str, | |
| parameters: dict[str, str], | |
| fn: callable, | |
| returns: dict[str, str] | None = None, | |
| module: str | None = None, | |
| source_file: str | None = None, | |
| ) -> None: | |
| self.name: str = name | |
| self.description: str = description | |
| self.parameters: dict[str, dict[str, str]] = parameters | |
| self.returns: dict[str, str] | None = returns | |
| self.fn = fn | |
| self.module: str | None = module | |
| self.source_file: str | None = source_file | |
| def __str__(self) -> str: | |
| ret = "Name: " + self.name + "\n" | |
| ret += "Description: " + self.description + "\n" | |
| ret += "Parameters:\n" | |
| for key, props in self.parameters.items(): | |
| ret += f" {key}:\n" | |
| for prop, value in props.items(): | |
| ret += f" {prop}: {value}\n" | |
| return ret | |
| def to_dict(self) -> dict: | |
| return { | |
| "name": self.name, | |
| "description": self.description, | |
| "parameters": {"type": "object", "properties": self.parameters}, | |
| } | |
| def to_json(self) -> str: | |
| return json.dumps(self.to_dict(), indent=4) | |
| class RoutineRegistry: | |
| def __init__(self, routines: list[Routine] | None = None) -> None: | |
| self.routines: list[Routine] = routines if routines is not None else [] | |
| def add_routine(self, routine: Routine) -> None: | |
| self.routines.append(routine) | |
| def remove_routine(self, routine: Routine) -> None: | |
| self.routines.remove(routine) | |
| def get_routine_of_name(self, name: str) -> Routine: | |
| return next( | |
| ( | |
| routine | |
| for routine in self.routines | |
| if routine.name.lower() == name.lower() | |
| ), | |
| None, | |
| ) | |
| def to_list(self) -> list[dict]: | |
| return [routine.to_dict() for routine in self.routines] | |
| def to_json(self) -> str: | |
| return json.dumps(self.to_list(), indent=4) | |
| def load_from_file(self, module_name: str, file_path: str) -> None: | |
| module = import_from_path(module_name, file_path) | |
| for name, obj in module.__dict__.items(): | |
| if type(obj) == types.FunctionType: | |
| doc = docstring_parser.parse(obj.__doc__) | |
| parameters = { | |
| param.arg_name: { | |
| "type": param.type_name, | |
| "description": param.description, | |
| } | |
| for param in doc.params | |
| } | |
| returns = { | |
| "type": doc.returns.type_name, | |
| "description": doc.returns.description, | |
| } | |
| self.add_routine( | |
| Routine( | |
| name, | |
| doc.short_description | |
| + "\n" | |
| + ( | |
| doc.long_description | |
| if doc.long_description is not None | |
| else "" | |
| ), | |
| parameters, | |
| obj, | |
| returns, | |
| module_name, | |
| file_path, | |
| ) | |
| ) | |