Upload 36 files

.gitignore

@@ -0,0 +1,11 @@
+.env
+test_endpoint.py
+app/llm_workflows.py
+build
+__pycache__
+databricks_llm.egg-info
+dist
+hf_endpoint.py
+azureOAI_endpoint.py
+**/prompt.py
+prompts

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

README - raw.md

@@ -0,0 +1,42 @@
+# Introduction 
+This repo contains a CLI which can perform interactions with the Databricks API by using natural language.
+This functionality is enabled by Langchain and Large Language Models (LLMs).
+The code is tested with text-davinci-003 from OpenAI with the configuration in llm.yaml. Every OpenAI model can be used in the configuration llm.yaml.
+
+# General
+You can either install the CLI directly or use the repository directly with python.py databricks_llm.py \[OPTIONS\] COMMAND \[ARGS\].
+
+# Getting Started
+1. Dependencies
+* All necessary dependencies are included in the requirements.txt and should be installed in a virtual environment. (e.g. pip install -r requirements.txt)
+2. Installation of the CLI:
+* python setup.py bdist_wheel
+* pip install .\dist\databricks_llm-0.0.1-py3-none-any.whl
+3. The following environment variables always need to be set:
+* OPENAI_API_KEY = ... (Your access key from OpenAI or Azure OpenAI)
+* DATABRICKS_HOST = ... (The host name in the form "adb-\<workspace-id\>.azuredatabricks.net")
+* DBR_BEARER_TOKEN = ... (A personal access token for your workspace)
+4. If you want to use the **AzureOpenAI** Service, please add the following environment variables:
+* OPENAI_API_TYPE = "azure" (example)
+* OPENAI_API_VERSION= "2023-05-15" (example)
+* OPENAI_API_BASE = "https://\<resource-name\>.openai.azure.com/"
+5. The **llm.yaml** config:
+* This contains the parameters for the Large Langugage Model.
+* **For AzureOpenAI the entry "deployment_name" needs to be present, otherwise please remove it**
+
+# Functionalities
+This repo contains several functionalities around Databricks Jobs/Workflows and their permissions as well as functionalities around MLflow models that are stored within a Databricks workspace. The commands and their descriptions can be accessed with the --help option.
+
+# Example workflows
+* Jobs:
+    * Get particular job infos: databricks_llm jobs "Please return a list of all jobs; Figure out individual job configuration and only return the job names with tags 'NLP' and their jobID summarized:"
+* Permissions:
+    * Add a permission entry: databricks_llm permissions "Please add an entry in the list for the group: '\<group_name\>' with permission 'CAN_VIEW'" "\<workflow_id\>"
+    * Modify a permission entry: databricks_llm permissions "Please modify the entry in the list for the group: '\<group_name\>' to permission 'CAN_MANAGE'" "\<workflow_id\>"  
+* Machine Learning:
+    * Get model infos: databricks_llm ml get-model-info "Please give me all infos for the model '\<model_name\>'"
+    * Get infos about a run: databricks_llm ml get-run-info "\<run-id\>" "Please return the f1 score"
+    * Transition a model: databricks_llm ml transition-model "Please transition model '\<model_name\>' version 6 into stage 'Production' and archive the current model in production"
+
+# Overview of the workflows
+![api_chatter](./assets/API_Chatter_Overview.PNG)

app.py

@@ -0,0 +1,103 @@
+import click
+import logging
+import sys
+import yaml
+import os
+from langchain.llms import AzureOpenAI, OpenAI
+from app.api_funcs import get_job_infos, get_run, get_model, \
+    trans_model, batch_mod_permission, prepare_api_docs
+from pathlib import Path
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+handler = logging.StreamHandler(sys.stdout)
+handler.setLevel(logging.DEBUG)
+logger.addHandler(handler)
+
+PATH = Path(os.path.abspath(os.path.dirname(__file__)))
+
+# Open the YAML file
+conf_path = PATH / "app" / 'llm.yaml'
+with open(conf_path) as f:
+    config = yaml.safe_load(f)
+
+# Use AzureOpenAI, if config contains deployment name, otherwise OpenAI
+if config['model'].get('deployment_name', False):
+    llm = AzureOpenAI(**config['model'])
+else:
+    llm = OpenAI(**config['model'])
+
+
+headers = {"Authorization": f"Bearer {os.getenv('DBR_BEARER_TOKEN')}"}
+updated_api_docs = prepare_api_docs()
+
+
+def comma_list(comma_str: str):
+    return comma_str.split(',')
+
+def determine_api_text(updated_api_docs: dict, query: str):
+    pick_api_prompt = """Please return the file name from the list {api_docs}
+                        that best corresponds to the following query: {query}. \
+                        DO NOT EXPLAIN your answer!
+                        """
+    api_docs = os.listdir(PATH / "app" / "dbr_api_docs")
+    selected_api_doc = llm(pick_api_prompt.format(api_docs=api_docs, query=query)).lstrip().rstrip()
+    logger.info(f"\nSelecting the following api document: {selected_api_doc}")
+    api_text = updated_api_docs[selected_api_doc]
+    return api_text, selected_api_doc    
+
+# Add subcommands for commands
+@click.group()
+def cli():
+    pass
+
+@cli.group(help='Run machine learning model.')
+def ml():
+    pass
+
+# Add commands for specific subcommands of 'ml'
+@ml.command(help='Get information about a model.')
+@click.argument('query', type=str)
+def get_model_info(query):
+    # Instruction to get model infos
+    api_text, _ = determine_api_text(updated_api_docs, query)
+    logger.info(get_model(llm, query, api_text, headers))
+
+@ml.command(help='Get information about a model run.')
+@click.argument('run_id', type=str)
+@click.argument('query', type=str)
+def get_run_info(query, run_id):
+    # ID of the model run for which you'd like information.
+    # Which information should be pulled from the run?
+    api_text, _ = determine_api_text(updated_api_docs, query)
+    logger.info(get_run(llm, run_id, query, api_text, headers))
+
+@ml.command(help='Transition a model from one state to another.')
+@click.argument('query', type=str)
+def transition_model(query):
+    # Instruction to transition a model.
+    api_text, _ = determine_api_text(updated_api_docs, query)
+    trans_model(llm, query, api_text, headers)
+
+@cli.command(help='View job history.')
+@click.argument('query', type=str)
+def jobs(query):
+    if ";" not in query:
+        query = query + ";"
+    query, response_query = query.split(";")
+    api_text, _ = determine_api_text(updated_api_docs, query)
+    # The query for the LLM + an optional query for the API response
+    logger.info(get_job_infos(llm, query, response_query, api_text, headers))
+
+@cli.command(help='Manage user permissions.')
+@click.argument('query', type=str)
+@click.argument('jobs', type=comma_list)
+def permissions(jobs, query):
+    api_text, api_name = determine_api_text(updated_api_docs, query)
+    # Add/Get user permissions.
+    batch_mod_permission(
+        logger, llm, updated_api_docs, api_text, api_name, headers,
+        query, jobs=jobs
+    )
+
+if __name__ == '__main__':
+    cli()

