Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add a WikiData API Tool Component #4512

Merged
merged 8 commits into from
Nov 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading