| from collections import defaultdict |
| from pathlib import Path |
| import json |
|
|
|
|
| writeDocumention = Path(__file__).parent / "tool_docs.json" |
| queryExamples = Path(__file__).parent / "queryExamples.json" |
| queryParams = Path(__file__).parent / "queryParams.json" |
| queryDescription = Path(__file__).parent / "queryDescription.json" |
| queryTables = Path(__file__).parent / "queryTables.json" |
| functionDocumention = Path(__file__).parent / "functionDocumention.json" |
| retrieveTables = Path(__file__).parent.parent.parent / "docs/dataset_tables.json" |
|
|
| endpoints = {"query_lineage": "extended_lineage_new", "query_policy_results": "policy_initiative","query_products": "asbom_filtered","query_risk": "product_risk_1.0", "query_vulnerabilities": "API Vulnerabilities", "query_findings": "Findings"} |
|
|
|
|
|
|
| def addExample(function: str, question: str, querystr): |
| """ |
| The function can accept querystr as either a JSON object (dict) or a string. If using a JSON object, set format="json"; otherwise, the default is to treat querystr as a plain string. |
| """ |
| try: |
| with open(queryExamples, "r+") as f: |
| data = json.load(f) |
| except: |
| data = {} |
| if function not in endpoints.keys(): raise "Not a valid function" |
| try: |
| querystr = json.loads(querystr) |
| except: |
| pass |
| if function not in data: data[function] = [] |
| for i in range(len(data[function])): |
| if data[function][i]["input"] == question: return |
| data[function].append({"input": question,"output": querystr}) |
| with open(queryExamples, "w") as f: |
| json.dump(data, fp=f, indent=2) |
|
|
| def addDescription(function: str, description: str): |
| with open(queryDescription, "r+") as f: |
| data = json.load(f) |
| data[function] = description |
| with open(queryDescription, "w") as f: |
| json.dump(data, fp=f, indent=2) |
|
|
| def updateTables(): |
| with open(retrieveTables, "r") as f: |
| data = json.load(f) |
| tables = {} |
| for key, value in endpoints.items(): |
| del data[value]["md"] |
| tables[key] = data[value] |
| with open(queryTables, "w") as f: |
| json.dump(tables, fp=f, indent=2) |
|
|
| def getTables(key: str): |
| with open(retrieveTables, "r") as f: |
| res = [] |
| data = json.load(f) |
| for item in data[key]["columns"]: |
| res.append((item["column_name"], item["label"])) |
| return res |
|
|
| def updateToolDescription(): |
| with open(queryExamples, "r") as f: |
| examples = json.load(f) |
| with open(queryParams, "r") as f: |
| params = json.load(f) |
| with open(queryDescription, "r") as f: |
| descriptions = json.load(f) |
| with open(queryTables, "r") as f: |
| tables = json.load(f) |
| endpointDocumention = defaultdict(list) |
| with open(writeDocumention, "w") as f: |
| for endpoint in endpoints.keys(): |
| endpointDocumention[endpoint].append({"name": endpoint}) |
| endpointDocumention[endpoint].append({"description": descriptions[endpoint]}) |
| |
| endpointDocumention[endpoint].append({"examples": examples[endpoint]}) |
| endpointDocumention[endpoint].append({"ColumnTables": tables[endpoint]}) |
| json.dump(endpointDocumention, fp=f, indent=2) |
|
|
| def parseDoc(): |
| with open(functionDocumention, "r") as f: |
| data = json.load(f) |
| name, description, tableName = data["name"], data["description"], data["tableName"] |
| endpoints[name] = tableName |
| for item in data["examples"].values(): |
| addExample(function=name, question=item["question"], querystr=item["querystr"]) |
| addDescription(name, description=description) |
| updateTables() |
| updateToolDescription() |
| |
| if __name__ == "__main__": |
| updateTables() |
