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

tooluniverse/llm_clients.py

@@ -21,6 +21,29 @@ class BaseLLMClient:
     ) -> Optional[str]:
         raise NotImplementedError
 
+    def infer_stream(
+        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,
+    ):
+        """Default streaming implementation falls back to regular inference."""
+        result = self.infer(
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            return_json=return_json,
+            custom_format=custom_format,
+            max_retries=max_retries,
+            retry_delay=retry_delay,
+        )
+        if result is not None:
+            yield result
+
 
 class AzureOpenAIClient(BaseLLMClient):
     # Built-in defaults for model families (can be overridden by env)
@@ -305,6 +328,179 @@ class AzureOpenAIClient(BaseLLMClient):
         self.logger.error("Max retries exceeded. Unable to complete the request.")
         return None
 
+    def infer_stream(
+        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,
+    ):
+        if return_json or custom_format is not None:
+            yield from super().infer_stream(
+                messages,
+                temperature,
+                max_tokens,
+                return_json,
+                custom_format,
+                max_retries,
+                retry_delay,
+            )
+            return
+
+        retries = 0
+        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,
+                    "stream": True,
+                }
+                if temperature is not None:
+                    kwargs["temperature"] = temperature
+                if eff_max is not None:
+                    kwargs["max_tokens"] = eff_max
+
+                stream = self.client.chat.completions.create(**kwargs)
+                for chunk in stream:
+                    text = AzureOpenAIClient._extract_text_from_chunk(chunk)  # type: ignore[attr-defined]
+                    if text:
+                        yield text
+                return
+            except self._openai.RateLimitError:  # type: ignore[attr-defined]
+                self.logger.warning(
+                    f"OpenRouter streaming rate limit hit. Retrying in {retry_delay} seconds..."
+                )
+                retries += 1
+                time.sleep(retry_delay * retries)
+            except Exception as e:  # noqa: BLE001
+                self.logger.error(f"OpenRouter streaming error: {e}")
+                break
+
+        yield from super().infer_stream(
+            messages,
+            temperature,
+            max_tokens,
+            return_json,
+            custom_format,
+            max_retries,
+            retry_delay,
+        )
+
+    @staticmethod
+    def _extract_text_from_chunk(chunk) -> Optional[str]:
+        try:
+            choices = getattr(chunk, "choices", None)
+        except Exception:
+            choices = None
+        if not choices:
+            return None
+
+        first_choice = choices[0]
+        delta = getattr(first_choice, "delta", None)
+        if delta is None and isinstance(first_choice, dict):
+            delta = first_choice.get("delta")
+        if delta is None:
+            return None
+
+        content = getattr(delta, "content", None)
+        if content is None and isinstance(delta, dict):
+            content = delta.get("content")
+        if not content:
+            return None
+
+        if isinstance(content, str):
+            return content
+
+        if isinstance(content, list):
+            fragments: List[str] = []
+            for item in content:
+                text = getattr(item, "text", None)
+                if text is None and isinstance(item, dict):
+                    text = item.get("text")
+                if text:
+                    fragments.append(text)
+            return "".join(fragments) if fragments else None
+
+        return None
+
+    def infer_stream(
+        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,
+    ):
+        if return_json or custom_format is not None:
+            yield from super().infer_stream(
+                messages,
+                temperature,
+                max_tokens,
+                return_json,
+                custom_format,
+                max_retries,
+                retry_delay,
+            )
+            return
+
+        retries = 0
+        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,
+                    "stream": True,
+                }
+                if eff_temp is not None:
+                    kwargs["temperature"] = eff_temp
+                if eff_max is not None:
+                    kwargs["max_tokens"] = eff_max
+
+                stream = self.client.chat.completions.create(**kwargs)
+                for chunk in stream:
+                    text = self._extract_text_from_chunk(chunk)
+                    if text:
+                        yield text
+                return
+            except self._openai.RateLimitError:  # type: ignore[attr-defined]
+                self.logger.warning(
+                    f"Rate limit exceeded. Retrying in {retry_delay} seconds (streaming)..."
+                )
+                retries += 1
+                time.sleep(retry_delay * retries)
+            except Exception as e:  # noqa: BLE001
+                self.logger.error(f"Streaming error: {e}")
+                break
+
+        # Fallback to non-streaming if streaming fails
+        yield from super().infer_stream(
+            messages,
+            temperature,
+            max_tokens,
+            return_json,
+            custom_format,
+            max_retries,
+            retry_delay,
+        )
+
 
 class GeminiClient(BaseLLMClient):
     def __init__(self, model_name: str, logger):
@@ -367,3 +563,294 @@ class GeminiClient(BaseLLMClient):
                 retries += 1
                 time.sleep(retry_delay * retries)
         return None
