isage-middleware 0.2.4.3__cp311-none-any.whl

sage/middleware/operators/rag/searcher.py

@@ -0,0 +1,37 @@
+from typing import Any
+
+import requests
+
+from sage.common.core.functions import MapFunction as MapOperator
+
+
+class BochaWebSearch(MapOperator):
+    def __init__(self, config: dict[str, Any], **kwargs):
+        super().__init__(**kwargs)
+        self.api_key = config.get("api_key")
+        self.count = config.get("count", 10)
+        self.page = config.get("page", 1)
+        self.summary = config.get("summary", True)
+        self.url = "https://api.bochaai.com/v1/web-search"
+
+        if not self.api_key:
+            raise ValueError("BochaWebSearch requires an 'api_key' in config.")
+
+    def execute(self, data: str) -> dict[str, Any]:
+        query = data
+        headers = {"Authorization": self.api_key, "Content-Type": "application/json"}
+        payload = {
+            "query": query,
+            "summary": self.summary,
+            "count": self.count,
+            "page": self.page,
+        }
+
+        try:
+            response = requests.post(self.url, headers=headers, json=payload)
+            response.raise_for_status()
+            result = response.json()
+            return result
+        except Exception as e:
+            self.logger.error(f"BochaWebSearch error: {e}", exc_info=True)
+            return {}  # Return empty dict on error

sage/middleware/operators/rag/types.py

@@ -0,0 +1,28 @@
+"""Compatibility shim for RAG type definitions.
+
+Import from ``sage.libs.rag.types`` instead of middleware.
+"""
+
+from sage.libs.rag.types import (  # noqa: F401
+    RAGDocument,
+    RAGInput,
+    RAGOutput,
+    RAGQuery,
+    RAGResponse,
+    create_rag_response,
+    ensure_rag_response,
+    extract_query,
+    extract_results,
+)
+
+__all__ = [
+    "RAGDocument",
+    "RAGQuery",
+    "RAGResponse",
+    "RAGInput",
+    "RAGOutput",
+    "ensure_rag_response",
+    "extract_query",
+    "extract_results",
+    "create_rag_response",
+]

sage/middleware/operators/rag/writer.py

@@ -0,0 +1,80 @@
+from sage.common.core.functions import MapFunction as MapOperator
+
+
+class MemoryWriter(MapOperator):
+    def __init__(self, config: dict, **kwargs):
+        super().__init__(config, **kwargs)
+        self.state = None
+        self.config = config
+        # 初始化各类型集合
+        self.collections = {}
+
+        # 配置STM
+        if self.config.get("stm", False):
+            stm_config = self.config.get("stm_config", {})
+            self.collections["stm"] = {
+                "collection": self.config.get("stm_collection"),
+                "config": stm_config,
+            }
+
+        # 配置LTM
+        if self.config.get("ltm", False):
+            ltm_config = self.config.get("ltm_config", {})
+            self.collections["ltm"] = {
+                "collection": self.config.get("ltm_collection"),
+                "config": ltm_config,
+            }
+
+        # 配置DCM
+        if self.config.get("dcm", False):
+            dcm_config = self.config.get("dcm_config", {})
+            self.collections["dcm"] = {
+                "collection": self.config.get("dcm_collection"),
+                "config": dcm_config,
+            }
+        # TODO: 在runtime_context中增加状态管理
+        # Issue URL: https://github.com/intellistream/SAGE/issues/235
+
+    def execute(self, data: str | list[str] | tuple[str, str]):
+        input_data = data
+
+        # 统一数据类型处理
+        processed_data = []
+        if isinstance(input_data, list):
+            processed_data = input_data
+        elif isinstance(input_data, tuple) and len(input_data) == 2:
+            processed_data = [f"{input_data[0]}{input_data[1]}"]  # 拼接元组
+        elif isinstance(input_data, str):
+            processed_data = [input_data]
+        else:
+            self.logger.error(f"Unsupported data type: {type(input_data)}")
+            return data
+
+        # 写入所有启用的集合
+        for mem_type, settings in self.collections.items():
+            collection = settings["collection"]
+            config = settings["config"]
+            if not collection:
+                self.logger.warning(f"{mem_type.upper()} collection not initialized")
+                continue
+
+            try:
+                # TODO: 这里的实现实际上要成为由writer 这个function主动往memory manager function发送一个数据。
+                # 而 memory manager function拿到这个数据之后就会去执行 `execute' method 即可实现记忆的读写。
+                # 这里可能会有一个由于调度原因导致的阻塞 -- 可以被优化，请参考MorphStream！
+                if self.state is not None:
+                    self.state.store(
+                        collection=collection,
+                        documents=processed_data,
+                        collection_config=config,
+                    )
+                    self.logger.debug(f"Stored {len(processed_data)} chunks to {mem_type.upper()}")
+                else:
+                    self.logger.warning(
+                        f"State manager not initialized. Cannot store to {mem_type.upper()}. "
+                        "See TODO: https://github.com/intellistream/SAGE/issues/235"
+                    )
+            except Exception as e:
+                self.logger.error(f"Failed to store to {mem_type.upper()}: {str(e)}")
+
+        return data  # 返回原始数据