app/api_funcs.py

@@ -0,0 +1,206 @@
+
+import json
+import os
+import ast
+from pathlib import Path
+from typing import List
+import langchain.llms as LLM
+from app.custom_langchain import ModAPIChain, APIResponse, FlexAPIChain, FlexAPIChainPayload
+from app.utils import find_value, custom_api_prompt
+from app.custom_api_prompts import API_REQUEST_PROMPT, \
+    API_REQUEST_PROMPT2, API_RESPONSE_PROMPT, API_RESPONSE_PROMPT2, PERMISSION_PROMPT
+
+PATH = Path(os.path.abspath(os.path.dirname(__file__)))
+
+# Open the YAML file
+api_doc_path = PATH / 'dbr_api_docs'
+
+def prepare_api_docs():
+    """
+    Add environment variables to the individual
+    API docs.
+    """
+    updated_api_docs = {}
+    for filename in os.listdir(api_doc_path):
+        with open(os.path.join(api_doc_path, filename)) as f:
+            updated_api_docs[filename] = f.read()
+            updated_api_docs[filename] = updated_api_docs[filename].replace(
+                "{DATABRICKS_HOST}", os.getenv('DATABRICKS_HOST'))
+    return updated_api_docs
+
+
+def get_run(llm: LLM, runID: str, query: str, api_text: str, headers: dict) -> str:
+    """
+    Get infos from an MLflow run.
+
+    Args:
+        llm (LLM): Large language model
+        runID (str): Run ID
+        query (str): Query
+        api_text (str): API documentation for a particular endpoint
+        headers (dict): Authorization for the Databricks API
+    
+    Return:
+        (str): the response of the language model
+    """
+    mlflow_run_get_chain = APIResponse.from_llm_and_api_docs(llm, api_text, headers=headers, verbose=True)
+    run_infos = json.loads(mlflow_run_get_chain.run(
+        f"""Please give me the infos about the run '{runID}'"""))
+    # Please return the f1 score
+    return llm(f"""{query} of the metrics : {find_value(run_infos, 'metrics')}""")
+
+
+def trans_model(llm: LLM, query: str, api_text: str, headers: dict):
+    """
+    Transition a MLflow model to another stage.
+
+    Args:
+        llm (LLM): Large language model
+        query (str): Query for the model
+        api_text (str): API documentation for a particular endpoint
+        headers (dict): Authorization for the Databricks API
+    """
+    mlflow_model_transition_chain = FlexAPIChainPayload.from_llm_and_api_docs(
+        llm, api_text, headers=headers, api_url_prompt=API_REQUEST_PROMPT2,
+        api_response_prompt=API_RESPONSE_PROMPT2, verbose=True
+    )
+    mlflow_model_transition_chain.run(query)
+
+
+def get_model(llm: LLM, model_query: str, api_text: str, headers: dict) -> str:
+    """
+    Get all infos about a model and its particular stages
+
+    Args:
+        llm (LLM): Large language model
+        model_query (str): Query
+        api_text (str): API documentation for a particular endpoint
+        headers (dict): Authorization for the Databricks API
+    
+    Return:
+        (str): response of the language model
+    """
+    mlflow_model_get_chain = APIResponse.from_llm_and_api_docs(llm, api_text, headers=headers, verbose=True)
+    api_result = mlflow_model_get_chain.run(model_query)  
+
+    return llm(
+        f"""What is the runID and their version of the latest version with stage 'None' \
+            and stage 'Production' respectively, given the following info: {api_result}"""
+    )
+
+
+def batch_mod_permission(
+    logger, llm: LLM, updated_api_docs: dict, api_text: str, api_name: str,
+    headers: dict, permission_mod: str, jobs: List[str]
+):
+    """
+    Get or modify permissions for a batch of jobs.
+
+    Args:
+        logger (Logger): Logger object
+        llm (LLM): Large language model
+        updated_api_docs (dict): Updated API docs
+        api_text (str): API documentation for a particular endpoint
+        api_name (str): API name
+        headers (dict): Authorization for the Databricks API
+        permission_mod (str): Permission modification
+        jobs (list): List of jobs
+    """
+    get_permission_txt = updated_api_docs['get_permissions.txt']
+    for job in jobs:
+        logger.info(f"Permission for job: {job}")
+        mod_permissions(
+            logger, llm, permission_mod, job, headers,
+            get_permission_txt, api_text, api_name
+        )
+
+
+def get_permissions(logger, llm: LLM, api_text: str, headers: dict, jobID: str) -> dict:
+    """
+    Get the permissions of a particular job.
+    
+    Args:
+        logger (Logger): Logger object
+        llm (LLM): Large language model
+        api_text (str): API documentation for a particular endpoint
+        headers (dict): Authorization for the Databricks API
+        jobID (str): ID of the Databricks job
+    
+    Return:
+        (dict): Access control list
+    """
+    init_query = f"""Get the jobs permissions for jobID {jobID}"""
+    permission_get_chain = APIResponse.from_llm_and_api_docs(llm, api_text, headers=headers, verbose=True)
+
+    acc_control_list = {'access_control_list': json.loads(permission_get_chain.run(init_query))['access_control_list']}
+    logger.info(acc_control_list)
+    return acc_control_list    
+
+def mod_permissions(
+    logger,
+    llm: LLM,
+    permission_mod: str,
+    jobID: str, headers: dict,
+    get_permission_txt: str,
+    api_text: str, api_name: str
+):
+    """
+    Modify permissions for a job.
+
+    Args:
+        logger (Logger): Logger object
+        llm (LLM): Large language model
+        permission_mod (str): Permission modification
+        jobID (str): ID of the Databricks job
+        headers (dict): Authorization for the Databricks API
+        get_permission_txt (str): API text for getting permissions
+        api_text (str): API documentation for a particular endpoint
+        api_name (str): file name of the API documentation
+    """
+    acc_control_list = get_permissions(logger, llm, get_permission_txt, headers, jobID)
+
+    # Update the permission, if the query demands it
+    if api_name == 'update_permissions.txt':
+        permission_prompt= PERMISSION_PROMPT.format(
+            acc_control_list=acc_control_list, permission_mod=permission_mod)
+
+        output_acc_list = llm(permission_prompt)
+        patch_payload = ast.literal_eval(output_acc_list.rstrip().lstrip()) #
+        # Reformat the general permission response into the form for the patch payload
+        for idx, _ in enumerate(patch_payload['access_control_list']):
+            patch_payload['access_control_list'][idx]['permission_level'] = patch_payload['access_control_list'][idx]['all_permissions'][0]['permission_level']
+            patch_payload['access_control_list'][idx].pop('all_permissions')
+
+        permission_update_chain = FlexAPIChain.from_llm_and_api_docs(
+            llm, api_text, headers=headers, verbose=True,
+            api_url_prompt=API_REQUEST_PROMPT, api_response_prompt=API_RESPONSE_PROMPT
+        ) # use updated API Chain which can do post, put and patch also....
+
+        permission_update_chain.body = patch_payload
+        logger.info(f"Payload for permission update: {patch_payload}")
+        permission_update_chain(f"""Update the jobs permissions of jobID {jobID}""")    
+
+
+def get_job_infos(
+    llm: LLM, query: str, response_query: str, api_text: str, headers: dict
+) -> str:
+    """
+    Get particular infos about the list of existing jobs.
+
+    Args:
+        llm (LLM): Large language model
+        query (str): Query to get the list of jobs
+        response_query (str): Query to get the response
+        api_text (str): API documentation for a particular endpoint
+        headers (dict): Authorization for the Databricks API
+    
+    Return:
+        (str): response from the large language model
+    """
+    job_chain = ModAPIChain.from_llm_and_api_docs(
+        llm, api_text, headers=headers, verbose=True
+    )
+
+    custom_resp = f"This response contains a list of databricks jobs. {response_query}"
+    jobs_result = custom_api_prompt(llm, job_chain, query, custom_resp)
+    return jobs_result 