+
+    @staticmethod
+    def _extract_text_from_stream_chunk(chunk) -> Optional[str]:
+        if chunk is None:
+            return None
+        text = getattr(chunk, "text", None)
+        if text:
+            return text
+
+        candidates = getattr(chunk, "candidates", None)
+        if not candidates and isinstance(chunk, dict):
+            candidates = chunk.get("candidates")
+        if not candidates:
+            return None
+
+        candidate = candidates[0]
+        content = getattr(candidate, "content", None)
+        if content is None and isinstance(candidate, dict):
+            content = candidate.get("content")
+        if not content:
+            return None
+
+        parts = getattr(content, "parts", None)
+        if parts is None and isinstance(content, dict):
+            parts = content.get("parts")
+        if parts and isinstance(parts, list):
+            fragments: List[str] = []
+            for part in parts:
+                piece = getattr(part, "text", None)
+                if piece is None and isinstance(part, dict):
+                    piece = part.get("text")
+                if piece:
+                    fragments.append(piece)
+            return "".join(fragments) if fragments else None
+
+        final_text = getattr(content, "text", None)
+        if final_text is None and isinstance(content, dict):
+            final_text = content.get("text")
+        return final_text
+
+    def infer_stream(
+        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,
+    ):
+        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()
+                stream = model.generate_content(
+                    contents, generation_config=gen_cfg, stream=True
+                )
+                for chunk in stream:
+                    text = self._extract_text_from_stream_chunk(chunk)
+                    if text:
+                        yield text
+                return
+            except Exception as e:  # noqa: BLE001
+                self.logger.error(f"Gemini streaming error: {e}")
+                retries += 1
+                time.sleep(retry_delay * retries)
+
+        yield from super().infer_stream(
+            messages,
+            temperature,
+            max_tokens,
+            return_json,
+            custom_format,
+            max_retries,
+            retry_delay,
+        )
+
+
+class OpenRouterClient(BaseLLMClient):
+    """
+    OpenRouter client using OpenAI SDK with custom base URL.
+    Supports models from OpenAI, Anthropic, Google, Qwen, and many other providers.
+    """
+
+    # Default model limits based on latest OpenRouter offerings
+    DEFAULT_MODEL_LIMITS: Dict[str, Dict[str, int]] = {
+        "openai/gpt-5": {"max_output": 128_000, "context_window": 400_000},  
+        "openai/gpt-5-codex": {"max_output": 128_000, "context_window": 400_000},
+        "google/gemini-2.5-flash": {"max_output": 65_536, "context_window": 1_000_000},
+        "google/gemini-2.5-pro": {"max_output": 65_536, "context_window": 1_000_000},
+        "anthropic/claude-sonnet-4.5": {"max_output": 16_384, "context_window": 1_000_000},
+    }
+
+    def __init__(self, model_id: str, logger):
+        try:
+            from openai import OpenAI as _OpenAI  # type: ignore
+            import openai as _openai  # type: ignore
+        except Exception as e:  # pragma: no cover
+            raise RuntimeError("openai client is not available") from e
+        
+        self._OpenAI = _OpenAI
+        self._openai = _openai
+        self.model_name = model_id
+        self.logger = logger
+
+        api_key = os.getenv("OPENROUTER_API_KEY")
+        if not api_key:
+            raise ValueError("OPENROUTER_API_KEY not set")
+        
+        # Optional headers for OpenRouter
+        default_headers = {}
+        if site_url := os.getenv("OPENROUTER_SITE_URL"):
+            default_headers["HTTP-Referer"] = site_url
+        if site_name := os.getenv("OPENROUTER_SITE_NAME"):
+            default_headers["X-Title"] = site_name
+        
+        self.client = self._OpenAI(
+            base_url="https://openrouter.ai/api/v1",
+            api_key=api_key,
+            default_headers=default_headers if default_headers else None,
+        )
+
+        # Load env overrides for model limits
+        env_limits_raw = os.getenv("OPENROUTER_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)
+                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:
+                pass
+
+    def _resolve_default_max_tokens(self, model_id: str) -> Optional[int]:
+        """Resolve default max tokens for a model."""
+        # Highest priority: explicit env per-model tokens mapping
+        mapping_raw = os.getenv("OPENROUTER_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
+        
+        # Check for prefix match
+        for k, v in mapping.items():
+            try:
+                if model_id.startswith(k):
+                    return int(v)
+            except Exception:
+                continue
+        
+        # Next: built-in/default-limits map
+        if model_id in self._default_limits:
+            return int(self._default_limits[model_id].get("max_output", 0)) or None
+        
+        # Check for prefix match in default limits
+        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 test_api(self) -> None:
+        """Test API connectivity with minimal token usage."""
+        test_messages = [{"role": "user", "content": "ping"}]
+        token_attempts = [1, 4, 16, 32]
+        last_error: Optional[Exception] = None
+        
+        for tok in token_attempts:
+            try:
+                self.client.chat.completions.create(
+                    model=self.model_name,
+                    messages=test_messages,
+                    max_tokens=tok,
+                    temperature=0,
+                )
+                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"OpenRouter API test failed: {last_error}")
+        raise ValueError("OpenRouter 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]:
+        """Execute inference using OpenRouter."""
+        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_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 temperature is not None:
+                    kwargs["temperature"] = temperature
+                if eff_max is not None:
+                    kwargs["max_tokens"] = eff_max
+                
+                resp = call_fn(**kwargs)
+                
+                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"OpenRouter error: {e}")
+                import traceback
+                traceback.print_exc()
+                break
+        
+        self.logger.error("Max retries exceeded. Unable to complete the request.")
+        return None

tooluniverse/mcp_tool_registry.py

@@ -327,13 +327,12 @@ def _start_server_for_port(port: int, **kwargs):
 
     print(f"🚀 Starting MCP server on port {port} with {len(tools)} tools...")
 
-    # Create SMCP server with stateless mode for compatibility
+    # Create SMCP server for compatibility
     server = _get_smcp()(
         name=config["server_name"],
         auto_expose_tools=False,  # We'll add tools manually
         search_enabled=True,
         max_workers=config.get("max_workers", 5),
-        stateless_http=True,  # Enable stateless mode for MCPAutoLoaderTool compatibility
         **kwargs,
     )
 
@@ -347,8 +346,9 @@ def _start_server_for_port(port: int, **kwargs):
     # Start server in background thread
     def run_server():
         try:
+            # Enable stateless mode for MCPAutoLoaderTool compatibility
             server.run_simple(
-                transport=config["transport"], host=config["host"], port=port
+                transport=config["transport"], host=config["host"], port=port, stateless_http=True
             )
         except Exception as e:
             print(f"❌ Error running MCP server on port {port}: {e}")

tooluniverse/medrxiv_tool.py

@@ -0,0 +1,97 @@
+import requests
+from .base_tool import BaseTool
+from .tool_registry import register_tool
+
+
+@register_tool("MedRxivTool")
+class MedRxivTool(BaseTool):
+    """
+    Search medRxiv preprints using medRxiv's API (same interface as bioRxiv).
+
+    Arguments:
+        query (str): Search term
+        max_results (int): Max results to return (default 10, max 200)
+    """
+
+    def __init__(
+        self,
+        tool_config,
+        base_url="https://api.medrxiv.org/details",
+    ):
+        super().__init__(tool_config)
+        self.base_url = base_url
+
+    def run(self, arguments=None):
+        arguments = arguments or {}
+        query = arguments.get("query")
+        max_results = int(arguments.get("max_results", 10))
+        if not query:
+            return {"error": "`query` parameter is required."}
+        return self._search(query, max_results)
+
+    def _search(self, query, max_results):
+        # Use date range search for recent preprints
+        # Format: /medrxiv/{start_date}/{end_date}/{cursor}/json
+        from datetime import datetime, timedelta
+
+        # Search last 30 days
+        end_date = datetime.now()
+        start_date = end_date - timedelta(days=30)
+
+        url = (f"{self.base_url}/medrxiv/"
+               f"{start_date.strftime('%Y-%m-%d')}/"
+               f"{end_date.strftime('%Y-%m-%d')}/0/json")
+
+        try:
+            resp = requests.get(url, timeout=20)
+            resp.raise_for_status()
+            data = resp.json()
+        except requests.RequestException as e:
+            return {
+                "error": "Network/API error calling medRxiv",
+                "reason": str(e),
+            }
+        except ValueError:
+            return {"error": "Failed to decode medRxiv response as JSON"}
+
+        results = []
+        # The API returns a dictionary with a 'collection' key
+        collection = data.get("collection", [])
+        if not isinstance(collection, list):
+            return {"error": "Unexpected API response format"}
+
+        for item in collection:
+            title = item.get("title")
+            authors = item.get("authors", "")
+            if isinstance(authors, str):
+                authors = [a.strip() for a in authors.split(";") if a.strip()]
+            elif isinstance(authors, list):
+                authors = [str(a).strip() for a in authors if str(a).strip()]
+            else:
+                authors = []
+
+            year = None
+            date = item.get("date")
+            if date and len(date) >= 4 and date[:4].isdigit():
+                year = int(date[:4])
+
+            doi = item.get("doi")
+            url = f"https://www.medrxiv.org/content/{doi}" if doi else None
+
+            # Filter by query if provided
+            if query and query.lower() not in (title or "").lower():
+                continue
+
+            results.append(
+                {
+                    "title": title,
+                    "authors": authors,
+                    "year": year,
+                    "doi": doi,
+                    "url": url,
+                    "abstract": item.get("abstract", ""),
+                    "source": "medRxiv",
+                }
+            )
+
+        return results[:max_results]