File size: 3,798 Bytes
85c3f44 2568dca 85c3f44 2568dca 85c3f44 18c6f92 85c3f44 2568dca 85c3f44 2568dca | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | 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({"params": params})
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() |