app/convert_api_txt.py

@@ -0,0 +1 @@
+# Script for converting raw api texts (crawled, cop&paste from webpage etc.)

app/custom_api_prompts.py

@@ -0,0 +1,83 @@
+from langchain.prompts.prompt import PromptTemplate
+
+API_URL_PROMPT_TEMPLATE = """You are given the below API Documentation:
+{api_docs}
+Using this documentation, generate the full API url to call for answering the user question.
+You should build the API url in order to get a response that is as short as possible, while still getting the necessary information to answer the question. Pay attention to deliberately exclude any unnecessary pieces of data in the API call.
+AFTER the API url, you should extract the request METHOD (can be GET, POST, PATCH or PUT) from doc.
+
+Question:{question}
+"""
+
+API_REQUEST_PROMPT_TEMPLATE = API_URL_PROMPT_TEMPLATE + """Output the API url and METHOD, join them with `|`. DO NOT GIVE ANY EXPLANATION."""
+
+API_REQUEST_PROMPT = PromptTemplate(
+    input_variables=[
+        "api_docs",
+        "question",
+    ],
+    template=API_REQUEST_PROMPT_TEMPLATE,
+)
+
+API_RESPONSE_PROMPT_TEMPLATE = (
+    API_URL_PROMPT_TEMPLATE
+    + """API url: {api_url}
+
+Here is the response from the API:
+
+{api_response}
+
+Summarize this response to answer the original question.
+
+Summary:"""
+)
+
+API_RESPONSE_PROMPT = PromptTemplate(
+    input_variables=["api_docs", "question", "api_url", "api_response"],
+    template=API_RESPONSE_PROMPT_TEMPLATE,
+)
+
+API_URL_PROMPT_TEMPLATE2 = """You are given the below API Documentation:
+{api_docs}
+Using this documentation, generate the full API url to call for answering the user question.
+You should build the API url in order to get a response that is as short as possible, while still getting the necessary information to answer the question. Pay attention to deliberately exclude any unnecessary pieces of data in the API call.
+AFTER the API url, you should extract the request METHOD from doc, and generate the BODY data in JSON format according to the user question if necessary. The BODY data could be empty dict.
+
+Question:{question}
+"""
+
+API_REQUEST_PROMPT_TEMPLATE2 = API_URL_PROMPT_TEMPLATE2 + """Output the API url, METHOD and BODY, join them with `|`. DO NOT GIVE ANY EXPLANATION."""
+
+API_REQUEST_PROMPT2 = PromptTemplate(
+    input_variables=[
+        "api_docs",
+        "question",
+    ],
+    template=API_REQUEST_PROMPT_TEMPLATE2,
+)
+
+API_RESPONSE_PROMPT_TEMPLATE2 = (
+    API_URL_PROMPT_TEMPLATE2
+    + """API url: {api_url}
+
+Here is the response from the API:
+
+{api_response}
+
+Summarize this response to answer the original question.
+
+Summary:"""
+)
+
+API_RESPONSE_PROMPT2 = PromptTemplate(
+    input_variables=["api_docs", "question", "api_url", "api_response"],
+    template=API_RESPONSE_PROMPT_TEMPLATE2,
+)
+
+PERMISSION_PROMPT = """
+The following list is the content for an 'access control list' of Databricks Jobs:
+            {acc_control_list}
+Each entry of the 'access_control_list' denotes either a group or user and their permission level.
+{permission_mod} (This permission is not inherited!):
+"""
+

