feat: Add a WikiData API Tool Component (#4512)

* feat: Add a WikiData API Tool Component

* [autofix.ci] apply automated fixes

* Update wikidata_api.py

* Fix casing and import in init

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Eric Hare 2024-11-12 07:03:01 -08:00 committed by GitHub
commit 8cd87210ad
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 115 additions and 0 deletions

View file

@ -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",

View file

@ -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