isage-middleware 0.2.4.3__cp311-none-any.whl

sage/middleware/operators/agentic/config.py

@@ -0,0 +1,254 @@
+"""Agent Runtime Operator Configuration.
+
+Provides dataclass-based configuration for AgentRuntimeOperator.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+@dataclass
+class GeneratorConfig:
+    """Generator configuration for agent LLM calls.
+
+    Attributes:
+        engine_type: Engine type to use:
+            - "sagellm" (default): SageLLMGenerator with configurable backend
+            - "openai": OpenAIGenerator for OpenAI-compatible APIs
+            - "hf": HFGenerator for HuggingFace models
+        backend_type: Backend type for sagellm engine:
+            - "auto" (default): Automatically select best available backend
+            - "mock": Mock backend for testing without GPU
+            - "cuda": NVIDIA CUDA backend
+            - "ascend": Huawei Ascend NPU backend
+        model_path: Model path or HuggingFace model ID (sagellm only)
+        device_map: Device mapping strategy (auto/cuda:0/cpu)
+        dtype: Data type (auto/float16/bfloat16)
+        max_tokens: Maximum generation tokens
+        temperature: Sampling temperature
+        top_p: Nucleus sampling parameter
+        top_k: Top-k sampling parameter
+        timeout: Request timeout in seconds
+        default_options: Default generation options
+        model_name: Model name for OpenAI (openai only)
+        base_url: API base URL (openai only)
+        api_key: API key (openai only)
+    """
+
+    engine_type: Literal["sagellm", "openai", "hf"] = "sagellm"
+
+    # SageLLM options
+    backend_type: str = "auto"
+    model_path: str = ""
+    device_map: str = "auto"
+    dtype: str = "auto"
+    max_tokens: int = 2048
+    temperature: float = 0.7
+    top_p: float = 0.95
+    top_k: int = 50
+    timeout: float = 120.0
+    default_options: dict[str, Any] = field(default_factory=dict)
+
+    # OpenAI options
+    model_name: str = ""
+    base_url: str = ""
+    api_key: str = ""
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for operator initialization."""
+        return {
+            "engine_type": self.engine_type,
+            "backend_type": self.backend_type,
+            "model_path": self.model_path,
+            "device_map": self.device_map,
+            "dtype": self.dtype,
+            "max_tokens": self.max_tokens,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+            "timeout": self.timeout,
+            "default_options": self.default_options,
+            "model_name": self.model_name,
+            "base_url": self.base_url,
+            "api_key": self.api_key,
+        }
+
+
+@dataclass
+class ProfileConfig:
+    """Agent profile configuration.
+
+    Attributes:
+        name: Agent name
+        description: Agent description
+        role: Agent role (assistant/user/system)
+        system_prompt: System prompt for the agent
+    """
+
+    name: str = "DefaultAgent"
+    description: str = "A general-purpose AI assistant"
+    role: str = "assistant"
+    system_prompt: str = "You are a helpful assistant."
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for operator initialization."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "role": self.role,
+            "system_prompt": self.system_prompt,
+        }
+
+
+@dataclass
+class RuntimeSettings:
+    """Runtime settings for agent execution.
+
+    Attributes:
+        max_steps: Maximum execution steps
+        summarizer: Summarizer config (null/"reuse_generator"/dict)
+    """
+
+    max_steps: int = 6
+    summarizer: str | dict[str, Any] | None = "reuse_generator"
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for operator initialization."""
+        return {
+            "max_steps": self.max_steps,
+            "summarizer": self.summarizer,
+        }
+
+
+@dataclass
+class AgentRuntimeConfig:
+    """Complete configuration for AgentRuntimeOperator.
+
+    Example:
+        ```python
+        # Create config with mock backend for testing
+        config = AgentRuntimeConfig(
+            generator=GeneratorConfig(
+                engine_type="sagellm",
+                backend_type="mock",
+            ),
+            profile=ProfileConfig(name="TestBot"),
+        )
+        operator = AgentRuntimeOperator(config=config.to_dict())
+
+        # Create config with OpenAI
+        config = AgentRuntimeConfig(
+            generator=GeneratorConfig(
+                engine_type="openai",
+                model_name="gpt-4o-mini",
+                api_key="sk-xxx",  # pragma: allowlist secret
+            ),
+        )
+        ```
+
+    Attributes:
+        generator: Generator configuration
+        profile: Agent profile configuration
+        planner: Planner configuration (optional)
+        tools: List of tool specifications
+        runtime: Runtime settings
+    """
+
+    generator: GeneratorConfig = field(default_factory=GeneratorConfig)
+    profile: ProfileConfig = field(default_factory=ProfileConfig)
+    planner: dict[str, Any] = field(default_factory=dict)
+    tools: list[dict[str, Any]] = field(default_factory=list)
+    runtime: RuntimeSettings = field(default_factory=RuntimeSettings)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for operator initialization."""
+        return {
+            "generator": self.generator.to_dict(),
+            "profile": self.profile.to_dict(),
+            "planner": self.planner,
+            "tools": self.tools,
+            "runtime": self.runtime.to_dict(),
+        }
+
+    @classmethod
+    def for_mock_testing(cls, profile_name: str = "TestBot") -> AgentRuntimeConfig:
+        """Create a configuration for mock testing.
+
+        Args:
+            profile_name: Name for the test agent profile
+
+        Returns:
+            AgentRuntimeConfig configured for mock backend
+        """
+        return cls(
+            generator=GeneratorConfig(
+                engine_type="sagellm",
+                backend_type="mock",
+            ),
+            profile=ProfileConfig(name=profile_name),
+        )
+
+    @classmethod
+    def for_openai(
+        cls,
+        model_name: str = "gpt-4o-mini",
+        api_key: str = "",
+        base_url: str = "https://api.openai.com/v1",
+        profile_name: str = "OpenAIAgent",
+    ) -> AgentRuntimeConfig:
+        """Create a configuration for OpenAI.
+
+        Args:
+            model_name: OpenAI model name
+            api_key: OpenAI API key
+            base_url: API base URL
+            profile_name: Name for the agent profile
+
+        Returns:
+            AgentRuntimeConfig configured for OpenAI
+        """
+        return cls(
+            generator=GeneratorConfig(
+                engine_type="openai",
+                model_name=model_name,
+                api_key=api_key,
+                base_url=base_url,
+            ),
+            profile=ProfileConfig(name=profile_name),
+        )
+
+    @classmethod
+    def for_sagellm(
+        cls,
+        model_path: str,
+        backend_type: str = "auto",
+        profile_name: str = "SageLLMAgent",
+    ) -> AgentRuntimeConfig:
+        """Create a configuration for SageLLM.
+
+        Args:
+            model_path: Model path or HuggingFace model ID
+            backend_type: Backend type (auto/mock/cuda/ascend)
+            profile_name: Name for the agent profile
+
+        Returns:
+            AgentRuntimeConfig configured for SageLLM
+        """
+        return cls(
+            generator=GeneratorConfig(
+                engine_type="sagellm",
+                backend_type=backend_type,
+                model_path=model_path,
+            ),
+            profile=ProfileConfig(name=profile_name),
+        )
+
+
+__all__ = [
+    "AgentRuntimeConfig",
+    "GeneratorConfig",
+    "ProfileConfig",
+    "RuntimeSettings",
+]