app/custom_langchain.py

@@ -0,0 +1,104 @@
+from typing import Dict, Any
+import json
+from langchain.chains import APIChain
+from typing import Any, Dict, Optional
+from langchain.callbacks.manager import (
+    CallbackManagerForChainRun,
+)
+
+
+class ModAPIChain(APIChain):
+    """Chain that makes API calls and summarizes the responses to answer a question,
+    but doesn't display any API outputs for confidentiality! (only difference to APIChain)"""
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+        question = inputs[self.question_key]
+        api_url = self.api_request_chain.predict(
+            question=question,
+            api_docs=self.api_docs,
+            callbacks=_run_manager.get_child(),
+        )
+        _run_manager.on_text(api_url, color="green", end="\n", verbose=self.verbose)
+        api_response = self.requests_wrapper.get(api_url)
+        answer = self.api_answer_chain.predict(
+            question=question,
+            api_docs=self.api_docs,
+            api_url=api_url,
+            api_response=api_response,
+            callbacks=_run_manager.get_child(),
+        )
+        return {self.output_key: answer}
+
+
+class APIResponse(APIChain):
+    """Chain that makes API calls and summarizes the responses to answer a question."""
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager = None,
+    ) -> Dict[str, str]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+        question = inputs[self.question_key]
+        api_url = self.api_request_chain.predict(
+            question=question,
+            api_docs=self.api_docs,
+            callbacks=_run_manager.get_child(),
+        )
+        _run_manager.on_text(api_url, color="green", end="\n", verbose=self.verbose)
+        api_response = self.requests_wrapper.get(api_url)
+        answer = api_response
+        return {self.output_key: answer}
+
+
+class FlexAPIChain(APIChain):
+    """
+    Flexible API Chain which can create all request types
+    whereby the body is passed as an attribute. Relies on
+    specific Prompt Templates.
+    """
+    body: dict = {}
+    def _call(self, inputs: Dict[str, str], run_manager= None) -> Dict[str, str]:
+        question = inputs[self.question_key]
+        request_info = self.api_request_chain.predict(
+            question=question, api_docs=self.api_docs
+        )
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+
+        api_url, request_method = request_info.split('|')
+        _run_manager.on_text(request_method, color="green", end="\n", verbose=self.verbose)
+        _run_manager.on_text(api_url, color="green", end="\n", verbose=self.verbose)
+
+        # get the method with same name
+        request_func = getattr(self.requests_wrapper, request_method.lower().rstrip().lstrip())
+
+        api_response = request_func(api_url, self.body)
+        return {self.output_key: api_response}
+
+
+class FlexAPIChainPayload(APIChain):
+    """
+    Flexible API Chain which can create all request types
+    including the necessary payload. Relies on
+    specific Prompt Templates.
+    """
+    def _call(self, inputs: Dict[str, str], run_manager = None) -> Dict[str, str]:
+        question = inputs[self.question_key]
+        request_info = self.api_request_chain.predict(
+            question=question, api_docs=self.api_docs
+        )
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+
+        api_url, request_method, body = request_info.split('|')
+        _run_manager.on_text(body, color="green", end="\n", verbose=self.verbose)
+        _run_manager.on_text(request_method, color="green", end="\n", verbose=self.verbose)
+        _run_manager.on_text(api_url, color="green", end="\n", verbose=self.verbose)
+
+        # get the method with same name
+        request_func = getattr(self.requests_wrapper, request_method.lower().rstrip().lstrip())
+
+        api_response = request_func(api_url, json.loads(body))
+        return {self.output_key: api_response}

