Update app.py

Updated the API call for the certificate inventory to retrieve greater data volumes than offered by the previous API.
Added error handling to skip API calls that may result in errors due to environment configurations.
Also made prompt engineering more dynamic, using dictionaries to avoid errors and determine specifically what to instruct the LLM on and avoid confusion.

app.py

@@ -78,8 +78,16 @@ def prompt_analyze_questions(prompt, passed_args):
 
 # Data Retrieval
 
+# Some API calls may fail if customers haven't fully configured their environment
+
 def getData(tlspc_api_key, openai_api_key):
-    #Get Cert Requests
+    # Store OpenAI API Key
+    openai.api_key = openai_api_key
+    
+    # Create dictionary for adding variables
+    dfs = {}
+    
+    # Get Cert Requests
     def getCertRequests():
         currentPage = 0
         cert_requests_url = "https://api.venafi.cloud/outagedetection/v1/certificaterequestssearch"
@@ -103,38 +111,74 @@ def getData(tlspc_api_key, openai_api_key):
             data = response.json()
             cert_requests += data['certificateRequests']
         return cert_requests
-    try:
-        # Store OpenAI API Key
-        openai.api_key = openai_api_key
-        # Get Certificate Data
-        cert_url = "https://api.venafi.cloud/outagedetection/v1/certificates?ownershipTree=false&excludeSupersededInstances=false&limit=10000"
 
-        headers = {
-            "accept": "application/json",
-            "tppl-api-key": tlspc_api_key
-        }
-
-        cert_response = requests.get(cert_url, headers=headers)
-
-        certs_df = pd.json_normalize(cert_response.json()['certificates']).convert_dtypes()
-        certs_df.rename(columns = {'id':'certificateId'}, inplace = True)
-        certs_df.drop(['companyId'],axis=1,inplace=True)
+        try:
+            cert_requests_json = getCertRequests()
+            certificate_requests_df = pd.json_normalize(cert_requests_json).convert_dtypes()
+            certificate_requests_df.rename(columns = {'id':'cert_request_id', 'creationDate':'cert_request_creationDate'}, inplace = True)
+            certificate_requests_df.drop(['companyId'],axis=1,inplace=True)
+            dfs['certificate_requests_df'] = certificate_requests_df
+        except:
+            pass
+        
+    # Certificate inventory
+    cert_url = "https://api.venafi.cloud/outagedetection/v1/certificatesearch?ownershipTree=true&excludeSupersededInstances=false"
+
+    def getCerts():
+          currentPage = 0
+          payload = {
+            "paging": 
+            { "pageNumber": currentPage }}
+          response = requests.post(url=cert_url, headers=headers,json=payload)
+          if(response.status_code != 200):
+            raise Exception('Error retrieving certificates:' + "\n" + response.text + "\n=============\n")
+          data = response.json()
+          certs = data['certificates']
+          while data['count'] >= 1:
+            currentPage += 1
+            print('Getting page ' + str(currentPage) + ' - Count:' + str(data['count']))
+            payload['paging']['pageNumber'] = currentPage
+            response = requests.post(url=cert_url, headers=headers,json=payload)
+            data = response.json()
+            certs += data['certificates']
+          return certs
 
-        certs_df['validityStart'] = pd.to_datetime(certs_df['validityStart']).dt.date
-        certs_df['validityEnd'] = pd.to_datetime(certs_df['validityEnd']).dt.date
+    try:
+        certs_json = getCerts()
+        certificates_df = pd.json_normalize(certs_json).convert_dtypes()
+        certificates_df.rename(columns = {'id':'certificateId'}, inplace = True)
+        dfs['certificates_df'] = certificates_df
+    except Exception:
+        pass
 
-        # Application Data and Formatting
-        application_url = "https://api.venafi.cloud/outagedetection/v1/applications"
+    # Application Data and Formatting
+    application_url = "https://api.venafi.cloud/outagedetection/v1/applications"
 
-        headers = {
-            "accept": "application/json",
-            "tppl-api-key": tlspc_api_key
-        }
+    headers = {
+        "accept": "application/json",
+        "tppl-api-key": tlspc_api_key
+    }
 
+    try:
         application_response = requests.get(application_url, headers=headers)
 
         application_df = pd.json_normalize(application_response.json()['applications']).convert_dtypes()
 
