src/backend/base/langflow/components/tools/wikidata_api.py

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 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):
    display_name = "Wikidata API"
    description = "Performs a search using the Wikidata API."
    name = "WikidataAPI"
    icon = "Wikipedia"

    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=result,
            )
            for result in results
        ]

        self.status = data  # type: ignore[assignment]

        return data