app/dbr_api_docs/get_permissions.txt

@@ -0,0 +1,21 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+API Endpoint: Get Object Permissions
+
+Path:
+GET /api/2.0/permissions/{request_object_type}/{request_object_id}
+
+Path Parameters:
+- request_object_type (required, string)
+- request_object_id (required, string)
+
+Responses:
+- 200
+  - object_id (string)
+  - object_type (string)
+  - access_control_list (Array of objects):
+    - user_name (string): Name of the user
+    - group_name (string): Name of the group
+    - service_principal_name (string): Name of the service principal
+    - all_permissions (Array of objects): All permissions

app/dbr_api_docs/jobs_cancel_now.txt

@@ -0,0 +1,16 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+## Cancel a Job Run
+
+**POST** `/api/2.1/jobs/runs/cancel`
+
+Cancels a job run. The run is canceled asynchronously, so it may still be running when this request completes.
+
+**Parameters**
+
+- `run_id` (required, integer, int64): This field is required.
+
+**Responses**
+
+- `200` Run was cancelled successfully.

app/dbr_api_docs/jobs_create.txt

@@ -0,0 +1,13 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+**POST /api/2.1/jobs/create**
+Create a new job.
+
+**Payload**
+A valid job dictionary.
+
+**Responses**
+- 200 Job was created successfully
+  - job_id: integer, int64
+    The canonical identifier for the newly created job.

app/dbr_api_docs/jobs_get.txt

@@ -0,0 +1,21 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+## Get a Single Job
+
+**GET** `/api/2.1/jobs/get`
+
+Retrieves the details for a single job.
+
+**Query Parameters**
+
+- `job_id` (required, integer, int64): The canonical identifier of the job to retrieve information about. This field is required.
+
+**Responses**
+
+- `200` Job was retrieved successfully.
+  - `created_time` (integer, int64): The time at which this job was created in epoch milliseconds (milliseconds since 1/1/1970 UTC).
+  - `job_id` (integer, int64): The canonical identifier for this job.
+  - `trigger_history` (object): History of the file arrival trigger associated with the job.
+  - `creator_user_name` (string): The creator user name. This field won’t be included in the response if the user has already been deleted.
+  - `run_as_user_name` (string): The user name that the job runs as. run_as_user_name is based on the current job settings, and is set to the creator of the

app/dbr_api_docs/jobs_list.txt

@@ -0,0 +1,23 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+**List all jobs**
+
+**GET** /api/2.1/jobs/list
+
+Retrieves a list of jobs.
+
+**Query Parameters**
+
+- `limit`: integer [1..25], default: 20. The number of jobs to return. This value must be greater than 0 and less or equal to 25.
+- `offset`: integer, default: 0. The offset of the first job to return, relative to the most recently created job.
+- `name`: string. A filter on the list based on the exact (case insensitive) job name.
+- `expand_tasks`: boolean, default: false. Whether to include task and cluster details in the response.
+
+**Responses**
+
+- 200 List of jobs was retrieved successfully.
+  - `jobs`: Array of objects. The list of jobs.
+    - `job_id`: integer, int64. The canonical identifier for this job.
+    - `creator_user_name`: string. The creator user name. This field won’t be included in the response if the user has already been deleted.
+    - `settings`: object

app/dbr_api_docs/jobs_run_now.txt

@@ -0,0 +1,13 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+POST /api/2.1/jobs/run-now
+Run a job and return the run_id of the triggered run.
+
+Parameters:
+- job_id (required, integer, int64): The ID of the job to be executed
+
+Responses:
+- 200 Run was started successfully.
+  - run_id (integer, int64): The globally unique ID of the newly triggered run.
+  - number_in_job (integer, int64): Deprecated. A unique identifier for this job run. This is set to the same value as run_id.

app/dbr_api_docs/mlflow_registered_models_get.txt

@@ -0,0 +1,24 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+## API Documentation
+**Endpoint**
+
+GET /api/2.0/mlflow/databricks/registered-models/get
+
+**Query Parameters**
+
+- `name` (required): string, Registered model unique name identifier.
+
+**Responses**
+
+- `200` Model details were returned successfully.
+
+**Response Body**
+
+- `registered_model` (object):
+  - `name` (string): [1..100] characters, Name of the model.
+  - `description` (string): <= 65535 characters, User-specified description for the object.
+  - `permission_level` (string): Enum: "CAN_MANAGE" "CAN_EDIT" "CAN_READ" "CAN_MANAGE_STAGING_VERSIONS" "CAN_MANAGE_PRODUCTION_VERSIONS", Permission level of the requesting user on the object. For what is allowed at each level, see MLflow Model permissions.
+  - `tags` (Array of objects): Array of tags associated with the model.
+  - `latest_versions` (Array of objects): Array of model versions, each the

app/dbr_api_docs/mlflow_runs_get.txt

