tooluniverse 1.0.3__py3-none-any.whl → 1.0.4__py3-none-any.whl

tooluniverse/compose_scripts/output_summarizer.py

@@ -81,31 +81,37 @@ def compose(arguments: Dict[str, Any], tooluniverse, call_tool) -> Dict[str, Any
             else:
                 print(f"❌ Chunk {i+1} summarization failed")
 
-        # Step 3: Merge summaries
+        # Step 3: Merge summaries (or gracefully fall back)
         if chunk_summaries:
             final_summary = _merge_summaries(
                 chunk_summaries, query_context, tool_name, max_summary_length, call_tool
             )
+            print(
+                f"✅ Summarization completed. Final length: {len(final_summary)} characters"
+            )
+            return {
+                "success": True,
+                "original_length": len(tool_output),
+                "summary_length": len(final_summary),
+                "chunks_processed": len(chunks),
+                "summary": final_summary,
+                "tool_name": tool_name,
+            }
         else:
-            final_summary = "Unable to generate summary due to processing errors."
+            # Treat as a non-fatal failure so upstream falls back to original output
             print("❌ No chunk summaries were generated. This usually indicates:")
             print("   1. ToolOutputSummarizer tool is not available")
             print("   2. The output_summarization tools are not loaded")
             print("   3. There was an error in the summarization process")
             print("   Please check that the SMCP server is started with hooks enabled.")
-
-        print(
-            f"✅ Summarization completed. Final length: {len(final_summary)} characters"
-        )
-
-        return {
-            "success": True,
-            "original_length": len(tool_output),
-            "summary_length": len(final_summary),
-            "chunks_processed": len(chunks),
-            "summary": final_summary,
-            "tool_name": tool_name,
-        }
+            return {
+                "success": False,
+                "error": "No chunk summaries generated",
+                "original_length": len(tool_output),
+                "chunks_processed": len(chunks),
+                "original_output": tool_output,
+                "tool_name": tool_name,
+            }
 
     except Exception as e:
         error_msg = f"Error in output summarization: {str(e)}"

tooluniverse/data/output_summarization_tools.json

@@ -73,7 +73,7 @@
           "type": "integer",
           "description": "Size of each chunk for processing",
           "required": false,
-          "default": 2000
+          "default": 30000
         },
         "focus_areas": {
           "type": "string",
@@ -85,7 +85,7 @@
           "type": "integer",
           "description": "Maximum length of final summary",
           "required": false,
-          "default": 3000
+          "default": 10000
         }
       },
       "required": ["tool_output", "query_context", "tool_name"]

tooluniverse/llm_clients.py

@@ -0,0 +1,369 @@
+from __future__ import annotations
+from typing import Any, Dict, List, Optional
+import os
+import time
+import json as _json
+
+
+class BaseLLMClient:
+    def test_api(self) -> None:
+        raise NotImplementedError
+
+    def infer(
+        self,
+        messages: List[Dict[str, str]],
+        temperature: Optional[float],
+        max_tokens: Optional[int],
+        return_json: bool,
+        custom_format: Any = None,
+        max_retries: int = 5,
+        retry_delay: int = 5,
+    ) -> Optional[str]:
+        raise NotImplementedError
+
+
+class AzureOpenAIClient(BaseLLMClient):
+    # Built-in defaults for model families (can be overridden by env)
+    DEFAULT_MODEL_LIMITS: Dict[str, Dict[str, int]] = {
+        # GPT-4.1 series
+        "gpt-4.1": {"max_output": 32768, "context_window": 1_047_576},
+        "gpt-4.1-mini": {"max_output": 32768, "context_window": 1_047_576},
+        "gpt-4.1-nano": {"max_output": 32768, "context_window": 1_047_576},
+        # GPT-4o series
+        "gpt-4o-1120": {"max_output": 16384, "context_window": 128_000},
+        "gpt-4o-0806": {"max_output": 16384, "context_window": 128_000},
+        "gpt-4o-mini-0718": {"max_output": 16384, "context_window": 128_000},
+        "gpt-4o": {"max_output": 16384, "context_window": 128_000},  # general prefix
+        # O-series
+        "o4-mini-0416": {"max_output": 100_000, "context_window": 200_000},
+        "o3-mini-0131": {"max_output": 100_000, "context_window": 200_000},
+        "o4-mini": {"max_output": 100_000, "context_window": 200_000},
+        "o3-mini": {"max_output": 100_000, "context_window": 200_000},
+        # Embeddings (for completeness)
+        "embedding-ada": {"max_output": 8192, "context_window": 8192},
+        "text-embedding-3-small": {"max_output": 8192, "context_window": 8192},
+        "text-embedding-3-large": {"max_output": 8192, "context_window": 8192},
+    }
+
+    def __init__(self, model_id: str, api_version: Optional[str], logger):
+        try:
+            from openai import AzureOpenAI as _AzureOpenAI  # type: ignore
+            import openai as _openai  # type: ignore
+        except Exception as e:  # pragma: no cover
+            raise RuntimeError("openai AzureOpenAI client is not available") from e
+        self._AzureOpenAI = _AzureOpenAI
+        self._openai = _openai
+
+        self.model_name = model_id
+        self.logger = logger
+
+        resolved_version = api_version or self._resolve_api_version(model_id)
+        self.logger.debug(
+            f"Resolved Azure API version for {model_id}: {resolved_version}"
+        )
+
+        api_key = os.getenv("AZURE_OPENAI_API_KEY")
+        if not api_key:
+            raise ValueError("AZURE_OPENAI_API_KEY not set")
+        endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "https://azure-ai.hms.edu")
+        self.client = self._AzureOpenAI(
+            azure_endpoint=endpoint, api_key=api_key, api_version=resolved_version
+        )
+        self.api_version = resolved_version
+
+        # Load env overrides for model limits (JSON dict of {prefix: {max_output, context_window}})
+        env_limits_raw = os.getenv("AZURE_DEFAULT_MODEL_LIMITS")
+        self._default_limits: Dict[str, Dict[str, int]] = (
+            self.DEFAULT_MODEL_LIMITS.copy()
+        )
+        if env_limits_raw:
+            try:
+                env_limits = _json.loads(env_limits_raw)
+                # shallow merge by keys
+                for k, v in env_limits.items():
+                    if isinstance(v, dict):
+                        base = self._default_limits.get(k, {}).copy()
+                        base.update(
+                            {
+                                kk: int(vv)
+                                for kk, vv in v.items()
+                                if isinstance(vv, (int, float, str))
+                            }
+                        )
+                        self._default_limits[k] = base
+            except Exception:
+                # ignore bad env format
+                pass
+
+    # --------- helpers (Azure specific) ---------
+    def _resolve_api_version(self, model_id: str) -> str:
+        mapping_raw = os.getenv("AZURE_OPENAI_API_VERSION_BY_MODEL")
+        mapping: Dict[str, str] = {}
+        if mapping_raw:
+            try:
+                mapping = _json.loads(mapping_raw)
+            except Exception:
+                mapping = {}
+        if model_id in mapping:
+            return mapping[model_id]
+        for k, v in mapping.items():
+            try:
+                if model_id.startswith(k):
+                    return v
+            except Exception:
+                continue
+        try:
+            if model_id.startswith("o3-mini") or model_id.startswith("o4-mini"):
+                return "2024-12-01-preview"
+        except Exception:
+            pass
+        return os.getenv("AZURE_OPENAI_API_VERSION", "2024-12-01-preview")
+
+    def _resolve_default_max_tokens(self, model_id: str) -> Optional[int]:
+        # Highest priority: explicit env per-model tokens mapping
+        mapping_raw = os.getenv("AZURE_MAX_TOKENS_BY_MODEL")
+        mapping: Dict[str, Any] = {}
+        if mapping_raw:
+            try:
+                mapping = _json.loads(mapping_raw)
+            except Exception:
+                mapping = {}
+        if model_id in mapping:
+            try:
+                return int(mapping[model_id])
+            except Exception:
+                pass
+        for k, v in mapping.items():
+            try:
+                if model_id.startswith(k):
+                    return int(v)
+            except Exception:
+                continue
+        # Next: built-in/default-limits map (with env merged)
+        if model_id in self._default_limits:
+            return int(self._default_limits[model_id].get("max_output", 0)) or None
+        for k, v in self._default_limits.items():
+            try:
+                if model_id.startswith(k):
+                    return int(v.get("max_output", 0)) or None
+            except Exception:
+                continue
+        return None
+
+    def _normalize_temperature(
+        self, model_id: str, temperature: Optional[float]
+    ) -> Optional[float]:
+        if isinstance(model_id, str) and (
+            model_id.startswith("o3-mini") or model_id.startswith("o4-mini")
+        ):
+            if temperature is not None:
+                self.logger.warning(
+                    f"Model {model_id} does not support 'temperature'; ignoring provided value."
+                )
+            return None
+        return temperature
+
+    # --------- public API ---------
+    def test_api(self) -> None:
+        test_messages = [{"role": "user", "content": "ping"}]
+        token_attempts = [1, 4, 16, 32]
+        last_error: Optional[Exception] = None
+        for tok in token_attempts:
+            try:
+                try:
+                    self.client.chat.completions.create(
+                        model=self.model_name,
+                        messages=test_messages,
+                        max_tokens=tok,
+                        temperature=0,
+                    )
+                    return
+                except self._openai.BadRequestError:  # type: ignore[attr-defined]
+                    self.client.chat.completions.create(
+                        model=self.model_name,
+                        messages=test_messages,
+                        max_completion_tokens=tok,
+                    )
+                    return
+            except Exception as e:  # noqa: BLE001
+                last_error = e
+                msg = str(e).lower()
+                if (
+                    "max_tokens" in msg
+                    or "model output limit" in msg
+                    or "finish the message" in msg
+                ) and tok != token_attempts[-1]:
+                    continue
+                break
+        if last_error:
+            raise ValueError(f"ChatGPT API test failed: {last_error}")
+        raise ValueError("ChatGPT API test failed: unknown error")
+
+    def infer(
+        self,
+        messages: List[Dict[str, str]],
+        temperature: Optional[float],
+        max_tokens: Optional[int],
+        return_json: bool,
+        custom_format: Any = None,
+        max_retries: int = 5,
+        retry_delay: int = 5,
+    ) -> Optional[str]:
+        retries = 0
+        call_fn = (
+            self.client.chat.completions.parse
+            if custom_format is not None
+            else self.client.chat.completions.create
+        )
+        response_format = (
+            custom_format
+            if custom_format is not None
+            else ({"type": "json_object"} if return_json else None)
+        )
+        eff_temp = self._normalize_temperature(self.model_name, temperature)
+        eff_max = (
+            max_tokens
+            if max_tokens is not None
+            else self._resolve_default_max_tokens(self.model_name)
+        )
+        while retries < max_retries:
+            try:
+                kwargs: Dict[str, Any] = {
+                    "model": self.model_name,
+                    "messages": messages,
+                }
+                if response_format is not None:
+                    kwargs["response_format"] = response_format
+                if eff_temp is not None:
+                    kwargs["temperature"] = eff_temp
+                try:
+                    if eff_max is not None:
+                        resp = call_fn(max_tokens=eff_max, **kwargs)
+                    else:
+                        resp = call_fn(**kwargs)
+                except self._openai.BadRequestError as be:  # type: ignore[attr-defined]
+                    if eff_max is not None:
+                        resp = call_fn(max_completion_tokens=eff_max, **kwargs)
+                    else:
+                        be_msg = str(be).lower()
+                        fallback_limits = [
+                            8192,
+                            4096,
+                            2048,
+                            1024,
+                            512,
+                            256,
+                            128,
+                            64,
+                            32,
+                        ]
+                        if any(
+                            k in be_msg
+                            for k in [
+                                "max_tokens",
+                                "output limit",
+                                "finish the message",
+                                "max_completion_tokens",
+                            ]
+                        ):
+                            last_exc: Optional[Exception] = be
+                            for lim in fallback_limits:
+                                try:
+                                    try:
+                                        resp = call_fn(
+                                            max_completion_tokens=lim, **kwargs
+                                        )
+                                        last_exc = None
+                                        break
+                                    except Exception as inner_e:  # noqa: BLE001
+                                        last_exc = inner_e
+                                        resp = call_fn(max_tokens=lim, **kwargs)
+                                        last_exc = None
+                                        break
+                                except Exception as inner2:  # noqa: BLE001
+                                    last_exc = inner2
+                                    continue
+                            if last_exc is not None:
+                                raise last_exc
+                        else:
+                            raise be
+                if custom_format is not None:
+                    return resp.choices[0].message.parsed.model_dump()
+                return resp.choices[0].message.content
+            except self._openai.RateLimitError:  # type: ignore[attr-defined]
+                self.logger.warning(
+                    f"Rate limit exceeded. Retrying in {retry_delay} seconds..."
+                )
+                retries += 1
+                time.sleep(retry_delay * retries)
+            except Exception as e:  # noqa: BLE001
+                self.logger.error(f"An error occurred: {e}")
+                import traceback
+
+                traceback.print_exc()
+                break
+        self.logger.error("Max retries exceeded. Unable to complete the request.")
+        return None
+
+
+class GeminiClient(BaseLLMClient):
+    def __init__(self, model_name: str, logger):
+        try:
+            import google.generativeai as genai  # type: ignore
+        except Exception as e:  # pragma: no cover
+            raise RuntimeError("google.generativeai not available") from e
+        api_key = os.getenv("GEMINI_API_KEY")
+        if not api_key:
+            raise ValueError("GEMINI_API_KEY not found")
+        self._genai = genai
+        self._genai.configure(api_key=api_key)
+        self.model_name = model_name
+        self.logger = logger
+
+    def _build_model(self):
+        return self._genai.GenerativeModel(self.model_name)
+
+    def test_api(self) -> None:
+        model = self._build_model()
+        model.generate_content(
+            "ping",
+            generation_config={
+                "max_output_tokens": 8,
+                "temperature": 0,
+            },
+        )
+
+    def infer(
+        self,
+        messages: List[Dict[str, str]],
+        temperature: Optional[float],
+        max_tokens: Optional[int],
+        return_json: bool,
+        custom_format: Any = None,
+        max_retries: int = 5,
+        retry_delay: int = 5,
+    ) -> Optional[str]:
+        if return_json:
+            raise ValueError("Gemini JSON mode not supported here")
+        contents = ""
+        for m in messages:
+            if m["role"] in ("user", "system"):
+                contents += f"{m['content']}\n"
+        retries = 0
+        while retries < max_retries:
+            try:
+                gen_cfg: Dict[str, Any] = {
+                    "temperature": (temperature if temperature is not None else 0)
+                }
+                if max_tokens is not None:
+                    gen_cfg["max_output_tokens"] = max_tokens
+                model = self._build_model()
+                resp = model.generate_content(contents, generation_config=gen_cfg)
+                return getattr(resp, "text", None) or getattr(resp, "candidates", [{}])[
+                    0
+                ].get("content")
+            except Exception as e:  # noqa: BLE001
+                self.logger.error(f"Gemini error: {e}")
+                retries += 1
+                time.sleep(retry_delay * retries)
+        return None

tooluniverse/output_hook.py

@@ -226,6 +226,9 @@ class SummarizationHook(OutputHook):
         self.chunk_size = hook_config.get("chunk_size", 2000)
         self.focus_areas = hook_config.get("focus_areas", "key_findings_and_results")
         self.max_summary_length = hook_config.get("max_summary_length", 3000)
+        # Optional timeout to prevent hangs in composer / LLM calls
+        # If the composer does not return within this window, we gracefully fall back
+        self.composer_timeout_sec = hook_config.get("composer_timeout_sec", 20)
 
     def process(
         self,
@@ -252,6 +255,19 @@ class SummarizationHook(OutputHook):
             Any: The summarized output, or original output if summarization fails
         """
         try:
+            # Debug: basic context
+            try:
+                _len = len(str(result))
+            except Exception:
+                _len = -1
+            import sys as _sys
+
+            print(
+                f"[SummarizationHook] process: tool={tool_name}, result_len={_len}, "
+                f"chunk_size={self.chunk_size}, max_summary_length={self.max_summary_length}",
+                file=_sys.stderr,
+                flush=True,
+            )
             # Check if the required tools are available
             if (
                 self.composer_tool_name not in self.tooluniverse.callable_functions
@@ -277,9 +293,49 @@ class SummarizationHook(OutputHook):
             }
 
             # Call Compose Summarizer Tool through ToolUniverse
-            composer_result = self.tooluniverse.run_one_function(
-                {"name": self.composer_tool_name, "arguments": composer_args}
+            print(
+                f"[SummarizationHook] calling composer tool: {self.composer_tool_name} (timeout={self.composer_timeout_sec}s)",
+                file=_sys.stderr,
+                flush=True,
             )
+            # Run composer with timeout to avoid hangs
+            try:
+                from concurrent.futures import (
+                    ThreadPoolExecutor,
+                )
+
+                def _call_composer():
+                    return self.tooluniverse.run_one_function(
+                        {"name": self.composer_tool_name, "arguments": composer_args}
+                    )
+
+                with ThreadPoolExecutor(max_workers=1) as _pool:
+                    _future = _pool.submit(_call_composer)
+                    composer_result = _future.result(timeout=self.composer_timeout_sec)
+            except Exception as _e_timeout:
+                # Timeout or execution error; log and fall back to original output
+                print(
+                    f"[SummarizationHook] composer execution failed/timeout: {_e_timeout}",
+                    file=_sys.stderr,
+                    flush=True,
+                )
+                return result
+            # Debug: show composer result meta
+            try:
+                if isinstance(composer_result, dict):
+                    success = composer_result.get("success", False)
+                    summary_len = len(composer_result.get("summary", ""))
+                    print(
+                        f"[SummarizationHook] composer_result: success={success} summary_len={summary_len}",
+                        file=_sys.stderr,
+                        flush=True,
+                    )
+            except Exception as _e_dbg:
+                print(
+                    f"[SummarizationHook] debug error inspecting composer_result: {_e_dbg}",
+                    file=_sys.stderr,
+                    flush=True,
+                )
 
             # Process Compose Tool result
             if isinstance(composer_result, dict) and composer_result.get("success"):
@@ -294,7 +350,13 @@ class SummarizationHook(OutputHook):
 
         except Exception as e:
             error_msg = str(e)
-            print(f"Error in summarization hook: {error_msg}")
+            import sys as _sys
+
+            print(
+                f"Error in summarization hook: {error_msg}",
+                file=_sys.stderr,
+                flush=True,
+            )
 
             # Check if the error is due to missing tools
             if "not found" in error_msg.lower() or "ToolOutputSummarizer" in error_msg:
@@ -365,6 +427,16 @@ class HookManager:
         self.config_path = config.get("config_path", "template/hook_config.json")
         self._pending_tools_to_load: List[str] = []
         self._load_hook_config()
+
+        # Validate LLM API keys before loading hooks
+        if not self._validate_llm_api_keys():
+            print("⚠️  Warning: LLM API keys not available. Hooks will be disabled.")
+            print(
+                "   To enable hooks, please set AZURE_OPENAI_API_KEY environment variable."
+            )
+            self.enabled = False
+            return
+
         self._load_hooks()
 
     def apply_hooks(
@@ -415,6 +487,23 @@ class HookManager:
 
         return result
 
+    def _validate_llm_api_keys(self) -> bool:
+        """
+        Validate that LLM API keys are available for hook tools.
+
+        Returns:
+            bool: True if API keys are available, False otherwise
+        """
+        from .agentic_tool import AgenticTool
+
+        if AgenticTool.has_any_api_keys():
+            print("✅ LLM API keys validated successfully")
+            return True
+        else:
+            print("❌ LLM API key validation failed: No API keys available")
+            print("   To enable hooks, please set API key environment variables.")
+            return False
+
     def enable_hook(self, hook_name: str):
         """
         Enable a specific hook by name.

tooluniverse/smcp_server.py

@@ -146,8 +146,8 @@ Examples:
   # Start with specific categories
   tooluniverse-stdio --categories uniprot ChEMBL opentarget
 
-  # Disable hooks
-  tooluniverse-stdio --no-hooks
+  # Enable hooks
+  tooluniverse-stdio --hooks
 
   # Use FileSaveHook instead of SummarizationHook
   tooluniverse-stdio --hook-type FileSaveHook
@@ -272,18 +272,17 @@ Examples:
         "--verbose", "-v", action="store_true", help="Enable verbose logging"
     )
 
