Add optional model document name field. When this field is provided, that is the document the library will connect to.

Author: mdeyell-valid-mind

validmind/api_client.py

@@ -67,12 +67,15 @@ def get_api_model() -> Optional[str]:
 
 
 def _get_api_headers() -> Dict[str, str]:
-    return {
+    headers = {
         "X-API-KEY": _api_key,
         "X-API-SECRET": _api_secret,
         "X-MODEL-CUID": _model_cuid,
         "X-MONITORING": str(_monitoring),
     }
+    if _model_document:
+        headers["X-DOCUMENT-TYPE"] = _model_document
+    return headers
 
 
 def _get_session() -> aiohttp.ClientSession:
@@ -194,6 +197,7 @@ def init(
     model: Optional[str] = None,
     monitoring: bool = False,
     generate_descriptions: Optional[bool] = None,
+    model_document: Optional[str] = None,
 ):
     """
     Initializes the API client instances and calls the /ping endpoint to ensure
@@ -209,11 +213,12 @@ def init(
         api_secret (str, optional): The API secret. Defaults to None.
         api_host (str, optional): The API host. Defaults to None.
         monitoring (bool): The ongoing monitoring flag. Defaults to False.
-        generate_descriptions (bool): Whether to use GenAI to generate test result descriptions. Defaults to True.
+        generate_descriptions (bool, optional): Whether to use GenAI to generate test result descriptions. Defaults to True.
+        model_document (str, optional): The name of the document.
     Raises:
         ValueError: If the API key and secret are not provided
     """
-    global _api_key, _api_secret, _api_host, _model_cuid, _monitoring
+    global _api_key, _api_secret, _api_host, _model_cuid, _monitoring, _model_document
 
     if api_key == "...":
         # special case to detect when running a notebook placeholder (...)
@@ -238,6 +243,7 @@ def init(
     if generate_descriptions is not None:
         os.environ["VALIDMIND_LLM_DESCRIPTIONS_ENABLED"] = str(generate_descriptions)
 
+    _model_document = model_document
     reload()