sage/middleware/operators/agentic/planning_operator.py

@@ -0,0 +1,125 @@
+"""
+Planning Operator
+
+Middleware operator for planning using runtime components.
+
+Supports engine_type switching:
+- sagellm (default): SageLLMGenerator with configurable backend
+  - backend_type="auto": Automatically select best available backend
+  - backend_type="mock": Mock backend for testing without GPU
+  - backend_type="cuda": NVIDIA CUDA backend
+  - backend_type="ascend": Huawei Ascend NPU backend
+- openai: OpenAIGenerator for OpenAI-compatible APIs
+- hf: HFGenerator for HuggingFace models
+"""
+
+from typing import Any, Optional
+
+from sage_libs.sage_agentic.agents.runtime import BenchmarkAdapter, Orchestrator, RuntimeConfig
+from sage_libs.sage_agentic.agents.runtime.config import PlannerConfig
+
+from sage.common.core.functions import MapFunction
+
+from .runtime import _build_generator
+
+
+class PlanningOperator(MapFunction):
+    """
+    Operator for planning.
+
+    Wraps runtime planner in a middleware operator interface.
+
+    Args:
+        planner: Planner instance (optional)
+        config: Configuration dictionary with optional keys:
+            - planner: Planner-specific config
+            - generator: Generator config with engine_type/backend_type
+            - engine_type: Shorthand for generator.engine_type (sagellm/openai/hf)
+            - backend_type: Shorthand for generator.backend_type (auto/mock/cuda/ascend)
+
+    Example:
+        ```python
+        # Using sagellm with mock backend (for testing)
+        operator = PlanningOperator(config={
+            "generator": {
+                "engine_type": "sagellm",
+                "backend_type": "mock",
+            },
+        })
+
+        # Using default sagellm with auto backend
+        operator = PlanningOperator(config={
+            "engine_type": "sagellm",
+            "backend_type": "auto",
+        })
+        ```
+    """
+
+    def __init__(
+        self,
+        planner: Optional[Any] = None,
+        config: Optional[dict[str, Any]] = None,
+    ):
+        """Initialize planning operator.
+
+        Args:
+            planner: Planner instance (optional)
+            config: Configuration dictionary
+        """
+        super().__init__()
+
+        # Parse configuration
+        if config is None:
+            config = {}
+
+        self.config = config
+
+        # Build generator if config provided
+        generator_conf = config.get("generator", {})
+        # Allow shorthand engine_type/backend_type at top level
+        if "engine_type" in config and "engine_type" not in generator_conf:
+            generator_conf["engine_type"] = config["engine_type"]
+        if "backend_type" in config and "backend_type" not in generator_conf:
+            generator_conf["backend_type"] = config["backend_type"]
+
+        # Build generator (defaults to sagellm with auto backend)
+        if generator_conf or not planner:
+            engine_type = generator_conf.get("engine_type", "sagellm")
+            # Ensure we have at least minimal config
+            if not generator_conf:
+                generator_conf = {"engine_type": "sagellm", "backend_type": "auto"}
+            self.generator = _build_generator(generator_conf, engine_type=engine_type)
+        else:
+            self.generator = None
+
+        planner_config = PlannerConfig(**config.get("planner", {}))
+        runtime_config = RuntimeConfig(planner=planner_config)
+
+        # Create orchestrator
+        self.orchestrator = Orchestrator(config=runtime_config, planner=planner)
+
+        # Create adapter for easy use
+        self.adapter = BenchmarkAdapter(self.orchestrator)
+
+    def __call__(self, request: Any) -> Any:
+        """Execute planning.
+
+        Args:
+            request: Planning request
+
+        Returns:
+            Generated plan
+        """
+        return self.adapter.run_planning(request)
+
+    def execute(self, data: Any) -> Any:
+        """Execute map function interface."""
+        return self.__call__(data)
+
+    def get_metrics(self) -> dict[str, Any]:
+        """Get performance metrics.
+
+        Returns:
+            Dictionary of metrics
+        """
+        return self.adapter.get_metrics()