-    # Hook configuration options (default enabled for stdio)
+    # Hook configuration options (default disabled for stdio)
     hook_group = parser.add_argument_group("Hook Configuration")
     hook_group.add_argument(
-        "--no-hooks",
+        "--hooks",
         action="store_true",
-        help="Disable output processing hooks (default: enabled for stdio)",
+        help="Enable output processing hooks (default: disabled for stdio)",
     )
     hook_group.add_argument(
         "--hook-type",
         choices=["SummarizationHook", "FileSaveHook"],
-        default="SummarizationHook",
-        help="Hook type to use (default: SummarizationHook)",
+        help="Hook type to use (default: SummarizationHook when hooks are enabled)",
     )
     hook_group.add_argument(
         "--hook-config-file",
@@ -493,16 +492,23 @@ Examples:
                 hook_config = json.load(f)
             print(f"🔗 Hook config loaded from: {args.hook_config_file}")
 
-        # Determine hook settings (default enabled for stdio)
-        hooks_enabled = not args.no_hooks
+        # Determine hook settings (default disabled for stdio)
+        hooks_enabled = (
+            args.hooks or args.hook_type is not None or hook_config is not None
+        )
+
+        # Set default hook type if hooks are enabled but no type specified
+        hook_type = args.hook_type
+        if hooks_enabled and hook_type is None:
+            hook_type = "SummarizationHook"
         if hooks_enabled:
-            if args.hook_type:
-                print(f"🔗 Hooks enabled: {args.hook_type}")
+            if hook_type:
+                print(f"🔗 Hooks enabled: {hook_type}")
             elif hook_config:
                 hook_count = len(hook_config.get("hooks", []))
                 print(f"🔗 Hooks enabled: {hook_count} custom hooks")
             else:
-                print(f"🔗 Hooks enabled: {args.hook_type} (default)")
+                print("🔗 Hooks enabled: default configuration")
         else:
             print("🔗 Hooks disabled")
 
@@ -525,7 +531,7 @@ Examples:
             stateless_http=True,  # Enable stateless mode for MCPAutoLoaderTool compatibility
             hooks_enabled=hooks_enabled,
             hook_config=hook_config,
-            hook_type=args.hook_type,
+            hook_type=hook_type,
         )
 
         # Run server with stdio transport (forced)