sage/middleware/operators/tools/__init__.py

@@ -0,0 +1,71 @@
+"""
+Tool Operators
+
+This module contains domain-specific tool operators:
+- Search tools (web search, document search)
+- Data extraction tools
+
+These operators inherit from base operator classes in sage.kernel.operators
+and implement tool-specific business logic.
+
+Note: Some tools require heavy dependencies (torch, transformers).
+      They are loaded lazily and will raise ImportError if dependencies are missing.
+"""
+
+import warnings
+from typing import TYPE_CHECKING
+
+# Core tools (minimal dependencies)
+from sage.middleware.operators.tools.arxiv_paper_searcher import _Searcher_Tool
+from sage.middleware.operators.tools.arxiv_searcher import ArxivSearcher
+from sage.middleware.operators.tools.nature_news_fetcher import Nature_News_Fetcher_Tool
+from sage.middleware.operators.tools.searcher_tool import BochaSearchTool
+from sage.middleware.operators.tools.url_text_extractor import URL_Text_Extractor_Tool
+
+# Heavy tools (require torch/transformers) - lazy load
+_HEAVY_TOOLS_LOADED = False
+ImageCaptioner = None  # type: ignore
+text_detector = None  # type: ignore
+
+
+def _load_heavy_tools():
+    """Load tools that require torch/transformers."""
+    global _HEAVY_TOOLS_LOADED, ImageCaptioner, text_detector
+    if _HEAVY_TOOLS_LOADED:
+        return
+    try:
+        from sage.middleware.operators.tools.image_captioner import ImageCaptioner as _IC
+        from sage.middleware.operators.tools.text_detector import text_detector as _TD
+
+        ImageCaptioner = _IC
+        text_detector = _TD
+        _HEAVY_TOOLS_LOADED = True
+    except ImportError as e:
+        warnings.warn(
+            f"Heavy tool operators not available: {e}\n"
+            "Install with: pip install torch transformers",
+            UserWarning,
+            stacklevel=2,
+        )
+
+
+def __getattr__(name: str):
+    """Lazy load heavy tools on access."""
+    if name in ("ImageCaptioner", "text_detector"):
+        _load_heavy_tools()
+        if name == "ImageCaptioner":
+            return ImageCaptioner
+        if name == "text_detector":
+            return text_detector
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "BochaSearchTool",
+    "_Searcher_Tool",
+    "ArxivSearcher",
+    "Nature_News_Fetcher_Tool",
+    "ImageCaptioner",
+    "text_detector",
+    "URL_Text_Extractor_Tool",
+]

sage/middleware/operators/tools/arxiv_paper_searcher.py

