Added Tool HuggingFace API

Author: jobandhanjal

libs/community/langchain_community/tools/huggingface.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+from langchain.tools import BaseTool
+
+class HuggingFaceHubTool(BaseTool):
+    """Base tool for interacting with the Hugging Face Hub."""
+
+    api_client: "HfApi"
+    task: str = ""
+    top_k_results: int = 3
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        try:
+            from huggingface_hub import HfApi
+            self.api_client = HfApi()
+        except ImportError as e:
+            raise ImportError(
+                "huggingface_hub is not installed. Please install it with `pip install huggingface-hub`"
+            ) from e
+
+    def _run(self, query: str) -> str:
+        """Use the tool."""
+        try:
+            if self.task == "models":
+                results_list = self.api_client.list_models(
+                    search=query, top_k=self.top_k_results
+                )
+            elif self.task == "datasets":
+                results_list = self.api_client.list_datasets(
+                    search=query, top_k=self.top_k_results
+                )
+            else:
+                return "Invalid task specified for the Hugging Face Hub tool."
+
+            if not results_list:
+                return f"No {self.task} found on the Hugging Face Hub for '{query}'."
+
+            formatted_results = []
+            for result in results_list:
+                header = f"ID: {getattr(result, 'modelId', getattr(result, 'id', 'N/A'))}"
+                author = f"Author: {getattr(result, 'author', 'N/A')}"
+                tags = f"Tags: {', '.join(getattr(result, 'tags', []))}"
+                formatted_results.append(f"{header}\n{author}\n{tags}")
+
+            return "\n\n---\n\n".join(formatted_results)
+
+        except Exception as e:
+            return f"An error occurred: {e}"
+
+    async def _arun(self, query: str) -> str:
+        """Use the tool asynchronously."""
+        import asyncio
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self._run, query
+        )
+
+
+class HuggingFaceModelSearchTool(HuggingFaceHubTool):
+    """Tool that searches for models on the Hugging Face Hub."""
+
+    name: str = "hugging_face_model_search"
+    description: str = (
+        "Use this tool to search for models on the Hugging Face Hub. "
+        "The input should be a search query string. "
+        "The output will be a formatted string with the top results, "
+        "including their ID, author, and tags."
+    )
+    task: str = "models"
+
+
+class HuggingFaceDatasetSearchTool(HuggingFaceHubTool):
+    """Tool that searches for datasets on the Hugging Face Hub."""
+
+    name: str = "hugging_face_dataset_search"
+    description: str = (
+        "Use this tool to search for datasets on the Hugging Face Hub. "
+        "The input should be a search query string. "
+        "The output will be a formatted string with the top results, "
+        "including their ID, author, and tags."
+    )
+    task: str = "datasets"

libs/community/pyproject.toml

@@ -26,6 +26,11 @@ version = "0.3.29"
 description = "Community contributed LangChain integrations."
 readme = "README.md"
 
+[project.optional-dependencies]
+all = [
+    "huggingface-hub",
+]
+
 [project.urls]
 "Source Code" = "https://github.com/langchain-ai/langchain-community/tree/main/libs/community"
 "Release Notes" = "https://github.com/langchain-ai/langchain/releases?q=tag%3A%22langchain-community%3D%3D0%22&expanded=true"

libs/community/tests/unit_tests/tools/test_huggingface.py

@@ -0,0 +1,68 @@
+from unittest.mock import MagicMock
+import pytest
+from langchain_community.tools.huggingface import (
+    HuggingFaceModelSearchTool,
+    HuggingFaceDatasetSearchTool,
+)
+
+@pytest.fixture
+def mock_hf_api(mocker):
+    """Fixture to mock the HfApi client."""
+    mock_api = MagicMock()
+    # Mock the list_models and list_datasets methods
+    mock_api.list_models.return_value = [
+        MagicMock(modelId="gpt2", author="openai", tags=["text-generation"]),
+        MagicMock(modelId="distilbert-base-uncased", author="distilbert", tags=["fill-mask"]),
+    ]
+    mock_api.list_datasets.return_value = [
+        MagicMock(id="squad", author="stanford", tags=["question-answering"]),
+        MagicMock(id="imdb", author="stanford", tags=["text-classification"]),
+    ]
+    # Patch the HfApi constructor to return our mock object
+    mocker.patch(
+        "langchain_community.tools.huggingface.HfApi", return_value=mock_api
+    )
+    return mock_api
+
+
+def test_huggingface_model_search_tool(mock_hf_api):
+    """Test the model search tool with mocked API."""
+    tool = HuggingFaceModelSearchTool()
+    result = tool.run("test query")
+
+    # Assert that the list_models method was called
+    mock_hf_api.list_models.assert_called_once_with(search="test query", top_k=3)
+
+    # Assert that the output contains the mocked data
+    assert "ID: gpt2" in result
+    assert "Author: openai" in result
+    assert "Tags: text-generation" in result
+    assert "ID: distilbert-base-uncased" in result
+    assert "---" in result
+
+
+def test_huggingface_dataset_search_tool(mock_hf_api):
+    """Test the dataset search tool with mocked API."""
+    tool = HuggingFaceDatasetSearchTool()
+    result = tool.run("another query")
+
+    # Assert that the list_datasets method was called
+    mock_hf_api.list_datasets.assert_called_once_with(search="another query", top_k=3)
+
+    # Assert that the output contains the mocked data
+    assert "ID: squad" in result
+    assert "Author: stanford" in result
+    assert "Tags: question-answering" in result
+    assert "ID: imdb" in result
+    assert "---" in result
+
+
+def test_huggingface_model_search_no_results(mock_hf_api):
+    """Test the model search tool when no results are found."""
+    # Configure the mock to return an empty list for this test
+    mock_hf_api.list_models.return_value = []
+    
+    tool = HuggingFaceModelSearchTool()
+    result = tool.run("empty query")
+
+    assert result == "No models found on the Hugging Face Hub for 'empty query'."