Refactor Azure AI Document Intelligence provider to standardize crede…

…ntial keys and update YAML configuration for consistency
fujita-h · Jan 17, 2025 · c6ac604 · c6ac604
1 parent 537fd31
commit c6ac604
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 11 deletions.
diff --git a/provider/azure_ai_document_intelligence.py b/provider/azure_ai_document_intelligence.py
@@ -8,14 +8,14 @@
 class AzureAiDocumentIntelligenceProvider(ToolProvider):
     def _validate_credentials(self, credentials: dict[str, Any]) -> None:
         try:
-            api_endpoint = credentials.get("azure_ai_document_intelligence_api_endpoint")
-            api_key = credentials.get("azure_ai_document_intelligence_api_key")
+            api_endpoint = credentials.get("api_endpoint")
+            api_key = credentials.get("api_key")
 
             # Ensure API key and endpoint are provided
             if not api_endpoint:
-                raise Exception("Azure AI Document Intelligence API endpoint is missing")
+                raise Exception("API endpoint is missing")
             if not api_key:
-                raise Exception("Azure AI Document Intelligence API key is missing")
+                raise Exception("API key is missing")
 
             # Validate the API key and endpoint
             headers = {"Ocp-Apim-Subscription-Key": api_key}
@@ -30,7 +30,7 @@ def _validate_credentials(self, credentials: dict[str, Any]) -> None:
             json_response = response.json()
             if json_response["code"] == "NotFound":
                 return
-            raise Exception("Invalid Azure AI Document Intelligence API key or endpoint." + str(json_response))
+            raise Exception("API endpoint or key." + str(json_response))
 
         except Exception as e:
             raise ToolProviderCredentialValidationError(str(e)) from e
diff --git a/provider/azure_ai_document_intelligence.yaml b/provider/azure_ai_document_intelligence.yaml
@@ -11,7 +11,7 @@ identity:
     - utilities
     - image
 credentials_for_provider:
-  azure_ai_document_intelligence_api_endpoint:
+  api_endpoint:
     type: text-input
     required: true
     label:
@@ -21,7 +21,7 @@ credentials_for_provider:
       en_US: Enter your API Endpoint
       ja_JP: APIエンドポイントを入力
     url: https://portal.azure.com  
-  azure_ai_document_intelligence_api_key:
+  api_key:
     type: secret-input
     required: true
     label:

diff --git a/tools/extract_document.py b/tools/extract_document.py
@@ -15,14 +15,14 @@ def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessag
             raise Exception("Tool runtime or credentials are missing")
 
         # Get endpoint and api key
-        endpoint = str(self.runtime.credentials.get("azure_ai_document_intelligence_api_endpoint"))
-        api_key = str(self.runtime.credentials.get("azure_ai_document_intelligence_api_key"))
+        endpoint = str(self.runtime.credentials.get("api_endpoint"))
+        api_key = str(self.runtime.credentials.get("api_key"))
 
         # Create credential
         credential = AzureKeyCredential(api_key)
 
         # Get file
-        file = tool_parameters.get("file")
+        file = tool_parameters.get("input_file")
         if not file:
             raise ValueError("File is required")
         file_binary = io.BytesIO(file.blob)

diff --git a/tools/extract_document.yaml b/tools/extract_document.yaml
@@ -10,7 +10,7 @@ description:
     ja_JP: ドキュメント/画像をテキストまたはMarkdownに展開
   llm: This tool extract document or image file to text or markdown format.
 parameters:
-  - name: file
+  - name: input_file
     type: file
     required: true
     label: