From cfaac694dcb8984e59bf7edb3776233d4bb2eed5 Mon Sep 17 00:00:00 2001 From: Raphael Valdetaro <79842132+raphaelchristi@users.noreply.github.com> Date: Tue, 7 Jan 2025 19:22:54 -0300 Subject: [PATCH] refactor: Update WikidataAPI component to standard output pattern (#5431) * refactor: Update WikidataAPI component to standard output pattern * [autofix.ci] apply automated fixes * Fix formatting errors in wikidata_api.py * fix: Fix linting error in tools init * test: add unit tests for WikidataComponent * refactor: rename WikidataComponent to WikidataAPIComponent --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Eric Hare --- .../langflow/components/tools/wikidata_api.py | 136 ++++++++---------- .../components/tools/test_wikidata_api.py | 106 ++++++++++++++ 2 files changed, 167 insertions(+), 75 deletions(-) create mode 100644 src/backend/tests/unit/components/tools/test_wikidata_api.py diff --git a/src/backend/base/langflow/components/tools/wikidata_api.py b/src/backend/base/langflow/components/tools/wikidata_api.py index d42679090..5356568f7 100644 --- a/src/backend/base/langflow/components/tools/wikidata_api.py +++ b/src/backend/base/langflow/components/tools/wikidata_api.py @@ -1,58 +1,15 @@ -from typing import Any - import httpx -from langchain_core.tools import StructuredTool, ToolException -from pydantic import BaseModel, Field +from httpx import HTTPError +from langchain_core.tools import ToolException -from langflow.base.langchain_utilities.model import LCToolComponent -from langflow.field_typing import Tool -from langflow.inputs import MultilineInput +from langflow.custom import Component +from langflow.helpers.data import data_to_text +from langflow.io import MultilineInput, Output from langflow.schema import Data +from langflow.schema.message import Message -class WikidataSearchSchema(BaseModel): - query: str = Field(..., description="The search query for Wikidata") - - -class WikidataAPIWrapper(BaseModel): - """Wrapper around Wikidata API.""" - - wikidata_api_url: str = "https://www.wikidata.org/w/api.php" - - def results(self, query: str) -> list[dict[str, Any]]: - # Define request parameters for Wikidata API - params = { - "action": "wbsearchentities", - "format": "json", - "search": query, - "language": "en", - } - - # Send request to Wikidata API - response = httpx.get(self.wikidata_api_url, params=params) - response.raise_for_status() - response_json = response.json() - - # Extract and return search results - return response_json.get("search", []) - - def run(self, query: str) -> list[dict[str, Any]]: - try: - results = self.results(query) - if results: - return results - - error_message = "No search results found for the given query." - - raise ToolException(error_message) - - except Exception as e: - error_message = f"Error in Wikidata Search API: {e!s}" - - raise ToolException(error_message) from e - - -class WikidataAPIComponent(LCToolComponent): +class WikidataAPIComponent(Component): display_name = "Wikidata API" description = "Performs a search using the Wikidata API." name = "WikidataAPI" @@ -64,38 +21,67 @@ class WikidataAPIComponent(LCToolComponent): display_name="Query", info="The text query for similarity search on Wikidata.", required=True, + tool_mode=True, ), ] - def build_tool(self) -> Tool: - wrapper = WikidataAPIWrapper() + outputs = [ + Output(display_name="Data", name="data", method="fetch_content"), + Output(display_name="Text", name="text", method="fetch_content_text"), + ] - # Define the tool using StructuredTool and wrapper's run method - tool = StructuredTool.from_function( - name="wikidata_search_api", - description="Perform similarity search on Wikidata API", - func=wrapper.run, - args_schema=WikidataSearchSchema, - ) + def fetch_content(self) -> list[Data]: + try: + # Define request parameters for Wikidata API + params = { + "action": "wbsearchentities", + "format": "json", + "search": self.query, + "language": "en", + } - self.status = "Wikidata Search API Tool for Langchain" + # Send request to Wikidata API + wikidata_api_url = "https://www.wikidata.org/w/api.php" + response = httpx.get(wikidata_api_url, params=params) + response.raise_for_status() + response_json = response.json() - return tool + # Extract search results + results = response_json.get("search", []) - def run_model(self) -> list[Data]: - tool = self.build_tool() + if not results: + return [Data(data={"error": "No search results found for the given query."})] - results = tool.run({"query": self.query}) + # Transform the API response into Data objects + data = [ + Data( + text=f"{result['label']}: {result.get('description', '')}", + data={ + "label": result["label"], + "id": result.get("id"), + "url": result.get("url"), + "description": result.get("description", ""), + "concepturi": result.get("concepturi"), + }, + ) + for result in results + ] - # Transform the API response into Data objects - data = [ - Data( - text=result["label"], - metadata=result, - ) - for result in results - ] + self.status = data + except HTTPError as e: + error_message = f"HTTP Error in Wikidata Search API: {e!s}" + raise ToolException(error_message) from None + except KeyError as e: + error_message = f"Data parsing error in Wikidata API response: {e!s}" + raise ToolException(error_message) from None + except ValueError as e: + error_message = f"Value error in Wikidata API: {e!s}" + raise ToolException(error_message) from None + else: + return data - self.status = data # type: ignore[assignment] - - return data + def fetch_content_text(self) -> Message: + data = self.fetch_content() + result_string = data_to_text("{text}", data) + self.status = result_string + return Message(text=result_string) diff --git a/src/backend/tests/unit/components/tools/test_wikidata_api.py b/src/backend/tests/unit/components/tools/test_wikidata_api.py new file mode 100644 index 000000000..b7a2a1ab3 --- /dev/null +++ b/src/backend/tests/unit/components/tools/test_wikidata_api.py @@ -0,0 +1,106 @@ +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from langchain_core.tools import ToolException +from langflow.components.tools import WikidataAPIComponent +from langflow.custom import Component +from langflow.custom.utils import build_custom_component_template +from langflow.schema import Data +from langflow.schema.message import Message + + +def test_wikidata_initialization(): + component = WikidataAPIComponent() + assert component.display_name == "Wikidata API" + assert component.description == "Performs a search using the Wikidata API." + assert component.icon == "Wikipedia" + + +def test_wikidata_template(): + wikidata = WikidataAPIComponent() + component = Component(_code=wikidata._code) + frontend_node, _ = build_custom_component_template(component) + + # Verify basic structure + assert isinstance(frontend_node, dict) + + # Verify inputs + assert "template" in frontend_node + input_names = [input_["name"] for input_ in frontend_node["template"].values() if isinstance(input_, dict)] + assert "query" in input_names + + +@patch("langflow.components.tools.wikidata_api.httpx.get") +def test_fetch_content_success(mock_httpx): + component = WikidataAPIComponent() + component.query = "test query" + + # Mock successful API response + mock_response = MagicMock() + mock_response.json.return_value = { + "search": [ + { + "label": "Test Label", + "id": "Q123", + "url": "https://test.com", + "description": "Test Description", + "concepturi": "https://test.com/concept", + } + ] + } + mock_httpx.return_value = mock_response + + result = component.fetch_content() + + assert isinstance(result, list) + assert len(result) == 1 + assert result[0].text == "Test Label: Test Description" + assert result[0].data["label"] == "Test Label" + assert result[0].data["id"] == "Q123" + + +@patch("langflow.components.tools.wikidata_api.httpx.get") +def test_fetch_content_empty_response(mock_httpx): + component = WikidataAPIComponent() + component.query = "test query" + + # Mock empty API response + mock_response = MagicMock() + mock_response.json.return_value = {"search": []} + mock_httpx.return_value = mock_response + + result = component.fetch_content() + + assert isinstance(result, list) + assert len(result) == 1 + assert "error" in result[0].data + assert "No search results found" in result[0].data["error"] + + +@patch("langflow.components.tools.wikidata_api.httpx.get") +def test_fetch_content_error_handling(mock_httpx): + component = WikidataAPIComponent() + component.query = "test query" + + # Mock HTTP error + mock_httpx.side_effect = httpx.HTTPError("API Error") + + with pytest.raises(ToolException): + component.fetch_content() + + +def test_fetch_content_text(): + component = WikidataAPIComponent() + component.fetch_content = MagicMock( + return_value=[ + Data(text="First result", data={"label": "Label 1"}), + Data(text="Second result", data={"label": "Label 2"}), + ] + ) + + result = component.fetch_content_text() + + assert isinstance(result, Message) + assert "First result" in result.text + assert "Second result" in result.text