tooluniverse 1.0.5__py3-none-any.whl → 1.0.7__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,75 @@ 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_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):
@@ -368,6 +460,96 @@ class GeminiClient(BaseLLMClient):
                 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):
     """
@@ -377,11 +559,14 @@ class OpenRouterClient(BaseLLMClient):
 
     # 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": {"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},
+        "anthropic/claude-sonnet-4.5": {
+            "max_output": 16_384,
+            "context_window": 1_000_000,
+        },
     }
 
     def __init__(self, model_id: str, logger):
@@ -390,7 +575,7 @@ class OpenRouterClient(BaseLLMClient):
             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
@@ -399,14 +584,14 @@ class OpenRouterClient(BaseLLMClient):
         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,
@@ -445,13 +630,13 @@ class OpenRouterClient(BaseLLMClient):
                 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:
@@ -459,11 +644,11 @@ class OpenRouterClient(BaseLLMClient):
                     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:
@@ -471,7 +656,7 @@ class OpenRouterClient(BaseLLMClient):
                     return int(v.get("max_output", 0)) or None
             except Exception:
                 continue
-        
+
         return None
 
     def test_api(self) -> None:
@@ -479,7 +664,7 @@ class OpenRouterClient(BaseLLMClient):
         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(
@@ -499,7 +684,7 @@ class OpenRouterClient(BaseLLMClient):
                 ) 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")
@@ -521,39 +706,39 @@ class OpenRouterClient(BaseLLMClient):
             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..."
@@ -563,8 +748,86 @@ class OpenRouterClient(BaseLLMClient):
             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
+
+
+class VLLMClient(BaseLLMClient):
+    def __init__(self, model_name: str, server_url: str, logger):
+        try:
+            from openai import OpenAI
+        except Exception as e:
+            raise RuntimeError("openai package not available for vLLM client") from e
+
+        if not server_url:
+            raise ValueError("VLLM_SERVER_URL must be provided")
+
+        self.model_name = model_name
+        # Ensure server_url ends with /v1 for OpenAI-compatible API
+        if not server_url.endswith("/v1"):
+            server_url = server_url.rstrip("/") + "/v1"
+        self.server_url = server_url
+        self.logger = logger
+
+        self.client = OpenAI(
+            api_key="EMPTY",
+            base_url=self.server_url,
+        )
+
+    def test_api(self) -> None:
+        test_messages = [{"role": "user", "content": "ping"}]
+        try:
+            self.client.chat.completions.create(
+                model=self.model_name,
+                messages=test_messages,
+                max_tokens=8,
+                temperature=0,
+            )
+        except Exception as e:
+            raise ValueError(f"vLLM API test failed: {e}")
+
+    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 custom_format is not None:
+            self.logger.warning("vLLM does not support custom format, ignoring")
+
+        retries = 0
+        while retries < max_retries:
+            try:
+                kwargs: Dict[str, Any] = {
+                    "model": self.model_name,
+                    "messages": messages,
+                }
+
+                if temperature is not None:
+                    kwargs["temperature"] = temperature
+
+                if max_tokens is not None:
+                    kwargs["max_tokens"] = max_tokens
+
+                if return_json:
+                    kwargs["response_format"] = {"type": "json_object"}
+
+                resp = self.client.chat.completions.create(**kwargs)
+                return resp.choices[0].message.content
+
+            except Exception as e:
+                self.logger.error(f"vLLM error: {e}")
+                retries += 1
+                if retries < max_retries:
+                    time.sleep(retry_delay * retries)
+
+        self.logger.error("Max retries exceeded for vLLM request")
+        return None

tooluniverse/mcp_tool_registry.py

@@ -348,7 +348,10 @@ def _start_server_for_port(port: int, **kwargs):
         try:
             # Enable stateless mode for MCPAutoLoaderTool compatibility
             server.run_simple(
-                transport=config["transport"], host=config["host"], port=port, stateless_http=True
+                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,116 @@
+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 365 days to get more comprehensive results
+        end_date = datetime.now()
+        start_date = end_date - timedelta(days=365)
+
+        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 - search in both title and abstract
+            if query:
+                title_match = query.lower() in (title or "").lower()
+                abstract_match = (
+                    query.lower() in (item.get("abstract", "") or "").lower()
+                )
+                if not (title_match or abstract_match):
+                    continue
+
+            results.append(
+                {
+                    "title": title or "Title not available",
+                    "authors": (
+                        authors if authors else "Author information not available"
+                    ),
+                    "year": year,
+                    "doi": doi or "DOI not available",
+                    "url": url or "URL not available",
+                    "abstract": item.get("abstract", "Abstract not available"),
+                    "source": "medRxiv",
+                    "data_quality": {
+                        "has_abstract": bool(
+                            item.get("abstract")
+                            and item.get("abstract") != "Abstract not available"
+                        ),
+                        "has_authors": bool(authors),
+                        "has_year": bool(year),
+                        "has_doi": bool(doi),
+                        "has_url": bool(url),
+                    },
+                }
+            )
+
+        return results[:max_results]