[ai] Fix typing errors to make ai SDKs mypy-compliant (#33638)

* Add first batch of mypy fixes

* Turn on mypy gate

* Add missing Optional import

* Make __init__ return type None

* Make Optional explicit

* Make __init__ return type None

* Update pyproject.toml

* Make Optional explicit

* Ignore incorrect mypy error - code is fine

* Exclude _restclient from mypy

* Fix type mismatch

* Fix connection_subtypes errors

* Fix tokens errors

* Exclude vendor folder

* Fix _index_data_source erros

* Fix entities/data errors

* Fix some _deployment_utils errors

* Fix _ai_client_utils errors

* Fix logging_handler errors

* Fix _utils/logging errors

* Fix ai_client errors

* Fix _deployment operations errors

* Exclude tests from mypy check

* Fix _index._docstore errors

* Fix _model_operations errors

* Fix faiss errors

* Fix _model_operations errors

* Add type annotation

* Fix _ai_resource_operations errors

* Fix entities._mlindex errors

* Fix _dockerfile_utils errors

* Fix data operations errors

* Fix _index._langchain._docstore errors

* Fix entities.configs errors

* Fix _index._mlindex errors

* Fix some _index._embeddings.__init__ errors

* Fix _index._utils.connections errors

* fix operations._ingest_data_to_index error

* Fix connection_operations errors

* Update pyproject.toml files to not follow imports

* Update pyproject.toml files

* Fix some index._embeddings.__init__ errors

* Resolve merge conflicts

* Create work items and ignore errors

* Update pyproject.toml to call out unused ignore

* Fix deployment utils and deployment operations bugs

* Revert "Resolve merge conflicts"

This reverts commit 09bdc79.

* Resolve merge conflict

* Fix typing error in _documents/document.py

* Add cspell.json files

* Disabled following stubs

* Disabled following stubs

* Fix _index/_langchain/acs error

* Fix data_index_func errors

* Fix typing error in _utils/connections

* Ignore missing stubs from imports error

* Fix dsl/pipeline_decorator errors

* Fix _langchain/acs errors

* Fix _index/_embeddings/openai

* Ignore import stub errors

* Fix _index/models errors

* Fix evaluate/_utils errors

* Fix _model_tools/images

* Fix _model_tools/images

* Fix _index_data_source errors

* Fix resources/_index/_embeddings/__init__

* Fix _mlindex_operations

* Ignore _langchain/docstore error and assign bug

* Fix errors

* Fix dsl/_pipeline_decorator errors

* Fix errors

* Fix evaluate errors

* Fix import errors

* Fix _embeddings/openai errors

* Exclude vendor code from mypy check

* Fix _documents/chunking errors

* Fix _builder/data_index_func errors

* Fix _tasks/crack_and_chunk_and_embed_and_index error

* Fix _tasks/crack_and_chunk_and_embed errors

* Update excluded files/directories

* Fix _documents/cracking errors

* Fix index/_tasks/update_acs errors

* Fix _documents/__init__ errors

* Fix _documnets/__init__ errors

* Fix _tasks/update_acs errors

* Fix index/_mlindex errors

* Fix synthetic/qa errors

* Fix index/_docstore errors

* Fix faiss errors

* Fix _langchain/acs errors

* Fix index/_models and _conversation/conversation_request errors

* Fix evaluate/_evaluation_result error

* Fix _tasks/crack_and_chunk errors

* Ignore _tasks/generate_qa error

* Fix _tasks/update_pinecone errors

* Fix index/_embeddings/__init__ errors

* Fix _langchain/docstore errors

* Fix type: ignore comment

* Fix remaining errors and add TODO

Author: diondrapeck

sdk/ai/azure-ai-generative/azure/ai/generative/_telemetry/logging_handler.py

@@ -9,6 +9,7 @@
 import logging
 import platform
 import traceback
+from typing import Union, Dict
 import os
 
 from opencensus.ext.azure.log_exporter import AzureLogHandler
@@ -45,7 +46,7 @@ def __init__(self, name: str):
         self.package_logger: logging.Logger = logging.getLogger(GEN_AI_INTERNAL_LOGGER_NAMESPACE + name)
         self.package_logger.propagate = False
         self.module_logger = logging.getLogger(name)
-        self.custom_dimensions = {}
+        self.custom_dimensions: Dict[str, Union[str, Dict]] = {}
 
     def update_info(self) -> None:
         self.package_logger.addHandler(get_appinsights_log_handler(USER_AGENT))

sdk/ai/azure-ai-generative/azure/ai/generative/evaluate/_evaluate.py

@@ -8,7 +8,7 @@
 import time
 import logging
 from pathlib import Path
-from typing import Callable, Optional, Dict, List
+from typing import Callable, Optional, Dict, List, Mapping
 
 import mlflow
 import pandas as pd
@@ -87,13 +87,13 @@ def _log_metrics(run_id, metrics):
 @monitor_with_activity(package_logger, "Evaluate", ActivityType.PUBLICAPI)
 def evaluate(
         *,
-        evaluation_name: str = None,
+        evaluation_name: Optional[str] = None,
         target: Optional[Callable] = None,
         data: Optional[str] = None,
-        task_type: str = None,
+        task_type: Optional[str] = None,
         metrics_list: Optional[List[str]] = None,
-        model_config: Dict[str, str] = None,
-        data_mapping: Dict[str, str] = None,
+        model_config: Optional[Dict[str, str]] = None,
+        data_mapping: Optional[Mapping] = None,
         output_path: Optional[str] = None,
         **kwargs
 ):
@@ -114,7 +114,7 @@ def evaluate(
     :keyword model_config: GPT configuration details needed for AI-assisted metrics.
     :paramtype model_config: Dict[str, str]
     :keyword data_mapping: GPT configuration details needed for AI-assisted metrics.
-    :paramtype data_mapping: Dict[str, str]
+    :paramtype data_mapping: typing.Mapping
     :keyword output_path: The local folder path to save evaluation artifacts to if set
     :paramtype output_path: Optional[str]
     :keyword tracking_uri: Tracking uri to log evaluation results to AI Studio

sdk/ai/azure-ai-generative/azure/ai/generative/evaluate/_evaluation_result.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Optional
 
 from azure.ai.generative.evaluate._utils import _get_ai_studio_url
 
@@ -8,8 +8,8 @@ class EvaluationResult(object):
     def __init__(self, metrics_summary: Dict[str, float], artifacts: Dict[str, str], **kwargs):
         self._metrics_summary = metrics_summary
         self._artifacts = artifacts
-        self._tracking_uri = kwargs.get("tracking_uri")
-        self._evaluation_id = kwargs.get("evaluation_id")
+        self._tracking_uri: Optional[str] = kwargs.get("tracking_uri")
+        self._evaluation_id: str = kwargs.get("evaluation_id", "")
         if self._tracking_uri:
             self._studio_url = _get_ai_studio_url(self._tracking_uri, self._evaluation_id)
 
@@ -23,13 +23,13 @@ def artifacts(self) -> Dict[str, str]:
 
     @property
     def tracking_uri(self) -> str:
-        return self._tracking_uri
+        return self._tracking_uri  # type: ignore[return-value]
 
     @property
     def studio_url(self) -> str:
         return self._studio_url
 
-    def download_evaluation_artifacts(self, path: str) -> str:
+    def download_evaluation_artifacts(self, path: str) -> None:
         from mlflow.artifacts import download_artifacts
         for artifact, artifact_uri in self.artifacts.items():
             download_artifacts(

sdk/ai/azure-ai-generative/azure/ai/generative/evaluate/_utils.py

@@ -7,6 +7,7 @@
 import re
 import shutil
 from pathlib import Path
+from typing import Optional
 import pandas as pd
 
 _SUB_ID = "sub-id"
@@ -108,20 +109,20 @@ def _write_properties_to_run_history(properties: dict, logger) -> None:
         logger.error("Fail writing properties '%s' to run history: %s", properties, e)
 
 
-def _get_ai_studio_url(tracking_uri, evaluation_id):
+def _get_ai_studio_url(tracking_uri: str, evaluation_id: str):
     _PROJECT_INFO_REGEX = (
         r".*/subscriptions/(.+)/resourceGroups/(.+)"
         r"/providers/Microsoft.MachineLearningServices/workspaces/([^/]+)"
     )
 
     pattern = re.compile(_PROJECT_INFO_REGEX)
 
-    mo = pattern.match(tracking_uri)
+    mo: Optional[re.Match[str]] = pattern.match(tracking_uri)
 
     ret = {}
-    ret[_SUB_ID] = mo.group(1)
-    ret[_RES_GRP] = mo.group(2)
-    ret[_WS_NAME] = mo.group(3)
+    ret[_SUB_ID] = mo.group(1) if mo else mo
+    ret[_RES_GRP] = mo.group(2) if mo else mo
+    ret[_WS_NAME] = mo.group(3) if mo else mo
 
     studio_base_url = os.getenv("AI_STUDIO_BASE_URL", "https://ai.azure.com")
 

sdk/ai/azure-ai-generative/azure/ai/generative/index/_build_mlindex.py

@@ -3,9 +3,9 @@
 # ---------------------------------------------------------
 
 from pathlib import Path
-from typing import Union
+from typing import Dict, Optional, Union
 
-import yaml
+import yaml  # type: ignore[import]
 
 from azure.ai.resources.entities.mlindex import Index
 from azure.ai.resources.operations._index_data_source import ACSSource, LocalSource
@@ -15,24 +15,20 @@
 
 def build_index(
     *,
-    ######## required args ##########
     output_index_name: str,
     vector_store: str,
-    ######## embedding model information ##########
-    embeddings_model: str = None,
-    aoai_connection_id: str = None,
-    ######## chunking information ##########
-    data_source_url: str = None,
+    index_input_config: Union[ACSSource, LocalSource],
+    acs_config: ACSOutputConfig,  # todo better name?
+    embeddings_model: str,
+    aoai_connection_id: Optional[str] = None,
+    data_source_url: Optional[str] = None,
     chunk_size: int = 1024,
     chunk_overlap: int = 0,
     input_glob: str = "**/*",
-    max_sample_files: int = None,
-    chunk_prepend_summary: bool = None,
-    document_path_replacement_regex: str = None,
-    embeddings_cache_path: str = None,
-    ######## data source info ########
-    index_input_config: Union[ACSSource, LocalSource] = None,
-    acs_config: ACSOutputConfig = None,  # todo better name?
+    max_sample_files: Optional[int] = None,
+    chunk_prepend_summary: Optional[bool] = None,
+    document_path_replacement_regex: Optional[Dict[str, str]] = None,
+    embeddings_cache_path: Optional[str] = None,
 ) -> Index:
 
     """Generates embeddings locally and stores Index reference in memory
@@ -53,12 +49,13 @@ def build_index(
     if vector_store == "azure_cognitive_search" and isinstance(index_input_config, ACSSource):
         return _create_mlindex_from_existing_acs(
             output_index_name=output_index_name,
-            embedding_model=embeddings_model,
+            # TODO: Fix Bug 2818331
+            embedding_model=embeddings_model,  # type: ignore[no-redef,arg-type]
             aoai_connection=aoai_connection_id,
             acs_config=index_input_config,
         )
-    embeddings_cache_path = Path(embeddings_cache_path) if embeddings_cache_path else Path.cwd()
-    save_path = embeddings_cache_path/f"{output_index_name}-mlindex"
+    embeddings_cache_path = str(Path(embeddings_cache_path) if embeddings_cache_path else Path.cwd())
+    save_path = str(Path(embeddings_cache_path)/f"{output_index_name}-mlindex")
     splitter_args= {
         'chunk_size': chunk_size,
         'chunk_overlap': chunk_overlap,
@@ -101,6 +98,7 @@ def build_index(
         uri=embeddings_model,
         **connection_args,
     )
+
     embeddings = embedder.embed(chunked_docs)
 
     if vector_store.lower() == "faiss":
@@ -165,7 +163,7 @@ def build_index(
 def _create_mlindex_from_existing_acs(
     output_index_name: str,
     embedding_model: str,
-    aoai_connection: str,
+    aoai_connection: Optional[str],
     acs_config: ACSSource,
 ) -> Index:
     try:
@@ -209,6 +207,7 @@ def _create_mlindex_from_existing_acs(
     if acs_config.acs_metadata_key:
         mlindex_config["index"]["field_mapping"]["metadata"] = acs_config.acs_metadata_key
 
+    model_connection_args: Dict[str, Optional[Union[str, Dict]]]
     if not aoai_connection:
         import openai
         model_connection_args = {
@@ -220,7 +219,7 @@ def _create_mlindex_from_existing_acs(
             "connection": {"id": aoai_connection}
         }
 
-    embedding = EmbeddingsContainer.from_uri(embedding_model, **model_connection_args)
+    embedding = EmbeddingsContainer.from_uri(embedding_model, credential=None, **model_connection_args)
     mlindex_config["embeddings"] = embedding.get_metadata()
 
     path = Path.cwd() / f"import-acs-{acs_config.acs_index_name}-mlindex"

sdk/ai/azure-ai-generative/azure/ai/generative/index/_dataindex/data_index/models.py

@@ -3,6 +3,7 @@
 # ---------------------------------------------------------
 """DataIndex embedding model helpers."""
 import re
+from typing import Optional
 
 OPEN_AI_PROTOCOL_TEMPLATE = "azure_open_ai://deployment/{}/model/{}"
 OPEN_AI_PROTOCOL_REGEX_PATTERN = OPEN_AI_PROTOCOL_TEMPLATE.format(".*", ".*")
@@ -13,7 +14,7 @@
 HUGGINGFACE_PROTOCOL_REGEX_PATTERN = HUGGINGFACE_PROTOCOL_TEMPLATE.format(".*")
 
 
-def build_model_protocol(model: str = None):
+def build_model_protocol(model: Optional[str] = None):
     """Build a model protocol from user input."""
     if not model or re.match(OPEN_AI_PROTOCOL_REGEX_PATTERN, model, re.IGNORECASE):
         return model

sdk/ai/azure-ai-generative/azure/ai/generative/index/_dataindex/dsl/_pipeline_decorator.py

@@ -57,7 +57,8 @@
 
 # Overload the returns a decorator when func is None
 @overload
-def pipeline(
+def pipeline(  # type: ignore[misc]
+    # TODO: Bug 2876412
     func: None = None,
     *,
     name: Optional[str] = None,
@@ -74,7 +75,7 @@ def pipeline(
 # Overload the returns a decorated function when func isn't None
 @overload
 def pipeline(
-    func: Callable[P, T] = None,
+    func: Optional[Callable[P, T]] = None,
     *,
     name: Optional[str] = None,
     version: Optional[str] = None,
@@ -137,7 +138,8 @@ def pipeline(
     get_component = kwargs.get("get_component", False)
 
     def pipeline_decorator(func: Callable[P, T]) -> Callable[P, PipelineJob]:
-        if not isinstance(func, Callable):  # pylint: disable=isinstance-second-argument-not-valid-type
+        # pylint: disable=isinstance-second-argument-not-valid-type
+        if not isinstance(func, Callable):  # type: ignore
             raise UserErrorException(f"Dsl pipeline decorator accept only function type, got {type(func)}.")
 
         non_pipeline_inputs = kwargs.get("non_pipeline_inputs", []) or kwargs.get("non_pipeline_parameters", [])
@@ -246,9 +248,9 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> PipelineJob:
 
             return built_pipeline
 
-        wrapper._is_dsl_func = True
-        wrapper._job_settings = job_settings
-        wrapper._pipeline_builder = pipeline_builder
+        wrapper._is_dsl_func = True  # type: ignore[attr-defined]
+        wrapper._job_settings = job_settings  # type: ignore[attr-defined]
+        wrapper._pipeline_builder = pipeline_builder  # type: ignore[attr-defined]
         return wrapper
 
     # enable use decorator without "()" if all arguments are default values

sdk/ai/azure-ai-generative/azure/ai/generative/index/_dataindex/entities/_assets/_artifacts/data.py

@@ -5,7 +5,8 @@
 from azure.ai.ml.entities._assets._artifacts.data import Data
 
 
-@classmethod
+@classmethod  # type: ignore[misc]
+# TODO: Bug 2874139
 def _resolve_cls_and_type(cls, data, params_override):
     from azure.ai.ml.entities._data_import.data_import import DataImport
     from azure.ai.generative.index._dataindex.entities.data_index import DataIndex