|
|
import json |
|
|
from typing import Any |
|
|
from urllib.parse import urljoin |
|
|
|
|
|
import httpx |
|
|
from langchain_core.tools import StructuredTool, ToolException |
|
|
from pydantic import BaseModel |
|
|
from pydantic.v1 import Field |
|
|
|
|
|
from langflow.base.langchain_utilities.model import LCToolComponent |
|
|
from langflow.field_typing import Tool |
|
|
from langflow.inputs import IntInput, MultilineInput, NestedDictInput, SecretStrInput, StrInput |
|
|
from langflow.schema import Data |
|
|
|
|
|
|
|
|
class GleanSearchAPISchema(BaseModel): |
|
|
query: str = Field(..., description="The search query") |
|
|
page_size: int = Field(10, description="Maximum number of results to return") |
|
|
request_options: dict[str, Any] | None = Field(default_factory=dict, description="Request Options") |
|
|
|
|
|
|
|
|
class GleanAPIWrapper(BaseModel): |
|
|
"""Wrapper around Glean API.""" |
|
|
|
|
|
glean_api_url: str |
|
|
glean_access_token: str |
|
|
act_as: str = "langflow-component@datastax.com" |
|
|
|
|
|
def _prepare_request( |
|
|
self, |
|
|
query: str, |
|
|
page_size: int = 10, |
|
|
request_options: dict[str, Any] | None = None, |
|
|
) -> dict: |
|
|
|
|
|
url = self.glean_api_url |
|
|
if not url.endswith("/"): |
|
|
url += "/" |
|
|
|
|
|
return { |
|
|
"url": urljoin(url, "search"), |
|
|
"headers": { |
|
|
"Authorization": f"Bearer {self.glean_access_token}", |
|
|
"X-Scio-ActAs": self.act_as, |
|
|
}, |
|
|
"payload": { |
|
|
"query": query, |
|
|
"pageSize": page_size, |
|
|
"requestOptions": request_options, |
|
|
}, |
|
|
} |
|
|
|
|
|
def results(self, query: str, **kwargs: Any) -> list[dict[str, Any]]: |
|
|
results = self._search_api_results(query, **kwargs) |
|
|
|
|
|
if len(results) == 0: |
|
|
msg = "No good Glean Search Result was found" |
|
|
raise AssertionError(msg) |
|
|
|
|
|
return results |
|
|
|
|
|
def run(self, query: str, **kwargs: Any) -> list[dict[str, Any]]: |
|
|
try: |
|
|
results = self.results(query, **kwargs) |
|
|
|
|
|
processed_results = [] |
|
|
for result in results: |
|
|
if "title" in result: |
|
|
result["snippets"] = result.get("snippets", [{"snippet": {"text": result["title"]}}]) |
|
|
if "text" not in result["snippets"][0]: |
|
|
result["snippets"][0]["text"] = result["title"] |
|
|
|
|
|
processed_results.append(result) |
|
|
except Exception as e: |
|
|
error_message = f"Error in Glean Search API: {e!s}" |
|
|
raise ToolException(error_message) from e |
|
|
|
|
|
return processed_results |
|
|
|
|
|
def _search_api_results(self, query: str, **kwargs: Any) -> list[dict[str, Any]]: |
|
|
request_details = self._prepare_request(query, **kwargs) |
|
|
|
|
|
response = httpx.post( |
|
|
request_details["url"], |
|
|
json=request_details["payload"], |
|
|
headers=request_details["headers"], |
|
|
) |
|
|
|
|
|
response.raise_for_status() |
|
|
response_json = response.json() |
|
|
|
|
|
return response_json.get("results", []) |
|
|
|
|
|
@staticmethod |
|
|
def _result_as_string(result: dict) -> str: |
|
|
return json.dumps(result, indent=4) |
|
|
|
|
|
|
|
|
class GleanSearchAPIComponent(LCToolComponent): |
|
|
display_name = "Glean Search API" |
|
|
description = "Call Glean Search API" |
|
|
name = "GleanAPI" |
|
|
icon = "Glean" |
|
|
|
|
|
inputs = [ |
|
|
StrInput( |
|
|
name="glean_api_url", |
|
|
display_name="Glean API URL", |
|
|
required=True, |
|
|
), |
|
|
SecretStrInput(name="glean_access_token", display_name="Glean Access Token", required=True), |
|
|
MultilineInput(name="query", display_name="Query", required=True), |
|
|
IntInput(name="page_size", display_name="Page Size", value=10), |
|
|
NestedDictInput(name="request_options", display_name="Request Options", required=False), |
|
|
] |
|
|
|
|
|
def build_tool(self) -> Tool: |
|
|
wrapper = self._build_wrapper( |
|
|
glean_api_url=self.glean_api_url, |
|
|
glean_access_token=self.glean_access_token, |
|
|
) |
|
|
|
|
|
tool = StructuredTool.from_function( |
|
|
name="glean_search_api", |
|
|
description="Search Glean for relevant results.", |
|
|
func=wrapper.run, |
|
|
args_schema=GleanSearchAPISchema, |
|
|
) |
|
|
|
|
|
self.status = "Glean Search API Tool for Langchain" |
|
|
|
|
|
return tool |
|
|
|
|
|
def run_model(self) -> list[Data]: |
|
|
tool = self.build_tool() |
|
|
|
|
|
results = tool.run( |
|
|
{ |
|
|
"query": self.query, |
|
|
"page_size": self.page_size, |
|
|
"request_options": self.request_options, |
|
|
} |
|
|
) |
|
|
|
|
|
|
|
|
data = [Data(data=result, text=result["snippets"][0]["text"]) for result in results] |
|
|
self.status = data |
|
|
|
|
|
return data |
|
|
|
|
|
def _build_wrapper( |
|
|
self, |
|
|
glean_api_url: str, |
|
|
glean_access_token: str, |
|
|
): |
|
|
return GleanAPIWrapper( |
|
|
glean_api_url=glean_api_url, |
|
|
glean_access_token=glean_access_token, |
|
|
) |
|
|
|
