WebSearchFree/integrations/langchain/wsf_tool.py at main · drmikecrypto/WebSearchFree · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
"""LangChain tool wrapper for WebSearchFree HTTP API.

Requires optional dependency: pip install langchain-core

Usage:
    from integrations.langchain.wsf_tool import WebSearchFreeTool
    tool = WebSearchFreeTool(base_url="http://127.0.0.1:8080")
    print(tool.invoke({"query": "open source metasearch"}))
"""

from __future__ import annotations

import json
import os
from typing import Any, Optional, Type

try:
    from langchain_core.tools import BaseTool
    from pydantic import BaseModel, Field
except ImportError as ex:  # pragma: no cover
    raise ImportError(
        "langchain-core is required for WebSearchFreeTool. "
        "Install with: pip install langchain-core"
    ) from ex

import urllib.error
import urllib.request


class WebSearchFreeInput(BaseModel):
    query: str = Field(..., description="Web search query")
    max_results: int = Field(5, description="Maximum number of results")
    include_raw_content: bool = Field(
        False, description="If true, fetch and attach page text for each result"
    )
    search_depth: str = Field("basic", description="basic or advanced (auto-extract)")
    topic: str = Field("general", description="general or news")
    include_domains: list[str] = Field(default_factory=list)
    exclude_domains: list[str] = Field(default_factory=list)
    include_answer: bool = Field(True, description="Extractive answer from snippets")


class WebSearchFreeTool(BaseTool):
    name: str = "websearchfree"
    description: str = (
        "Free keyless web search via local WebSearchFree (Tavily-shaped). "
        "Use for live web results without API keys."
    )
    args_schema: Type[BaseModel] = WebSearchFreeInput
    base_url: str = Field(
        default_factory=lambda: os.environ.get("WSF_BASE_URL", "http://127.0.0.1:8080").rstrip("/")
    )

    def _run(
        self,
        query: str,
        max_results: int = 5,
        include_raw_content: bool = False,
        search_depth: str = "basic",
        topic: str = "general",
        include_domains: Optional[list[str]] = None,
        exclude_domains: Optional[list[str]] = None,
        include_answer: bool = True,
        run_manager: Optional[Any] = None,
    ) -> str:
        del run_manager
        payload: dict[str, Any] = {
            "query": query,
            "max_results": max_results,
            "include_raw_content": include_raw_content,
            "search_depth": search_depth,
            "topic": topic,
            "include_answer": include_answer,
        }
        if include_domains:
            payload["include_domains"] = include_domains
        if exclude_domains:
            payload["exclude_domains"] = exclude_domains
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(
            self.base_url + "/search",
            data=data,
            method="POST",
            headers={"Content-Type": "application/json"},
        )
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                body = json.loads(resp.read().decode("utf-8"))
        except urllib.error.URLError as ex:
            return json.dumps({"error": f"WebSearchFree unreachable at {self.base_url}: {ex}"})
        return json.dumps(body, indent=2)

    async def _arun(self, *args: Any, **kwargs: Any) -> str:
        return self._run(*args, **kwargs)


class WebSearchFreeExtractInput(BaseModel):
    urls: list[str] = Field(..., description="Page URLs to extract (max 10)")


class WebSearchFreeExtractTool(BaseTool):
    name: str = "websearchfree_extract"
    description: str = "Extract main text from web pages via local WebSearchFree."
    args_schema: Type[BaseModel] = WebSearchFreeExtractInput
    base_url: str = Field(
        default_factory=lambda: os.environ.get("WSF_BASE_URL", "http://127.0.0.1:8080").rstrip("/")
    )

    def _run(self, urls: list[str], run_manager: Optional[Any] = None) -> str:
        del run_manager
        payload = {"urls": urls}
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(
            self.base_url + "/extract",
            data=data,
            method="POST",
            headers={"Content-Type": "application/json"},
        )
        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                body = json.loads(resp.read().decode("utf-8"))
        except urllib.error.URLError as ex:
            return json.dumps({"error": f"WebSearchFree unreachable at {self.base_url}: {ex}"})
        return json.dumps(body, indent=2)

    async def _arun(self, *args: Any, **kwargs: Any) -> str:
        return self._run(*args, **kwargs)