@@ -0,0 +1,175 @@
+import logging
+import re
+
+import requests
+from bs4 import BeautifulSoup
+from bs4.element import Tag
+
+from sage.libs.foundation.tools.tool import BaseTool
+
+
+class _Searcher_Tool(BaseTool):
+    def __init__(self):
+        super().__init__(
+            tool_name="_Searcher_Tool",
+            tool_description="A tool that searches arXiv for papers based on a given query.",
+            input_types={
+                "query": "str - The search query for arXiv papers.",
+                "size": "int - The number of results per page (25, 50, 100, or 200). If None, use 25.",
+                "max_results": "int - The maximum number of papers to return (default: 25). Should be less than or equal to 100.",
+            },
+            output_type="list - A list of dictionaries containing paper information.",
+            demo_commands=[
+                {
+                    "command": 'execution = tool.execute(query="tool agents with large language models")',
+                    "description": "Search for papers about tool agents with large language models.",
+                },
+                {
+                    "command": 'execution = tool.execute(query="quantum computing", size=100, max_results=50)',
+                    "description": "Search for quantum computing papers, with 100 results per page, returning a maximum of 50 papers.",
+                },
+                {
+                    "command": 'execution = tool.execute(query="machine learning", max_results=75)',
+                    "description": "Search for machine learning papers, returning a maximum of 75 papers.",
+                },
+            ],
+        )
+        # Store additional metadata as instance variables
+        self.tool_version = "1.0.0"
+        self.valid_sizes = [25, 50, 100, 200]
+        self.base_url = "https://arxiv.org/search/"
+
+    def build_tool(self):
+        """
+        No specific build required for this tool.
+        """
+        pass
+
+    def execute(self, query, size=None, max_results=25):
+        """
+        Executes the arXiv search tool to find papers based on the given query.
+
+        Parameters:
+            query (str): The search query for arXiv papers.
+            size (int): The number of results per page.
+            max_results (int): The maximum number of papers to return.
+
+        Returns:
+            list: A list of dictionaries containing paper information.
+        """
+        valid_sizes = self.valid_sizes
+        base_url = self.base_url
+
+        if size is None:
+            size = 25
+        elif size not in valid_sizes:
+            size = min(valid_sizes, key=lambda x: abs(x - size))
+
+        results = []
+        start = 0
+
+        max_results = min(max_results, 100)  # NOTE: For traffic reasons, limit to 100 results
+
+        while len(results) < max_results:
+            params = {
+                "searchtype": "all",
+                "query": query,
+                "abstracts": "show",
+                "order": "",
+                "size": str(size),
+                "start": str(start),
+            }
+
+            try:
+                response = requests.get(base_url, params=params)
+                soup = BeautifulSoup(response.content, "html.parser")
+
+                papers = soup.find_all("li", class_="arxiv-result")  # type: ignore
+                if not papers:
+                    break
+
+                for paper in papers:
+                    if len(results) >= max_results:
+                        break
+
+                    title_elem = paper.find("p", class_="title")  # type: ignore
+                    title = title_elem.text.strip() if title_elem else "No title found"
+
+                    authors_elem = paper.find("p", class_="authors")  # type: ignore
+                    authors = authors_elem.text.strip() if authors_elem else "No authors found"
+                    authors = re.sub(r"^Authors:\s*", "", authors)
+                    authors = re.sub(r"\s+", " ", authors).strip()
+
+                    abstract_elem = paper.find("span", class_="abstract-full")  # type: ignore
+                    abstract = (
+                        abstract_elem.text.strip() if abstract_elem else "No abstract available"
+                    )
+                    abstract = abstract.replace("△ Less", "").strip()
+
+                    link_elem = paper.find("p", class_="list-title")  # type: ignore
+                    link_tag = link_elem.find("a") if isinstance(link_elem, Tag) else None  # type: ignore
+                    link = (
+                        link_tag["href"]
+                        if isinstance(link_tag, Tag) and link_tag.has_attr("href")
+                        else "No link found"
+                    )
+
+                    results.append(
+                        {
+                            "title": title,
+                            "authors": authors,
+                            "abstract": abstract,
+                            "link": link,
+                        }
+                    )
+
+                start += size
+
+            except Exception as e:
+                logging.error(f"Error searching arXiv: {e}")
+                break
+
+        return results[:max_results]
+
+    def get_metadata(self):
+        """
+        Returns the metadata for the _Searcher_Tool.
+
+        Returns:
+            dict: A dictionary containing the tool's metadata.
+        """
+        metadata = super().get_metadata()
+        return metadata
+
+
+if __name__ == "__main__":
+    import json
+
+    print("ArXiv Search Tool Test")
+
+    # Example usage of the _Searcher_Tool
+    tool = _Searcher_Tool()
+
+    # Get tool metadata
+    metadata = tool.get_metadata()
+    print("Tool Metadata:")
+    print(metadata)
+
+    # Sample query for searching arXiv
+    query = ""
+    # Execute the tool
+    try:
+        execution = tool.execute(query=query, size=50, max_results=10)
+        print("\n==>> Execution:")
+        print(json.dumps(execution, indent=4))  # Pretty print JSON
+        print("\n==>> Search Results:")
+        for i, paper in enumerate(execution, 1):
+            print(f"{i}. {paper['title']}")
+            print(f"   Authors: {paper['authors']}")
+            print(f"   Abstract: {paper['abstract'][:2000]}")
+            print(f"   Link: {paper['link']}")
+            print()
+    except Exception as e:
+        print(f"Execution failed: {e}")
+
+    print("Done!")

