Make mypy happy

Author: pamelafox

app/backend/prepdocs.py

@@ -313,10 +313,10 @@ async def main(strategy: Strategy, setup_index: bool = True):
     )
     blob_manager = setup_blob_manager(
         azure_credential=azd_credential,
-        storage_account=os.getenv("AZURE_STORAGE_ACCOUNT"),
-        storage_container=os.getenv("AZURE_STORAGE_CONTAINER"),
-        storage_resource_group=os.getenv("AZURE_STORAGE_RESOURCE_GROUP"),
-        subscription_id=os.getenv("AZURE_SUBSCRIPTION_ID"),
+        storage_account=os.environ["AZURE_STORAGE_ACCOUNT"],
+        storage_container=os.environ["AZURE_STORAGE_CONTAINER"],
+        storage_resource_group=os.environ["AZURE_STORAGE_RESOURCE_GROUP"],
+        subscription_id=os.environ["AZURE_SUBSCRIPTION_ID"],
         search_images=use_gptvision,
         storage_key=clean_key_if_exists(args.storagekey),
     )
@@ -336,14 +336,17 @@ async def main(strategy: Strategy, setup_index: bool = True):
     elif not openai_host.startswith("azure") and os.getenv("OPENAI_API_KEY"):
         openai_key = os.getenv("OPENAI_API_KEY")
 
+    openai_dimensions = 1536
+    if os.getenv("AZURE_OPENAI_EMB_DIMENSIONS"):
+        openai_dimensions = int(os.environ["AZURE_OPENAI_EMB_DIMENSIONS"])
     openai_embeddings_service = setup_embeddings_service(
         azure_credential=azd_credential,
         openai_host=openai_host,
-        openai_model_name=os.getenv("AZURE_OPENAI_EMB_MODEL_NAME"),
+        openai_model_name=os.environ["AZURE_OPENAI_EMB_MODEL_NAME"],
         openai_service=os.getenv("AZURE_OPENAI_SERVICE"),
         openai_custom_url=os.getenv("AZURE_OPENAI_CUSTOM_URL"),
         openai_deployment=os.getenv("AZURE_OPENAI_EMB_DEPLOYMENT"),
-        openai_dimensions=os.getenv("AZURE_OPENAI_EMB_DIMENSIONS"),
+        openai_dimensions=openai_dimensions,
         openai_key=clean_key_if_exists(openai_key),
         openai_org=os.getenv("OPENAI_ORGANIZATION"),
         disable_vectors=dont_use_vectors,
@@ -358,7 +361,7 @@ async def main(strategy: Strategy, setup_index: bool = True):
             blob_manager=blob_manager,
             document_action=document_action,
             embeddings=openai_embeddings_service,
-            subscription_id=os.getenv("AZURE_SUBSCRIPTION_ID"),
+            subscription_id=os.environ["AZURE_SUBSCRIPTION_ID"],
             search_service_user_assigned_id=args.searchserviceassignedid,
             search_analyzer_name=os.getenv("AZURE_SEARCH_ANALYZER_NAME"),
             use_acls=use_acls,

scripts/adlsgen2setup.py

@@ -165,7 +165,7 @@ async def main(args: Any):
             data_access_control_format = json.load(f)
         command = AdlsGen2Setup(
             data_directory=args.data_directory,
-            storage_account_name=os.getenv("AZURE_ADLS_GEN2_STORAGE_ACCOUNT"),
+            storage_account_name=os.environ["AZURE_ADLS_GEN2_STORAGE_ACCOUNT"],
             filesystem_name="gptkbcontainer",
             security_enabled_groups=args.create_security_enabled_groups,
             credentials=credentials,

scripts/auth_init.py

@@ -172,12 +172,12 @@ async def main():
         print("Not setting up authentication.")
         exit(0)
 
-    if not os.getenv("AZURE_AUTH_TENANT_ID") and not os.getenv("AZURE_TENANT_ID"):
+    auth_tenant = os.getenv("AZURE_AUTH_TENANT_ID", os.getenv("AZURE_TENANT_ID"))
+    if not auth_tenant:
         print(
             "Error: No tenant ID set for authentication. Run `azd env set AZURE_AUTH_TENANT_ID tenant-id` to set the tenant ID."
         )
         exit(1)
-    auth_tenant = os.getenv("AZURE_AUTH_TENANT_ID", os.getenv("AZURE_TENANT_ID"))
     print("Setting up authentication for tenant", auth_tenant)
     credential = AzureDeveloperCliCredential(tenant_id=auth_tenant)
 

scripts/load_azd_env.py

@@ -0,0 +1,23 @@
+import json
+import logging
+import subprocess
+
+from dotenv import load_dotenv
+
+logger = logging.getLogger("scripts")
+
+
+def load_azd_env():
+    """Get path to current azd env file and load file using python-dotenv"""
+    result = subprocess.run("azd env list -o json", shell=True, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise Exception("Error loading azd env")
+    env_json = json.loads(result.stdout)
+    env_file_path = None
+    for entry in env_json:
+        if entry["IsDefault"]:
+            env_file_path = entry["DotEnvPath"]
+    if not env_file_path:
+        raise Exception("No default azd env file found")
+    logger.info(f"Loading azd env from {env_file_path}")
+    load_dotenv(env_file_path, override=True)