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