+        cols = ['id',
+         'name',
+         'description',
+         'fullyQualifiedDomainNames',
+         'ipRanges',
+         'ports',
+         'modificationDate',
+         'creationDate','ownership.owningUsers',
+         'ownership.owningTeams']
+
+        for c in cols:
+            df_cols = application_df.columns
+            if c not in df_cols:
+                application_df[c] = np.nan
+
         application_df_2 = application_df[['id',
          'name',
          'description',
@@ -150,19 +194,23 @@ def getData(tlspc_api_key, openai_api_key):
                                                record_path = ['ownerIdsAndTypes'], 
                                                  meta = ['id']).convert_dtypes()
 
-        application_df = pd.merge(application_df_2, application_owners, left_on = 'id', right_on = 'id')
-        application_df.rename(columns = {'id':'application_id', 
+        applications_df = pd.merge(application_df_2, application_owners, left_on = 'id', right_on = 'id')
+        applications_df.rename(columns = {'id':'application_id', 
                                          'creationDate':'application_creationDate',
                                          'modificationDate':'application_modificationDate'}, inplace = True)
+        dfs['applications_df'] = applications_df
+    except Exception:
+        pass
 
-        # User Data
-        users_url = "https://api.venafi.cloud/v1/users"
+    # User Data
+    users_url = "https://api.venafi.cloud/v1/users"
 
-        headers = {
-            "accept": "application/json",
-            "tppl-api-key": tlspc_api_key
-        }
+    headers = {
+        "accept": "application/json",
+        "tppl-api-key": tlspc_api_key
+    }
 
+    try:
         users_response = requests.get(users_url, headers=headers)
 
         users_df = pd.json_normalize(users_response.json()['users']).convert_dtypes()
@@ -171,14 +219,19 @@ def getData(tlspc_api_key, openai_api_key):
 
         users_df.drop(['companyId'],axis=1,inplace=True)
 
-        # Teams Data
-        teams_url = "https://api.venafi.cloud/v1/teams"
+        dfs['users_df'] = users_df
+    except Exception:
+        pass
 
-        headers = {
-            "accept": "application/json",
-            "tppl-api-key": tlspc_api_key
-        }
+    # Teams Data
+    teams_url = "https://api.venafi.cloud/v1/teams"
+
+    headers = {
+        "accept": "application/json",
+        "tppl-api-key": tlspc_api_key
+    }
 
+    try:
         teams_response = requests.get(teams_url, headers=headers)
 
         teams_df = pd.json_normalize(teams_response.json()['teams']).convert_dtypes()
@@ -186,15 +239,19 @@ def getData(tlspc_api_key, openai_api_key):
                                    'modificationDate':'teams_modificationDate'}, inplace = True)
 
         teams_df.drop(['companyId'],axis=1,inplace=True)
+        dfs['teams_df'] = teams_df
+    except Exception:
+        pass
 
-        # Machines Data
-        machines_url = "https://api.venafi.cloud/v1/machines"
+    # Machines Data
+    machines_url = "https://api.venafi.cloud/v1/machines"
 
-        headers = {
-            "accept": "application/json",
-            "tppl-api-key": tlspc_api_key
-        }
+    headers = {
+        "accept": "application/json",
+        "tppl-api-key": tlspc_api_key
+    }
 
+    try:
         machines_response = requests.get(machines_url, headers=headers)
 
         machines_df = pd.json_normalize(machines_response.json()['machines']).convert_dtypes()
@@ -203,15 +260,19 @@ def getData(tlspc_api_key, openai_api_key):
                                       'modificationDate':'machine_modificationDate'}, inplace = True)
 
         machines_df.drop(['companyId'],axis=1,inplace=True)
+        dfs['machines_df'] = machines_df
+    except Exception:
+        pass
 
-        # Machine Identities Data
-        machine_identities_url = "https://api.venafi.cloud/v1/machineidentities"
+    # Machine Identities Data
+    machine_identities_url = "https://api.venafi.cloud/v1/machineidentities"
 
-        headers = {
-            "accept": "application/json",
-            "tppl-api-key": tlspc_api_key
-        }
+    headers = {
+        "accept": "application/json",
+        "tppl-api-key": tlspc_api_key
+    }
 
+    try:
         machine_identities_response = requests.get(machine_identities_url, headers=headers)
 
         machine_identities_df = pd.json_normalize(machine_identities_response.json()['machineIdentities']).convert_dtypes().iloc[:,:7]
