Skip to content

Commit

Permalink
feat: Add a WikiData API Tool Component (langflow-ai#4512)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
2 people authored and diogocabral committed Nov 26, 2024
1 parent 296d006 commit 1815b8e
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 0 deletions.
2 changes: 2 additions & 0 deletions src/backend/base/langflow/components/tools/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
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
Expand Down Expand Up @@ -43,6 +44,7 @@
"SearXNGToolComponent",
"SerpAPIComponent",
"TavilySearchToolComponent",
"WikidataAPIComponent",
"WikipediaAPIComponent",
"WolframAlphaAPIComponent",
"YfinanceToolComponent",
Expand Down
113 changes: 113 additions & 0 deletions src/backend/base/langflow/components/tools/wikidata_api.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 1815b8e

Please sign in to comment.