@@ -0,0 +1,24 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+## API Documentation
+
+**Get a run**
+
+**GET** `/api/2.0/mlflow/runs/get`
+
+Gets the metadata, metrics, params, and tags for a run. In the case where multiple metrics with the same key are logged for a run, return only the value with the latest timestamp. If there are multiple values with the latest timestamp, return the maximum of these values.
+
+**Query Parameters**
+
+- `run_id` (required, string): ID of the run to fetch. Must be provided.
+- `run_uuid` (string): [Deprecated, use run_id instead] ID of the run to fetch. This field will be removed in a future MLflow version.
+
+**Responses**
+
+- `200`:
+  - `run` (object): Run metadata (name, start time, etc) and data (metrics, params, and tags).
+    - `info` (object): Run metadata.
+    - `data` (object): Run data.
+      - `metrics` (Array of objects): Run metrics.
+      - `params` (Array of objects):

app/dbr_api_docs/mlflow_transition_stage.txt

@@ -0,0 +1,16 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+API Endpoint: POST /api/2.0/mlflow/databricks/model-versions/transition-stage
+
+Description: Transition a model version's stage. This is a Databricks workspace version of the MLflow endpoint that also accepts a comment associated with the transition to be recorded.
+
+Parameters: 
+- stage (required, string, Enum: "None" "Staging" "Production" "Archived"): Target stage of the transition. Valid values are: None (The initial stage of a model version), Staging (Staging or pre-production stage), Production (Production stage), Archived (Archived stage).
+- name (required, string, [1..100] characters): Name of the model.
+- archive_existing_versions (required, boolean): Specifies whether to archive all current model versions in the target stage.
+- version (required, string): Version of the model.
+- comment (string, [1..65535] characters): User-provided comment on the action.
+
+Responses:
+- 200: Model version's stage was updated successfully.

app/dbr_api_docs/update_permissions.txt

@@ -0,0 +1,12 @@
+Endpoint PREFIX : https://{DATABRICKS_HOST}/ 
+
+
+### Update Permission
+
+**PATCH** `/api/2.0/permissions/{request_object_type}/{request_object_id}`
+
+Updates the permissions on an object. Payload needs to be a valid access control list.
+
+**Responses**
+
+- `200` - Permissions updated successfully.

app/dbr_api_docs_raw/get_permissions.txt

@@ -0,0 +1,40 @@
+Get object permissions
+GET
+/api/2.0/permissions/{request_object_type}/{request_object_id}
+Gets the permission of an object. Objects can inherit permissions from their parent objects or root objects.
+
+Path Parameters
+request_object_type
+required
+string
+request_object_id
+required
+string
+Responses
+200
+object_id
+string
+
+object_type
+string
+
+access_control_list
+Array of objects
+Array[ 
+  user_name
+  string
+  name of the user
+
+  group_name
+  string
+  name of the group
+
+  service_principal_name
+  string
+  name of the service principal
+
+  all_permissions
+  Array of objects
+  All permissions.
+]
+

app/dbr_api_docs_raw/jobs_cancel_now.txt

@@ -0,0 +1,13 @@
+Cancel a job run
+POST
+/api/2.1/jobs/runs/cancel
+Cancels a job run. The run is canceled asynchronously, so it may still be running when this request completes.
+
+run_id
+required
+integer
+int64
+This field is required.
+
+Responses
+200 Run was cancelled successfully.

app/dbr_api_docs_raw/jobs_create.txt

@@ -0,0 +1,13 @@
+Create a new job
+POST
+/api/2.1/jobs/create
+Create a new job.
+
+The payload needs to be a valid job dictionary.
+
+Responses
+200 Job was created successfully
+job_id
+integer
+int64
+The canonical identifier for the newly created job.

app/dbr_api_docs_raw/jobs_get.txt

@@ -0,0 +1,39 @@
+Get a single job
+GET
+/api/2.1/jobs/get
+Retrieves the details for a single job.
+
+Query Parameters
+job_id
+required
+integer
+int64
+The canonical identifier of the job to retrieve information about. This field is required.
+
+Responses
+200 Job was retrieved successfully.
+created_time
+integer
+int64
+The time at which this job was created in epoch milliseconds (milliseconds since 1/1/1970 UTC).
+
+job_id
+integer
+int64
+The canonical identifier for this job.
+
+trigger_history
+object
+History of the file arrival trigger associated with the job.
+
+creator_user_name
+string
+The creator user name. This field won’t be included in the response if the user has already been deleted.
+
+run_as_user_name
+string
+The user name that the job runs as. run_as_user_name is based on the current job settings, and is set to the creator of the job if job access control is disabled, or the is_owner permission if job access control is enabled.
+
+settings
+object
+Settings for this job and all of its runs. These settings can be updated using the resetJob method.

app/dbr_api_docs_raw/jobs_list.txt

@@ -0,0 +1,53 @@
+List all jobs
+GET
+/api/2.1/jobs/list
+Retrieves a list of jobs.
+
+Query Parameters
+limit
+integer
+[ 1 .. 25 ]
+Default: 20
+The number of jobs to return. This value must be greater than 0 and less or equal to 25. The default value is 20.
+
+offset
+integer
+Default: 0
+The offset of the first job to return, relative to the most recently created job.
+
+name
+string
+A filter on the list based on the exact (case insensitive) job name.
+
+expand_tasks
+boolean
+Default: false
+Whether to include task and cluster details in the response.
+
+Responses
+200 List of jobs was retrieved successfully.
+jobs
+Array of objects
+The list of jobs.
+
+Array 
+job_id
+integer
+int64
+The canonical identifier for this job.
+
+creator_user_name
+string
+The creator user name. This field won’t be included in the response if the user has already been deleted.
+
+settings
+object
+Settings for this job and all of its runs. These settings can be updated using the resetJob method.
+
+created_time
+integer
+int64
+The time at which this job was created in epoch milliseconds (milliseconds since 1/1/1970 UTC).
+
+has_more
+boolean