@@ -221,20 +282,19 @@ def getData(tlspc_api_key, openai_api_key):
                                                 'modificationDate':'machine_identities_modificationDate'}, inplace = True)
 
         machine_identities_df.drop(['companyId'],axis=1,inplace=True)
+        dfs['machine_identities_df'] = machine_identities_df
+    except Exception:
+        pass
 
-        cert_requests_json = getCertRequests()
-        cert_requests_df = pd.json_normalize(cert_requests_json).convert_dtypes()
-        cert_requests_df.rename(columns = {'id':'cert_request_id', 'creationDate':'cert_request_creationDate'}, inplace = True)
-        cert_requests_df.drop(['companyId'],axis=1,inplace=True)
-
-        # Issuing Templates
-        issuing_template_url = "https://api.venafi.cloud/v1/certificateissuingtemplates"
+    # Issuing Templates
+    issuing_template_url = "https://api.venafi.cloud/v1/certificateissuingtemplates"
 
-        headers = {
-            "accept": "application/json",
-            "tppl-api-key": tlspc_api_key
-        }
+    headers = {
+        "accept": "application/json",
+        "tppl-api-key": tlspc_api_key
+    }
 
+    try:
         issuing_template_response = requests.get(issuing_template_url, headers=headers)
 
         issuing_templates_df = pd.json_normalize(issuing_template_response.json()['certificateIssuingTemplates']).convert_dtypes()
@@ -242,47 +302,64 @@ def getData(tlspc_api_key, openai_api_key):
                                                'creationDate':'issuing_template_creationDate'}, inplace = True)
 
         issuing_templates_df.drop(['companyId'],axis=1,inplace=True)
+        dfs['issuing_templates_df'] = issuing_templates_df
+    except Exception:
+        pass
 