sage/middleware/operators/agentic/refined_searcher.py

@@ -0,0 +1,132 @@
+"""
+RefinedSearcherOperator - Search with optional context compression.
+
+Uses isage-refiner for context compression if enabled.
+
+Installation:
+    pip install isage-refiner  # Optional, only if refiner is used
+"""
+
+import logging
+from typing import Any, AsyncGenerator, Optional
+
+from sage_libs.sage_agentic.agents.bots.searcher_bot import SearcherBot
+
+from sage.libs.foundation.tools.tool import BaseTool
+
+logger = logging.getLogger(__name__)
+
+
+class RefinedSearcherOperator:
+    """
+    L4 Operator that wraps L3 SearcherBot and adds optional refiner capabilities.
+
+    Uses isage-refiner for context compression when refiner_config is provided.
+    """
+
+    name = "search_internet"
+    description = "Search the internet for information using multiple sources (Arxiv, etc)."
+    input_schema = {
+        "type": "object",
+        "properties": {"query": {"type": "string", "description": "The search query"}},
+        "required": ["query"],
+    }
+
+    input_types = {"query": "str - The search query"}
+
+    def __init__(
+        self, tools: list[BaseTool], refiner_config: Optional[dict[str, Any]] = None, **kwargs
+    ):
+        self.bot = SearcherBot(tools=tools, **kwargs)
+
+        self.compressor = None
+        if refiner_config:
+            try:
+                algorithm = refiner_config.get("algorithm", "long_refiner").lower()
+                self._init_compressor(algorithm, refiner_config)
+                logger.info(f"RefinedSearcherOperator: Initialized {algorithm} compressor")
+            except ImportError as e:
+                logger.warning(
+                    f"RefinedSearcherOperator: isage-refiner not installed: {e}\n"
+                    f"Install with: pip install isage-refiner"
+                )
+            except Exception as e:
+                logger.warning(f"RefinedSearcherOperator: Failed to init compressor: {e}")
+
+    def _init_compressor(self, algorithm: str, config: dict[str, Any]):
+        """Initialize compressor from isage-refiner."""
+        if algorithm == "long_refiner":
+            from sage_refiner import LongRefinerCompressor
+
+            self.compressor = LongRefinerCompressor(
+                base_model_path=config.get("base_model_path", "Qwen/Qwen2.5-3B-Instruct"),
+                score_model_path=config.get("score_model_path", "BAAI/bge-reranker-v2-m3"),
+                max_model_len=config.get("max_model_len", 25000),
+                gpu_memory_utilization=config.get("gpu_memory_utilization", 0.5),
+            )
+        elif algorithm == "reform":
+            from sage_refiner import REFORMCompressor
+
+            self.compressor = REFORMCompressor(**config.get("reform_config", {}))
+        elif algorithm == "provence":
+            from sage_refiner import ProvenceCompressor
+
+            self.compressor = ProvenceCompressor(**config.get("provence_config", {}))
+        else:
+            raise ValueError(f"Unsupported algorithm: {algorithm}")
+
+        self.budget = config.get("budget", 2048)
+
+    def call(self, arguments: dict) -> Any:
+        """MCP compatible call method"""
+        import asyncio
+
+        query = arguments.get("query")
+        try:
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                return asyncio.run(self.execute(query))
+        except RuntimeError:
+            return asyncio.run(self.execute(query))
+
+        return asyncio.run(self.execute(query))
+
+    async def execute(self, query: str) -> dict[str, Any]:
+        """Execute search and optionally compress results."""
+        data = query
+        # 1. Execute L3 Bot
+        raw_result = await self.bot.execute(data)
+        results = raw_result.get("results", [])
+
+        # 2. Compress if enabled
+        if self.compressor and results:
+            query_str = data if isinstance(data, str) else data.get("query", "")
+            try:
+                logger.info(f"Compressing {len(results)} results for query: {query_str}")
+
+                # Normalize documents to isage-refiner format
+                documents = [
+                    {"contents": r.get("contents") or r.get("text") or str(r)} for r in results
+                ]
+
+                compress_result = self.compressor.compress(
+                    question=query_str,
+                    document_list=documents,
+                    budget=self.budget,
+                )
+
+                return {
+                    "results": compress_result.get("compressed_context", ""),
+                    "original_count": len(results),
+                    "compressed": True,
+                }
+            except Exception as e:
+                logger.error(f"Compression failed: {e}")
+                return raw_result
+
+        return raw_result
+
+    async def execute_stream(self, data: Any) -> AsyncGenerator[dict[str, Any], None]:
+        """Stream execution. Compression is batch, so just stream search events."""
+        async for event in self.bot.execute_stream(data):
+            yield event