diff --git a/src/backend/base/langflow/components/tools/__init__.py b/src/backend/base/langflow/components/tools/__init__.py index 09723a2a9..2ffabdb20 100644 --- a/src/backend/base/langflow/components/tools/__init__.py +++ b/src/backend/base/langflow/components/tools/__init__.py @@ -16,6 +16,7 @@ from .search_api import SearchAPIComponent from .searxng import SearXNGToolComponent from .serp_api import SerpAPIComponent from .tavily_search import TavilySearchToolComponent +from .wikidata_api import WikidataAPIComponent from .wikipedia_api import WikipediaAPIComponent from .wolfram_alpha_api import WolframAlphaAPIComponent from .yahoo_finance import YfinanceToolComponent @@ -43,6 +44,7 @@ __all__ = [ "SearXNGToolComponent", "SerpAPIComponent", "TavilySearchToolComponent", + "WikidataAPIComponent", "WikipediaAPIComponent", "WolframAlphaAPIComponent", "YfinanceToolComponent", diff --git a/src/backend/base/langflow/components/tools/wikidata_api.py b/src/backend/base/langflow/components/tools/wikidata_api.py new file mode 100644 index 000000000..394bcdba4 --- /dev/null +++ b/src/backend/base/langflow/components/tools/wikidata_api.py @@ -0,0 +1,113 @@ +from typing import Any + +import httpx +from langchain_core.tools import StructuredTool, ToolException +from pydantic import BaseModel, Field + +from langflow.base.langchain_utilities.model import LCToolComponent +from langflow.field_typing import Tool +from langflow.inputs import MultilineInput +from langflow.schema import Data + + +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 not results: + msg = "No search results found for the given query." + + raise ToolException(msg) + + # Process and structure the results + return [ + { + "label": result.get("label", ""), + "description": result.get("description", ""), + "concepturi": result.get("concepturi", ""), + "id": result.get("id", ""), + } + for result in results + ] + + except Exception as e: + error_message = f"Error in Wikidata Search API: {e!s}" + + raise ToolException(error_message) from e + + +class WikidataAPIComponent(LCToolComponent): + display_name = "Wikidata API" + description = "Performs a search using the Wikidata API." + name = "WikidataAPI" + + inputs = [ + MultilineInput( + name="query", + display_name="Query", + info="The text query for similarity search on Wikidata.", + required=True, + ), + ] + + def build_tool(self) -> Tool: + wrapper = WikidataAPIWrapper() + + # 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, + ) + + self.status = "Wikidata Search API Tool for Langchain" + + return tool + + def run_model(self) -> list[Data]: + tool = self.build_tool() + + results = tool.run({"query": self.query}) + + # Transform the API response into Data objects + data = [ + Data( + text=result["label"], + metadata={ + "id": result["id"], + "concepturi": result["concepturi"], + "description": result["description"], + }, + ) + for result in results + ] + + self.status = data # type: ignore[assignment] + + return data