-        # Prompt Engineering
-
-        # Get data structure for each dataframe to be passed in initial prompt
-        users_data_description = users_df.dtypes.apply(lambda x: x.name).to_dict()
-        application_data_description = application_df.dtypes.apply(lambda x: x.name).to_dict()
-        certificate_data_description = certs_df.dtypes.apply(lambda x: x.name).to_dict()
-        teams_data_description = teams_df.dtypes.apply(lambda x: x.name).to_dict()
-        machines_data_description = machines_df.dtypes.apply(lambda x: x.name).to_dict()
-        machine_identities_data_description = machine_identities_df.dtypes.apply(lambda x: x.name).to_dict()
-        cert_requests_data_description = cert_requests_df.dtypes.apply(lambda x: x.name).to_dict()
-        issuing_templates_data_description = issuing_templates_df.dtypes.apply(lambda x: x.name).to_dict()
-
-        data_structure_overview = f"""I have multiple python pandas dataframes. 
-        One is named application_df which contains data on applications and has the following structure: {application_data_description}.
-        Another python pandas dataframe is named users_df and contains user information and has the following structure: {users_data_description}.
-        Another python pandas dataframe is named certs_df and contains certificate information and has the following structure: {certificate_data_description}.
-        Another python pandas dataframe is named teams_df and contains teams information and has the following structure: {teams_data_description}.
-        Another python pandas dataframe is named machines_df and contains machine information and has the following structure: {machines_data_description}.
-        Another python pandas dataframe is named machine_identities_df and contains machine identity information and has the following structure: {machine_identities_data_description}.
-        Another python pandas dataframe is named cert_requests_df and contains certificate request information and has the following structure: {cert_requests_data_description}
-        Another python pandas dataframe is named issuing_templates_df and contains issuing template information and has the following structure: {issuing_templates_data_description}
-        """
-
-        data_relationships_overview = """The dataframes relate to eachother in the following manner. 
-        The column values in the 'user_id' column in users_df match the column values in the 'ownerId' column in application_df.
-        The column values in the 'team_id' column in teams_df match the column values in the 'owningTeamId' column in machines_df.
-        The column values in the 'certificateOwnerUserId' column in cert_requests_df match the column values in the 'user_id' column in users_df.
-        The column values in the 'certificateIssuingTemplateId' column in cert_requests_df match the column values in the 'issuing_template_id' column in issuing_templates_df.
-        The column values in the 'certificateOwnerUserId' column in cert_requests_df match the column values in the 'user_id' column in users_df.
-        The column values in the 'certificateIssuingTemplateId' column in certs_request_df match the column values in the 'issuing_template_id' column in issuing_templates_df.
-        """
-        passed_args = {'application_df': application_df, 'users_df': users_df, 'certs_df': certs_df,
-                       'teams_df': teams_df, 'machines_df': machines_df, 'machine_identities_df': machine_identities_df,
-                       'cert_requests_df': cert_requests_df, 'issuing_templates_df': issuing_templates_df,
-                       'data_structure_overview': data_structure_overview,
-                       'data_relationships_overview': data_relationships_overview}
-
-        return 'Data successfully loaded!', passed_args
-    except Exception as e:
-        raise e
+    # Remove dictionary objects that aren't dataframes
+
+    def dataframe_filtering(pair):
+        key, value = pair
+        if isinstance(value, pd.DataFrame):
+            return True
+        else:
+            return False
+
+    dfs = dict(filter(dataframe_filtering, dfs.items()))
+
+    # Prompt Engineering
+
+    # Get data structure for each dataframe to be passed in initial prompt
+    dict_data_types = {}
+    for i in dfs.keys():
+        dict_data_types[i] = dfs[i].dtypes.apply(lambda x: x.name).to_dict()
+
+    data_structure_overview = 'I have multiple python pandas dataframes.'
+
+    for i in dict_data_types:
+        data_structure_overview += str('\nOne is named ' 
+              + i + 
+              ' and contains ' +  
+              i.split('df')[0].replace('_',' ') + 
+              ' information and has the following structure: ' 
+              + str(dict_data_types[i]))
+
+    # Define the relationships of the dataframes to eachother
+    data_relationships_overview = "The dataframes relate to eachother in the following manner:"
+
+    if 'users_df' in dfs.keys() and 'applications_df' in dfs.keys():
+        data_relationships_overview += "/nThe column values in the 'user_id' column in users_df match the column values in the 'ownerId' column in application_df."
+    if 'teams_df' in dfs.keys() and 'machines_df' in dfs.keys():
+        data_relationships_overview += "/nThe column values in the 'team_id' column in teams_df match the column values in the 'owningTeamId' column in machines_df."
+    if 'certificate_requests_df' in dfs.keys() and 'users_df' in dfs.keys():
+        data_relationships_overview += "/nThe column values in the 'certificateOwnerUserId' column in cert_requests_df match the column values in the 'user_id' column in users_df."
+    if 'certificate_requests_df' in dfs.keys() and 'issuing_templates_df' in dfs.keys():
+        data_relationships_overview += "/nThe column values in the 'certificateIssuingTemplateId' column in cert_requests_df match the column values in the 'issuing_template_id' column in issuing_templates_df."
+    if 'machine_identities_df' in dfs.keys() and 'certificates_df' in dfs.keys():
+        data_relationships_overview += "/nThe column values in the 'certificateId' column in machine_identities_df match the column values in the 'certificateId' column in certificates_df."
+    if 'machine_identities_df' in dfs.keys() and 'machines_df' in dfs.keys():
+        data_relationships_overview += "/nThe column values in the 'machine_id' column in machine_identities_df match the column values in the 'machine_id' column in machines_df."
+
+    # If the data relationship overview is blank, just use a '' to avoid confusing the LLM
+    if data_relationships_overview == "The dataframes relate to eachother in the following manner:":
+        data_relationships_overview = ''
+
+    dfs['data_structure_overview'] = data_structure_overview
+    dfs['data_relationships_overview'] = data_relationships_overview
+
+    passed_args = dfs
+
+    return 'Data successfully loaded!', passed_args
 
 # User facing application
 with gr.Blocks(theme='aliabid94/new-theme') as demo:
@@ -298,7 +375,6 @@ with gr.Blocks(theme='aliabid94/new-theme') as demo:
         To get started, navigate to the API Keys tab. This will connect to your TLS Protect Cloud and OpenAI accounts.  
         
         Signup for TLS Protect Cloud at [venafi.com/signup](https://venafi.com/signup).   
-
         Use of this project either when run locally in your environment or on Hugging Face may introduce risks. Running this project accesses data from your TLS Protect Cloud account. While this project does not store your TLS Protect Cloud data or send that data to OpenAI, the prompts entered are sent to OpenAI for Python code generation.           
                 
         As with every opensource project, application, or online service that uses your API keys, you are strongly recommended to rotate your API keys after use.