tooluniverse 1.0.5__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):
@@ -368,6 +564,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):
     """

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]

tooluniverse/openaire_tool.py

@@ -0,0 +1,145 @@
+import requests
+from .base_tool import BaseTool
+from .tool_registry import register_tool
+
+
+@register_tool("OpenAIRETool")
+class OpenAIRETool(BaseTool):
+    """
+    Search OpenAIRE Explore for research products (publications by default).
+
+    Parameters (arguments):
+        query (str): Query string
+        max_results (int): Max number of results (default 10, max 100)
+        type (str): product type filter: publications | datasets | software
+    """
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.base_url = "https://api.openaire.eu/search/publications"
+
+    def run(self, arguments=None):
+        arguments = arguments or {}
+        query = arguments.get("query")
+        max_results = int(arguments.get("max_results", 10))
+        prod_type = arguments.get("type", "publications")
+
+        if not query:
+            return {"error": "`query` parameter is required."}
+
+        endpoint = self._endpoint_for_type(prod_type)
+        if endpoint is None:
+            return {
+                "error": (
+                    "Unsupported type. Use publications/datasets/software."
+                ),
+            }
+
+        params = {
+            "format": "json",
+            "size": max(1, min(max_results, 100)),
+            "query": query,
+        }
+        try:
+            resp = requests.get(endpoint, params=params, timeout=20)
+            resp.raise_for_status()
+            data = resp.json()
+        except requests.RequestException as e:
+            return {
+                "error": "Network/API error calling OpenAIRE",
+                "reason": str(e),
+            }
+        except ValueError:
+            return {"error": "Failed to decode OpenAIRE response as JSON"}
+
+        return self._normalize(data, prod_type)
+
+    def _endpoint_for_type(self, prod_type):
+        if prod_type == "publications":
+            return "https://api.openaire.eu/search/publications"
+        if prod_type == "datasets":
+            return "https://api.openaire.eu/search/datasets"
+        if prod_type == "software":
+            return "https://api.openaire.eu/search/software"
+        return None
+
+    def _normalize(self, data, prod_type):
+        results = []
+        # OpenAIRE JSON has a root 'response' with 'results' → 'result' list
+        try:
+            items = (
+                data.get("response", {})
+                .get("results", {})
+                .get("result", [])
+            )
+        except Exception:
+            items = []
+
+        for it in items:
+            # header may contain identifiers, not used presently
+            _ = (
+                it.get("header", {})
+                if isinstance(it.get("header"), dict)
+                else {}
+            )
+            metadata = (
+                it.get("metadata", {})
+                if isinstance(it.get("metadata"), dict)
+                else {}
+            )
+            title = None
+            authors = []
+            year = None
+            doi = None
+            url = None
+
+            # Titles can be nested in 'oaf:result' structure
+            result_obj = metadata.get("oaf:result", {})
+            if isinstance(result_obj, dict):
+                t = result_obj.get("title")
+                if isinstance(t, list) and t:
+                    title = t[0].get("$")
+                elif isinstance(t, dict):
+                    title = t.get("$")
+
+                # Authors
+                creators = result_obj.get("creator", [])
+                if isinstance(creators, list):
+                    for c in creators:
+                        name = c.get("$")
+                        if name:
+                            authors.append(name)
+
+                # Year
+                date_obj = (
+                    result_obj.get("dateofacceptance")
+                    or result_obj.get("date")
+                )
+                if isinstance(date_obj, dict):
+                    year = date_obj.get("year") or date_obj.get("$")
+
+                # DOI and URL
+                pid = result_obj.get("pid", [])
+                if isinstance(pid, list):
+                    for p in pid:
+                        if p.get("@classid") == "doi":
+                            doi = p.get("$")
+                bestaccessright = result_obj.get("bestaccessright", {})
+                if isinstance(bestaccessright, dict):
+                    url_value = bestaccessright.get("$")
+                    if url_value:
+                        url = url_value
+
+            results.append(
+                {
+                    "title": title,
+                    "authors": authors,
+                    "year": year,
+                    "doi": doi,
+                    "url": url,
+                    "type": prod_type,
+                    "source": "OpenAIRE",
+                }
+            )
+
+        return results

tooluniverse/osf_preprints_tool.py

@@ -0,0 +1,67 @@
+import requests
+from .base_tool import BaseTool
+from .tool_registry import register_tool
+
+
+@register_tool("OSFPreprintsTool")
+class OSFPreprintsTool(BaseTool):
+    """
+    Search OSF Preprints via OSF API v2 filters.
+
+    Parameters (arguments):
+        query (str): Query string
+        max_results (int): Max results (default 10, max 100)
+        provider (str): Optional preprint provider (e.g., 'osf', 'psyarxiv')
+    """
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.base_url = "https://api.osf.io/v2/preprints/"
+
+    def run(self, arguments=None):
+        arguments = arguments or {}
+        query = arguments.get("query")
+        max_results = int(arguments.get("max_results", 10))
+        provider = arguments.get("provider")
+
+        if not query:
+            return {"error": "`query` parameter is required."}
+
+        params = {
+            "page[size]": max(1, min(max_results, 100)),
+            "filter[title]": query,
+        }
+        if provider:
+            params["filter[provider]"] = provider
+
+        try:
+            resp = requests.get(self.base_url, params=params, timeout=20)
+            resp.raise_for_status()
+            data = resp.json()
+        except requests.RequestException as e:
+            return {"error": "Network/API error calling OSF", "reason": str(e)}
+        except ValueError:
+            return {"error": "Failed to decode OSF response as JSON"}
+
+        results = []
+        for item in data.get("data", []):
+            attrs = item.get("attributes", {})
+            title = attrs.get("title")
+            date_published = attrs.get("date_published")
+            is_published = attrs.get("is_published")
+            doi = attrs.get("doi")
+            links_obj = item.get("links", {})
+            url = links_obj.get("html") or links_obj.get("self")
+
+            results.append(
+                {
+                    "title": title,
+                    "date_published": date_published,
+                    "published": is_published,
+                    "doi": doi,
+                    "url": url,
+                    "source": "OSF Preprints",
+                }
+            )
+
+        return results