Made the review changes, added an example, and unit test

libs/community/langchain_community/tools/azure_ai_services/azure_file_translation.py

@@ -1,20 +1,21 @@
 from __future__ import annotations
 
-import logging
 import os
-from typing import Any, Dict, Optional
+import logging
+from typing import Optional, Any, Dict
 
-from langchain_core.callbacks import CallbackManagerForToolRun
 from langchain_core.tools import BaseTool
+from langchain_core.callbacks import CallbackManagerForToolRun
 from langchain_core.utils import get_from_dict_or_env
+from pydantic import model_validator
 
 from langchain_community.document_loaders import (
-    UnstructuredExcelLoader,
-    UnstructuredHTMLLoader,
-    UnstructuredPDFLoader,
     UnstructuredPowerPointLoader,
     UnstructuredWordDocumentLoader,
+    UnstructuredPDFLoader,
+    UnstructuredExcelLoader,
     UnstructuredXMLLoader,
+    UnstructuredHTMLLoader
 )
 
 logger = logging.getLogger(__name__)
@@ -29,16 +30,21 @@ class AzureFileTranslateTool(BaseTool):
     text_translation_key: str = ""
     text_translation_endpoint: str = ""
     target_language: str = "en"
+    region: str = ""
+    file_path: str = ""
     translate_client: Any
 
-    name: str = "azure_document_translation"
-    description: str = """
-        A Wrapper around Azure AI Services can be used to
+    name: str = "azure_file_translation"
+    description: str = (
+        """
+        A Wrapper around Azure AI Services that can be used to
         translate a document into a specific language.
         It reads the text from a file, processes it,
         and then outputs with the desired language.
         """
+    )
 
+    @model_validator(mode="before")
     @classmethod
     def validate_environment(cls, values: Dict) -> Any:
         """
@@ -51,6 +57,11 @@ def validate_environment(cls, values: Dict) -> Any:
             values, "text_translation_endpoint", "AZURE_TRANSLATE_ENDPOINT"
         )
 
+        region = get_from_dict_or_env(
+            values, "region", "AZURE_REGION"
+        )
+
+
         try:
             from azure.ai.translation.text import TextTranslationClient
             from azure.core.credentials import AzureKeyCredential
@@ -59,6 +70,7 @@ def validate_environment(cls, values: Dict) -> Any:
             values["translate_client"] = TextTranslationClient(
                 endpoint=azure_translate_endpoint,
                 credential=AzureKeyCredential(azure_translate_key),
+                region=region
             )
 
         except ImportError:
@@ -93,13 +105,12 @@ def _read_text_from_file(self, file_path: str) -> str:
             ".pptx": UnstructuredPowerPointLoader,
             ".xlsx": UnstructuredExcelLoader,
             ".xml": UnstructuredXMLLoader,
-            ".html": UnstructuredHTMLLoader,
+            ".html": UnstructuredHTMLLoader
         }
 
         loader_class = loader_map.get(file_extension)
 
         if file_extension == ".txt":
-            # Handle plain text files directly
             return self._read_text(file_path)
         elif loader_class is None:
             raise ValueError(f"Unsupported file type: {file_extension}")