sage/middleware/operators/tools/arxiv_searcher.py

@@ -0,0 +1,102 @@
+"""
+Arxiv 论文搜索工具 (Real Implementation)
+"""
+
+import asyncio
+import logging
+import urllib.parse
+from typing import Any
+
+import aiohttp
+import feedparser
+
+from sage.libs.foundation.tools.tool import BaseTool
+
+logger = logging.getLogger(__name__)
+
+
+class ArxivSearcher(BaseTool):
+    """Arxiv 学术论文搜索工具"""
+
+    def __init__(self):
+        super().__init__(
+            tool_name="arxiv_searcher",
+            tool_description="Search Arxiv for academic papers. Returns title, authors, summary, and link.",
+            input_types=["str"],
+            output_type="list",
+            demo_commands=["search for transformer papers", "find papers about LLM agents"],
+            require_llm_engine=False,
+        )
+        self.base_url = "http://export.arxiv.org/api/query"
+
+    async def execute(self, query: str, max_results: int = 5) -> list[dict[str, Any]]:
+        """
+        Execute Arxiv search.
+        """
+        logger.info(f"Searching Arxiv for: {query}")
+
+        # Construct API query
+        # search_query=all:electron&start=0&max_results=10
+        params = {
+            "search_query": f"all:{query}",
+            "start": 0,
+            "max_results": max_results,
+            "sortBy": "relevance",
+            "sortOrder": "descending",
+        }
+
+        url = f"{self.base_url}?{urllib.parse.urlencode(params)}"
+
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.get(url) as response:
+                    if response.status != 200:
+                        logger.error(f"Arxiv API failed with status {response.status}")
+                        return []
+
+                    content = await response.text()
+
+            # Parse with feedparser
+            feed = feedparser.parse(content)
+
+            results = []
+            for entry in feed.entries:
+                paper = {
+                    "title": entry.title.replace("\n", " ").strip(),
+                    "authors": [author.name for author in entry.authors],
+                    "summary": entry.summary.replace("\n", " ").strip(),
+                    "published": entry.published,
+                    "link": entry.link,
+                    "pdf_link": next(
+                        (link.href for link in entry.links if link.title == "pdf"), None
+                    ),
+                }
+                results.append(paper)
+
+            logger.info(f"Found {len(results)} papers")
+            return results
+
+        except Exception as e:
+            logger.error(f"Arxiv search failed: {e}")
+            return []
+
+    def call(self, arguments: dict) -> Any:
+        """Sync wrapper for MCP"""
+        query = arguments.get("query")
+        if not query:
+            return []
+
+        # Check for running loop
+        try:
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                # If we are in a loop, we can't use asyncio.run.
+                # But AgentRuntime calls tools synchronously?
+                # If AgentRuntime is running in a thread, we can use asyncio.run.
+                # If AgentRuntime is running in the main loop, we are in trouble.
+                # But Gateway runs AgentRuntime in run_in_executor.
+                return asyncio.run(self.execute(query))
+        except RuntimeError:
+            return asyncio.run(self.execute(query))
+
+        return asyncio.run(self.execute(query))

