Update app.py

app.py

@@ -1,103 +1,106 @@
-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()
+#https://www.youtube.com/watch?v=AREBu2B5H3M&t=140s
+
+
+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()