app/dbr_api_docs_raw/jobs_run_now.txt

@@ -0,0 +1,23 @@
+Trigger a new job run
+POST
+/api/2.1/jobs/run-now
+Run a job and return the run_id of the triggered run.
+
+job_id
+required
+integer
+int64
+The ID of the job to be executed
+
+Responses
+200 Run was started successfully.
+run_id
+integer
+int64
+The globally unique ID of the newly triggered run.
+
+number_in_job
+integer
+int64
+Deprecated
+A unique identifier for this job run. This is set to the same value as run_id.

app/dbr_api_docs_raw/mlflow_registered_models_get.txt

@@ -0,0 +1,57 @@
+Get model
+GET
+/api/2.0/mlflow/databricks/registered-models/get
+Get the details of a model. This is a Databricks workspace version of the MLflow endpoint that also returns the model's Databricks workspace ID and the permission level of the requesting user on the model.
+
+Query Parameters
+name
+required
+string
+Registered model unique name identifier.
+
+Responses
+200 Model details were returned successfully.
+registered_model
+object
+name
+string
+[ 1 .. 100 ] characters
+Name of the model.
+
+description
+string
+<= 65535 characters
+User-specified description for the object.
+
+permission_level
+string
+Enum: "CAN_MANAGE" "CAN_EDIT" "CAN_READ" "CAN_MANAGE_STAGING_VERSIONS" "CAN_MANAGE_PRODUCTION_VERSIONS"
+Permission level of the requesting user on the object. For what is allowed at each level, see MLflow Model permissions.
+
+tags
+Array of objects
+Array of tags associated with the model.
+
+latest_versions
+Array of objects
+Array of model versions, each the latest version for its stage.
+
+user_id
+string
+email
+The username of the user that created the object.
+
+creation_timestamp
+integer
+int64
+Creation time of the object, as a Unix timestamp in milliseconds.
+
+id
+string
+uuid
+Unique identifier for the object.
+
+last_updated_timestamp
+integer
+int64
+Time of the object at last update, as a Unix timestamp in milliseconds.

app/dbr_api_docs_raw/mlflow_runs_get.txt

@@ -0,0 +1,47 @@
+Get a run
+GET
+/api/2.0/mlflow/runs/get
+Gets the metadata, metrics, params, and tags for a run. In the case where multiple metrics with the same key are logged for a run, return only the value with the latest timestamp.
+
+If there are multiple values with the latest timestamp, return the maximum of these values.
+
+Query Parameters
+run_id
+required
+string
+ID of the run to fetch. Must be provided.
+
+run_uuid
+string
+Deprecated
+[Deprecated, use run_id instead] ID of the run to fetch. This field will be removed in a future MLflow version.
+
+Responses
+200
+run
+object
+Run metadata (name, start time, etc) and data (metrics, params, and tags).
+
+info
+object
+Run metadata.
+
+data
+object
+Run data.
+
+metrics
+Array of objects
+Run metrics.
+
+params
+Array of objects
+Run parameters.
+
+tags
+Array of objects
+Additional metadata key-value pairs.
+
+inputs
+object
+Run inputs.

app/dbr_api_docs_raw/mlflow_transition_stage.txt

@@ -0,0 +1,44 @@
+Transition a stage
+POST
+/api/2.0/mlflow/databricks/model-versions/transition-stage
+Transition a model version's stage. This is a Databricks workspace version of the MLflow endpoint that also accepts a comment associated with the transition to be recorded.",
+
+Details required to transition a model version's stage.
+
+stage
+required
+string
+Enum: "None" "Staging" "Production" "Archived"
+Target stage of the transition. Valid values are:
+
+None: The initial stage of a model version.
+
+Staging: Staging or pre-production stage.
+
+Production: Production stage.
+
+Archived: Archived stage.
+
+name
+required
+string
+[ 1 .. 100 ] characters
+Name of the model.
+
+archive_existing_versions
+required
+boolean
+Specifies whether to archive all current model versions in the target stage.
+
+version
+required
+string
+Version of the model.
+
+comment
+string
+[ 1 .. 65535 ] characters
+User-provided comment on the action.
+
+Responses
+200 Model version's stage was updated successfully.

app/dbr_api_docs_raw/update_permissions.txt

@@ -0,0 +1,7 @@
+Update permission
+PATCH
+/api/2.0/permissions/{request_object_type}/{request_object_id}
+Updates the permissions on an object. Payload needs to be a valid access control list
+
+Responses
+200

app/llm.yaml

@@ -0,0 +1,5 @@
+# Model configuration
+model:
+  model_name: text-davinci-003
+  deployment_name: ptc-davinci-003
+  temperature: 0

app/utils.py

