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

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/scripts/filter_tool_files.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+"""
+Script to filter tool files by removing tools that don't exist in the current tool universe.
+
+This script:
+1. Gets all valid tool names from ToolUniverse using scan_all=True
+2. Filters tool_relationship_graph_FINAL.json to keep only valid tools
+3. Filters v4_all_tools_final.json to keep only valid tools
+4. Preserves all other data structure and content
+"""
+
+import json
+from pathlib import Path
+
+# Import after modifying sys.path
+from tooluniverse import ToolUniverse
+
+
+def load_json_file(file_path):
+    """Load JSON file and return the data."""
+    try:
+        with open(file_path, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception as e:
+        print(f"Error loading {file_path}: {e}")
+        return None
+
+
+def save_json_file(file_path, data):
+    """Save data to JSON file."""
+    try:
+        with open(file_path, "w", encoding="utf-8") as f:
+            json.dump(data, f, ensure_ascii=False, indent=2)
+        print(f"Successfully saved filtered data to {file_path}")
+        return True
+    except Exception as e:
+        print(f"Error saving {file_path}: {e}")
+        return False
+
+
+def filter_tool_relationship_graph(data, valid_tool_names):
+    """
+    Filter tool_relationship_graph_FINAL.json to keep only valid tools.
+
+    Args:
+        data: The loaded JSON data
+        valid_tool_names: Set of valid tool names
+
+    Returns:
+        Filtered data
+    """
+    if not isinstance(data, dict):
+        print("Warning: tool_relationship_graph data is not a dict")
+        return data
+
+    filtered_data = {}
+
+    # Handle nodes array
+    if "nodes" in data and isinstance(data["nodes"], list):
+        filtered_nodes = []
+        for node in data["nodes"]:
+            if isinstance(node, dict) and "name" in node:
+                if node["name"] in valid_tool_names:
+                    filtered_nodes.append(node)
+                else:
+                    print(f"Removing node from relationship graph: {node['name']}")
+            else:
+                # Keep non-tool nodes (if any)
+                filtered_nodes.append(node)
+        filtered_data["nodes"] = filtered_nodes
+        print(
+            f"Nodes: {len(data['nodes'])} -> {len(filtered_nodes)} ({len(data['nodes']) - len(filtered_nodes)} removed)"
+        )
+
+    # Handle edges array
+    if "edges" in data and isinstance(data["edges"], list):
+        filtered_edges = []
+        for edge in data["edges"]:
+            if isinstance(edge, dict) and "source" in edge and "target" in edge:
+                # Keep edge if both source and target are valid tools
+                if (
+                    edge["source"] in valid_tool_names
+                    and edge["target"] in valid_tool_names
+                ):
+                    filtered_edges.append(edge)
+                else:
+                    print(
+                        f"Removing edge from relationship graph: {edge.get('source', 'unknown')} -> {edge.get('target', 'unknown')}"
+                    )
+            else:
+                # Keep non-tool edges (if any)
+                filtered_edges.append(edge)
+        filtered_data["edges"] = filtered_edges
+        print(
+            f"Edges: {len(data['edges'])} -> {len(filtered_edges)} ({len(data['edges']) - len(filtered_edges)} removed)"
+        )
+
+    # Keep other fields as-is (like stats, metadata, etc.)
+    for key, value in data.items():
+        if key not in ["nodes", "edges"]:
+            filtered_data[key] = value
+
+    return filtered_data
+
+
+def filter_v4_all_tools(data, valid_tool_names):
+    """
+    Filter v4_all_tools_final.json to keep only valid tools.
+
+    Args:
+        data: The loaded JSON data
+        valid_tool_names: Set of valid tool names
+
+    Returns:
+        Filtered data
+    """
+    if not isinstance(data, list):
+        print("Warning: v4_all_tools data is not a list")
+        return data
+
+    filtered_data = []
+
+    for tool in data:
+        if isinstance(tool, dict) and "name" in tool:
+            if tool["name"] in valid_tool_names:
+                filtered_data.append(tool)
+            else:
+                print(f"Removing tool from v4_all_tools: {tool['name']}")
+        else:
+            # Keep non-tool entries (if any)
+            filtered_data.append(tool)
+
+    return filtered_data
+
+
+def main():
+    """Main function to filter the tool files."""
+    print("Starting tool file filtering process...")
+
+    # Initialize ToolUniverse and get all valid tool names
+    print("Getting all valid tool names from ToolUniverse...")
+    tu = ToolUniverse()
+    all_tool_names = tu.list_built_in_tools(mode="list_name", scan_all=True)
+    valid_tool_names = set(all_tool_names)
+    print(f"Found {len(valid_tool_names)} valid tools")
+
+    # Define file paths
+    project_root = Path(__file__).parent.parent.parent.parent
+    web_dir = project_root / "web"
+
+    relationship_graph_file = web_dir / "tool_relationship_graph_FINAL.json"
+    v4_tools_file = web_dir / "v4_all_tools_final.json"
+
+    # Check if files exist
+    if not relationship_graph_file.exists():
+        print(f"Error: {relationship_graph_file} not found")
+        return
+
+    if not v4_tools_file.exists():
+        print(f"Error: {v4_tools_file} not found")
+        return
+
+    # Process tool_relationship_graph_FINAL.json
+    print(f"\nProcessing {relationship_graph_file.name}...")
+    relationship_data = load_json_file(relationship_graph_file)
+    if relationship_data is not None:
+        len(relationship_data.get("nodes", []))
+        len(relationship_data.get("edges", []))
+        filtered_relationship_data = filter_tool_relationship_graph(
+            relationship_data, valid_tool_names
+        )
+
+        # Save filtered data
+        save_json_file(relationship_graph_file, filtered_relationship_data)
+
+    # Process v4_all_tools_final.json
+    print(f"\nProcessing {v4_tools_file.name}...")
+    v4_data = load_json_file(v4_tools_file)
+    if v4_data is not None:
+        original_count = len(v4_data)
+        filtered_v4_data = filter_v4_all_tools(v4_data, valid_tool_names)
+        filtered_count = len(filtered_v4_data)
+        print(
+            f"V4 tools: {original_count} -> {filtered_count} tools ({original_count - filtered_count} removed)"
+        )
+
+        # Save filtered data
+        save_json_file(v4_tools_file, filtered_v4_data)
+
+    print("\nTool file filtering completed!")
+
+
+if __name__ == "__main__":
+    main()