sage/middleware/operators/tools/duckduckgo_searcher.py

@@ -0,0 +1,105 @@
+"""
+DuckDuckGo web search tool (no API key required).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any
+
+import aiohttp
+from bs4 import BeautifulSoup
+from pydantic import BaseModel, Field
+
+from sage.libs.foundation.tools.tool import BaseTool
+
+logger = logging.getLogger(__name__)
+
+
+class DuckDuckGoSearchInput(BaseModel):
+    query: str = Field(..., description="Search query text")
+    max_results: int = Field(5, description="Number of results to return", ge=1, le=20)
+
+
+class DuckDuckGoSearcher(BaseTool):
+    """Simple HTML-based DuckDuckGo searcher.
+
+    Uses the public HTML endpoint (no API key) and extracts title/link/snippet.
+    Intended for lightweight research fallback when no commercial search API is configured.
+    """
+
+    def __init__(self):
+        super().__init__(
+            tool_name="duckduckgo_search",
+            tool_description="Search the web via DuckDuckGo (HTML endpoint). Returns title, link, and snippet.",
+            input_types={"query": "str - search query", "max_results": "int - number of results"},
+            output_type="list",
+            demo_commands=[
+                "search for latest vector database papers",
+                "find recent ML system posts",
+            ],
+            require_llm_engine=False,
+        )
+
+    async def execute(self, query: str, max_results: int = 5) -> list[dict[str, Any]]:
+        url = "https://duckduckgo.com/html"
+        params = {"q": query, "kl": "us-en"}
+
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url, data=params, timeout=15) as resp:
+                    if resp.status != 200:
+                        logger.warning("DuckDuckGo returned status %s", resp.status)
+                        return []
+                    html = await resp.text()
+        except Exception as exc:  # noqa: BLE001
+            logger.error("DuckDuckGo search failed: %s", exc)
+            return []
+
+        soup = BeautifulSoup(html, "html.parser")
+        results: list[dict[str, Any]] = []
+
+        for result in soup.select("div.result"):
+            if len(results) >= max_results:
+                break
+
+            link_tag = result.select_one("a.result__a")
+            snippet_tag = result.select_one("a.result__snippet") or result.select_one(
+                "div.result__snippet"
+            )
+
+            title = link_tag.get_text(strip=True) if link_tag else ""
+            href = link_tag.get("href") if link_tag else ""
+            snippet = snippet_tag.get_text(strip=True) if snippet_tag else ""
+
+            if not href:
+                continue
+
+            results.append(
+                {
+                    "title": title,
+                    "link": href,
+                    "content": snippet,
+                    "source": "duckduckgo",
+                }
+            )
+
+        return results
+
+    def call(self, arguments: dict) -> Any:
+        """Sync wrapper used by MCP/AgentRuntime."""
+        query = arguments.get("query")
+        if not query:
+            return []
+
+        max_results = arguments.get("max_results", 5)
+
+        try:
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                return asyncio.run(self.execute(query, max_results=max_results))
+        except RuntimeError:
+            return asyncio.run(self.execute(query, max_results=max_results))
+
+        return asyncio.run(self.execute(query, max_results=max_results))