@@ -0,0 +1,46 @@
+from langchain import LLMChain
+from langchain.prompts.prompt import PromptTemplate
+from langchain.chains.api.prompt import API_URL_PROMPT_TEMPLATE
+
+
+def find_value(dictionary: dict, key):
+    """
+    Recursive function to return the value of key in a dictionary
+    (Code was also created by an LLM)
+    """
+    if key in dictionary:
+        return dictionary[key]
+    for k, v in dictionary.items():
+        if isinstance(v, dict):
+            item = find_value(v, key)
+            if item is not None:
+                return item
+
+
+def create_custom_response_template(custom_resp: str):
+    # Customized Prompt
+    API_RESPONSE_PROMPT_TEMPLATE = (
+        API_URL_PROMPT_TEMPLATE
+        + """ {api_url}
+
+    Here is the response from the API:
+
+    {api_response}""" + f"\n\n{custom_resp}"
+    )
+
+    return PromptTemplate(
+        input_variables=["api_docs", "question", "api_url", "api_response"],
+        template=API_RESPONSE_PROMPT_TEMPLATE,
+    )
+
+
+def custom_api_prompt(llm, chain, question, custom_resp):
+    """
+    Helper function for separating the api-link creation and the post-processing on the response
+    """
+    resp_prompt = create_custom_response_template(custom_resp)
+    get_answer_chain = LLMChain(llm=llm, prompt=resp_prompt)
+    chain.api_answer_chain = get_answer_chain
+    return chain.run(question)
+
+

databricks_llm - raw.py

@@ -0,0 +1,103 @@
+import click
+import logging
+import sys
+import yaml
+import os
+from langchain.llms import AzureOpenAI, OpenAI
+from app.api_funcs import get_job_infos, get_run, get_model, \
+    trans_model, batch_mod_permission, prepare_api_docs
+from pathlib import Path
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+handler = logging.StreamHandler(sys.stdout)
+handler.setLevel(logging.DEBUG)
+logger.addHandler(handler)
+
+PATH = Path(os.path.abspath(os.path.dirname(__file__)))
+
+# Open the YAML file
+conf_path = PATH / "app" / 'llm.yaml'
+with open(conf_path) as f:
+    config = yaml.safe_load(f)
+
+# Use AzureOpenAI, if config contains deployment name, otherwise OpenAI
+if config['model'].get('deployment_name', False):
+    llm = AzureOpenAI(**config['model'])
+else:
+    llm = OpenAI(**config['model'])
+
+
+headers = {"Authorization": f"Bearer {os.getenv('DBR_BEARER_TOKEN')}"}
+updated_api_docs = prepare_api_docs()
+
+
+def comma_list(comma_str: str):
+    return comma_str.split(',')
+
+def determine_api_text(updated_api_docs: dict, query: str):
+    pick_api_prompt = """Please return the file name from the list {api_docs}
+                        that best corresponds to the following query: {query}. \
+                        DO NOT EXPLAIN your answer!
+                        """
+    api_docs = os.listdir(PATH / "app" / "dbr_api_docs")
+    selected_api_doc = llm(pick_api_prompt.format(api_docs=api_docs, query=query)).lstrip().rstrip()
+    logger.info(f"\nSelecting the following api document: {selected_api_doc}")
+    api_text = updated_api_docs[selected_api_doc]
+    return api_text, selected_api_doc    
+
+# Add subcommands for commands
+@click.group()
+def cli():
+    pass
+
+@cli.group(help='Run machine learning model.')
+def ml():
+    pass
+
+# Add commands for specific subcommands of 'ml'
+@ml.command(help='Get information about a model.')
+@click.argument('query', type=str)
+def get_model_info(query):
+    # Instruction to get model infos
+    api_text, _ = determine_api_text(updated_api_docs, query)
+    logger.info(get_model(llm, query, api_text, headers))
+
+@ml.command(help='Get information about a model run.')
+@click.argument('run_id', type=str)
+@click.argument('query', type=str)
+def get_run_info(query, run_id):
+    # ID of the model run for which you'd like information.
+    # Which information should be pulled from the run?
+    api_text, _ = determine_api_text(updated_api_docs, query)
+    logger.info(get_run(llm, run_id, query, api_text, headers))
+
+@ml.command(help='Transition a model from one state to another.')
+@click.argument('query', type=str)
+def transition_model(query):
+    # Instruction to transition a model.
+    api_text, _ = determine_api_text(updated_api_docs, query)
+    trans_model(llm, query, api_text, headers)
+
+@cli.command(help='View job history.')
+@click.argument('query', type=str)
+def jobs(query):
+    if ";" not in query:
+        query = query + ";"
+    query, response_query = query.split(";")
+    api_text, _ = determine_api_text(updated_api_docs, query)
+    # The query for the LLM + an optional query for the API response
+    logger.info(get_job_infos(llm, query, response_query, api_text, headers))
+
+@cli.command(help='Manage user permissions.')
+@click.argument('query', type=str)
+@click.argument('jobs', type=comma_list)
+def permissions(jobs, query):
+    api_text, api_name = determine_api_text(updated_api_docs, query)
+    # Add/Get user permissions.
+    batch_mod_permission(
+        logger, llm, updated_api_docs, api_text, api_name, headers,
+        query, jobs=jobs
+    )
+
+if __name__ == '__main__':
+    cli()

requirements.txt

@@ -0,0 +1,3 @@
+click==8.1.2
+langchain==0.0.171
+openai==0.27.6

setup.py

@@ -0,0 +1,17 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='databricks_llm', # name of your package
+    version='0.0.1', # version number
+    py_modules=['databricks_llm', 'app'], # list all .py files used to build your package
+    install_requires=[
+        'Click',
+    ],
+    packages = find_packages(),
+    include_package_data=True,
+    package_data={"app": ["dbr_api_docs/*.txt", "llm.yaml"]},
+    entry_points='''
+        [console_scripts]
+        databricks_llm=databricks_llm:cli
+    ''',
+)