@@ -131,32 +142,35 @@ def _translate_text(self, text: str, target_language: Optional[str] = None) -> s
         Raises:
             RuntimeError: If the translation request fails.
         """
+        if not self.translate_client:
+            values = self.validate_environment({})
+            self.translate_client = values["translate_client"]
+
         if target_language is None:
             target_language = self.target_language
-
         try:
             from azure.ai.translation.text.models import InputTextItem
         except ImportError:
-            raise ImportError("Run 'pip install azure-ai-translation-text'.")
+            raise ImportError(
+                "Run 'pip install azure-ai-translation-text'.")
 
         try:
             request_body = [InputTextItem(text=text)]
             response = self.translate_client.translate(
-                content=request_body, to=[target_language]
-            )
+                body=request_body, to_language=[target_language])
 
             translations = response[0].translations
             if translations:
                 return translations[0].text
-            return ""  # No translations found
+            return ""
         except Exception as e:
             raise RuntimeError(f"An error occurred during translation: {e}")
 
-    def _run(
-        self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
-    ) -> str:
-        """ "Run the tool"""
+    def _run(self, query: str,
+             run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
+        """"Run the tool"""
         try:
-            return self._translate_text(query)
+            text = self._read_text_from_file(self.file_path or query)
+            return self._translate_text(text)
         except Exception as e:
             raise RuntimeError(f"Error while running AzureFileTranslateTool: {e}")

libs/community/tests/unit_tests/tools/azure_ai_services/test_azure_file_translation.py

@@ -0,0 +1,105 @@
+from pathlib import Path
+from typing import Any
+
+import pytest
+from langchain_community.tools.azure_ai_services.azure_file_translation import (
+    AzureFileTranslateTool
+)
+
+_THIS_DIR = Path(__file__).parents[3]
+
+_EXAMPLES_DIR = _THIS_DIR / "examples"
+AZURE_PDF = _EXAMPLES_DIR / "test_azure.pdf"
+
+
+@pytest.mark.requires("azure-ai-translation-text")
+def test_tool_initialization(mocker: Any) -> None:
+    mocker.patch("azure.core.credentials.AzureKeyCredential", autospec=True)
+
+    mock_translate_client = mocker.Mock()
+    mocker.patch("azure.ai.translation.text.TextTranslationClient",
+                 return_value=mock_translate_client)
+
+    key = "key"
+    endpoint = "endpoint"
+    region = "westus2"
+
+    tool = AzureFileTranslateTool(
+        text_translation_key=key,
+        text_translation_endpoint=endpoint,
+        region=region,
+        translate_client=mock_translate_client
+    )
+
+    assert tool.text_translation_key == key
+    assert tool.text_translation_endpoint == endpoint
+    assert tool.region == region
+    assert tool.translate_client == mock_translate_client
+
+
+@pytest.mark.requires("azure-ai-translation-text")
+def test_translation_with_file(mocker: Any) -> None:
+    key = "key"
+    endpoint = "endpoint"
+    region = "westus2"
+
+    mocker.patch("azure.core.credentials.AzureKeyCredential",
+                 autospec=True)
+
+    mock_translate_client = mocker.Mock()
+    mocker.patch("azure.ai.translation.text.TextTranslationClient",
+                 return_value=mock_translate_client)
+
+    tool = AzureFileTranslateTool(
+        text_translation_key=key,
+        text_translation_endpoint=endpoint,
+        region=region,
+        translate_client=mock_translate_client
+    )
+
+    mock_translate_client.translate.return_value = [
+        {
+            'detectedLanguage': {'language': 'en', 'score': 1.0},
+            'translations': [
+                {'text': 'Hola, mi nombre es Azure', 'to': 'es'}
+            ]
+        }
+    ]
+
+    file_input: str = str(AZURE_PDF)
+    expected_output = "Hola, mi nombre es Azure"
+
+    result = tool._run(file_input)
+
+    assert result == expected_output
+
+
+@pytest.mark.requires("azure-ai-translation-text")
+def test_translation_with_no_file(mocker: Any) -> None:
+    key = "key"
+    endpoint = "endpoint"
+    region = "westus2"
+
+    mocker.patch("azure.core.credentials.AzureKeyCredential",
+                 autospec=True)
+
+    mock_translate_client = mocker.Mock()
+    mocker.patch("azure.ai.translation.text.TextTranslationClient",
+                 return_value=mock_translate_client)
+
+    tool = AzureFileTranslateTool(
+        text_translation_key=key,
+        text_translation_endpoint=endpoint,
+        region=region,
+        translate_client=mock_translate_client
+    )
+
+    file_input: str = ""
+    expected_output = "Error while running AzureFileTranslateTool"
+
+    try:
+        result = tool._run(file_input)
+    except RuntimeError as e:
+        result = str(e)
+
+    assert expected_output in result