test_api_docs.py

importast
fromtypingimportList, Dict, Set, Optional
importos
fromdataclassesimportdataclass
importargparse
importre
importsys

sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
importlitellm


@dataclass
classFunctionInfo:
"""Store function information."""

name: str
docstring: Optional[str]
parameters: Set[str]
file_path: str
line_number: int


classFastAPIDocVisitor(ast.NodeVisitor):
"""AST visitor to find FastAPI endpoint functions."""

def__init__(self, target_functions: Set[str]):
self.target_functions=target_functions
self.functions: Dict[str, FunctionInfo] = {}
self.current_file=""

defvisit_FunctionDef(self, node: ast.FunctionDef|ast.AsyncFunctionDef) ->None:
"""Visit function definitions (both async and sync) and collect info if they match target functions."""
ifnode.nameinself.target_functions:
# Extract docstring
docstring=ast.get_docstring(node)

# Extract parameters
parameters=set()
forarginnode.args.args:
ifarg.annotationisnotNone:
# Get the parameter type from annotation
ifisinstance(arg.annotation, ast.Name):
parameters.add((arg.arg, arg.annotation.id))
elifisinstance(arg.annotation, ast.Subscript):
ifisinstance(arg.annotation.value, ast.Name):
parameters.add((arg.arg, arg.annotation.value.id))

self.functions[node.name] =FunctionInfo(
name=node.name,
docstring=docstring,
parameters=parameters,
file_path=self.current_file,
line_number=node.lineno,
 )

# Also need to add this to handle async functions
defvisit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) ->None:
"""Handle async functions by delegating to the regular function visitor."""
returnself.visit_FunctionDef(node)


deffind_functions_in_file(
file_path: str, target_functions: Set[str]
) ->Dict[str, FunctionInfo]:
"""Find target functions in a Python file using AST."""
try:
withopen(file_path, "r", encoding="utf-8") asf:
content=f.read()

visitor=FastAPIDocVisitor(target_functions)
visitor.current_file=file_path
tree=ast.parse(content)
visitor.visit(tree)
returnvisitor.functions

exceptExceptionase:
print(f"Error parsing {file_path}: {str(e)}")
return {}


defextract_docstring_params(docstring: Optional[str]) ->Set[str]:
"""Extract parameter names from docstring."""
ifnotdocstring:
returnset()

params=set()
# Match parameters in format:
# - parameter_name: description
# or
# parameter_name: description
param_pattern=r"-?\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*(?:\([^)]*\))?\s*:"

formatchinre.finditer(param_pattern, docstring):
params.add(match.group(1))

returnparams


defanalyze_function(func_info: FunctionInfo) ->Dict:
"""Analyze function documentation and return validation results."""

docstring_params=extract_docstring_params(func_info.docstring)

print(f"func_info.parameters: {func_info.parameters}")
pydantic_params=set()

forname, type_nameinfunc_info.parameters:
iftype_name.endswith("Request") ortype_name.endswith("Response"):
pydantic_model=getattr(litellm.proxy._types, type_name, None)
ifpydantic_modelisnotNone:
forparaminpydantic_model.model_fields.keys():
pydantic_params.add(param)

print(f"pydantic_params: {pydantic_params}")

missing_params=pydantic_params-docstring_params

return {
"function": func_info.name,
"file_path": func_info.file_path,
"line_number": func_info.line_number,
"has_docstring": bool(func_info.docstring),
"pydantic_params": list(pydantic_params),
"documented_params": list(docstring_params),
"missing_params": list(missing_params),
"is_valid": len(missing_params) ==0,
 }


defprint_validation_results(results: Dict) ->None:
"""Print validation results in a readable format."""
print(f"\nChecking function: {results['function']}")
print(f"File: {results['file_path']}:{results['line_number']}")
print("-"*50)

ifnotresults["has_docstring"]:
print("❌ No docstring found!")
return

ifnotresults["pydantic_params"]:
print("ℹ️ No Pydantic input models found.")
return

ifresults["is_valid"]:
print("✅ All Pydantic parameters are documented!")
else:
print("❌ Missing documentation for parameters:")
forparaminsorted(results["missing_params"]):
print(f" - {param}")


defmain():
function_names= [
"new_end_user",
"end_user_info",
"update_end_user",
"delete_end_user",
"generate_key_fn",
"info_key_fn",
"update_key_fn",
"delete_key_fn",
"new_user",
"new_team",
"team_info",
"update_team",
"delete_team",
"new_organization",
"update_organization",
"delete_organization",
"list_organization",
"user_update",
"new_budget",
"info_budget",
"update_budget",
"delete_budget",
"list_budget",
 ]
# directory = "../../litellm/proxy/management_endpoints" # LOCAL
directory="./litellm/proxy/management_endpoints"

# Convert function names to set for faster lookup
target_functions=set(function_names)
found_functions: Dict[str, FunctionInfo] = {}

# Walk through directory
forroot, _, filesinos.walk(directory):
forfileinfiles:
iffile.endswith(".py"):
file_path=os.path.join(root, file)
found=find_functions_in_file(file_path, target_functions)
found_functions.update(found)

# Analyze and output results
forfunc_nameinfunction_names:
iffunc_nameinfound_functions:
result=analyze_function(found_functions[func_name])
ifnotresult["is_valid"]:
raiseException(print_validation_results(result))
# results.append(result)
# print_validation_results(result)

# # Exit with error code if any validation failed
# if any(not r["is_valid"] for r in results):
# exit(1)


if__name__=="__main__":
main()