Merge branch 'master' into ivana/toxgen/clean-more-stuff-up

Author: sentrivana

scripts/populate_tox/populate_tox.py

@@ -6,6 +6,7 @@
 
 import functools
 import hashlib
+import json
 import os
 import sys
 import time
@@ -36,18 +37,20 @@
 # CUTOFF = datetime.now(tz=timezone.utc) - timedelta(days=365 * 5)
 
 TOX_FILE = Path(__file__).resolve().parent.parent.parent / "tox.ini"
+RELEASES_CACHE_FILE = Path(__file__).resolve().parent / "releases.jsonl"
 ENV = Environment(
     loader=FileSystemLoader(Path(__file__).resolve().parent),
     trim_blocks=True,
     lstrip_blocks=True,
 )
 
-PYPI_COOLDOWN = 0.1  # seconds to wait between requests to PyPI
+PYPI_COOLDOWN = 0.05  # seconds to wait between requests to PyPI
 
 PYPI_PROJECT_URL = "https://pypi.python.org/pypi/{project}/json"
 PYPI_VERSION_URL = "https://pypi.python.org/pypi/{project}/{version}/json"
 CLASSIFIER_PREFIX = "Programming Language :: Python :: "
 
+CACHE = defaultdict(dict)
 
 IGNORE = {
     # Do not try auto-generating the tox entries for these. They will be
@@ -96,9 +99,33 @@ def fetch_package(package: str) -> Optional[dict]:
 
 @functools.cache
 def fetch_release(package: str, version: Version) -> Optional[dict]:
-    """Fetch release metadata from PyPI."""
+    """Fetch release metadata from cache or, failing that, PyPI."""
+    release = _fetch_from_cache(package, version)
+    if release is not None:
+        return release
+
     url = PYPI_VERSION_URL.format(project=package, version=version)
-    return fetch_url(url)
+    release = fetch_url(url)
+    if release is not None:
+        _save_to_cache(package, version, release)
+    return release
+
+
+def _fetch_from_cache(package: str, version: Version) -> Optional[dict]:
+    package = _normalize_name(package)
+    if package in CACHE and str(version) in CACHE[package]:
+        CACHE[package][str(version)]["_accessed"] = True
+        return CACHE[package][str(version)]
+
+    return None
+
+
+def _save_to_cache(package: str, version: Version, release: Optional[dict]) -> None:
+    with open(RELEASES_CACHE_FILE, "a") as releases_cache:
+        releases_cache.write(json.dumps(_normalize_release(release)) + "\n")
+
+    CACHE[_normalize_name(package)][str(version)] = release
+    CACHE[_normalize_name(package)][str(version)]["_accessed"] = True
 
 
 def _prefilter_releases(
@@ -612,6 +639,24 @@ def get_last_updated() -> Optional[datetime]:
     return timestamp
 
 
+def _normalize_name(package: str) -> str:
+    return package.lower().replace("-", "_")
+
+
+def _normalize_release(release: dict) -> dict:
+    """Filter out unneeded parts of the release JSON."""
+    normalized = {
+        "info": {
+            "classifiers": release["info"]["classifiers"],
+            "name": release["info"]["name"],
+            "requires_python": release["info"]["requires_python"],
+            "version": release["info"]["version"],
+            "yanked": release["info"]["yanked"],
+        },
+    }
+    return normalized
+
+
 def main(fail_on_changes: bool = False) -> None:
     """
     Generate tox.ini from the tox.jinja template.
@@ -649,6 +694,20 @@ def main(fail_on_changes: bool = False) -> None:
         f"The SDK supports Python versions {MIN_PYTHON_VERSION} - {MAX_PYTHON_VERSION}."
     )
 
+    # Load file cache
+    global CACHE
+
+    with open(RELEASES_CACHE_FILE) as releases_cache:
+        for line in releases_cache:
+            release = json.loads(line)
+            name = _normalize_name(release["info"]["name"])
+            version = release["info"]["version"]
+            CACHE[name][version] = release
+            CACHE[name][version][
+                "_accessed"
+            ] = False  # for cleaning up unused cache entries
+
+    # Process packages
     packages = defaultdict(list)
 
     for group, integrations in GROUPS.items():
@@ -714,6 +773,21 @@ def main(fail_on_changes: bool = False) -> None:
         packages, update_timestamp=not fail_on_changes, last_updated=last_updated
     )
 
+    # Sort the release cache file
+    releases = []
+    with open(RELEASES_CACHE_FILE) as releases_cache:
+        releases = [json.loads(line) for line in releases_cache]
+    releases.sort(key=lambda r: (r["info"]["name"], r["info"]["version"]))
+    with open(RELEASES_CACHE_FILE, "w") as releases_cache:
+        for release in releases:
+            if (
+                CACHE[_normalize_name(release["info"]["name"])][
+                    release["info"]["version"]
+                ]["_accessed"]
+                is True
+            ):
+                releases_cache.write(json.dumps(release) + "\n")
+
     if fail_on_changes:
         new_file_hash = get_file_hash()
         if old_file_hash != new_file_hash:

scripts/populate_tox/releases.jsonl

@@ -765,18 +765,20 @@ class SPANSTATUS:
     CANCELLED = "cancelled"
     DATA_LOSS = "data_loss"
     DEADLINE_EXCEEDED = "deadline_exceeded"
+    ERROR = "error"  # OTel status code: https://opentelemetry.io/docs/concepts/signals/traces/#span-status
     FAILED_PRECONDITION = "failed_precondition"
     INTERNAL_ERROR = "internal_error"
     INVALID_ARGUMENT = "invalid_argument"
     NOT_FOUND = "not_found"
-    OK = "ok"
+    OK = "ok"  # HTTP 200 and OTel status code: https://opentelemetry.io/docs/concepts/signals/traces/#span-status
     OUT_OF_RANGE = "out_of_range"
     PERMISSION_DENIED = "permission_denied"
     RESOURCE_EXHAUSTED = "resource_exhausted"
     UNAUTHENTICATED = "unauthenticated"
     UNAVAILABLE = "unavailable"
     UNIMPLEMENTED = "unimplemented"
     UNKNOWN_ERROR = "unknown_error"
+    UNSET = "unset"  # OTel status code: https://opentelemetry.io/docs/concepts/signals/traces/#span-status
 
 
 class OP:

sentry_sdk/consts.py

@@ -4,9 +4,10 @@
 import sentry_sdk
 from sentry_sdk.ai.monitoring import record_token_usage
 from sentry_sdk.ai.utils import set_data_normalized, get_start_span_function
-from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
 from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
@@ -52,6 +53,8 @@ def setup_once():
 
 def _capture_exception(exc):
     # type: (Any) -> None
+    set_span_errored()
+
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,
@@ -357,7 +360,13 @@ def _sentry_patched_create_sync(*args, **kwargs):
         integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
         kwargs["integration"] = integration
 
-        return _execute_sync(f, *args, **kwargs)
+        try:
+            return _execute_sync(f, *args, **kwargs)
+        finally:
+            span = sentry_sdk.get_current_span()
+            if span is not None and span.status == SPANSTATUS.ERROR:
+                with capture_internal_exceptions():
+                    span.__exit__(None, None, None)
 
     return _sentry_patched_create_sync
 
@@ -390,6 +399,12 @@ async def _sentry_patched_create_async(*args, **kwargs):
         integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
         kwargs["integration"] = integration
 
-        return await _execute_async(f, *args, **kwargs)
+        try:
+            return await _execute_async(f, *args, **kwargs)
+        finally:
+            span = sentry_sdk.get_current_span()
+            if span is not None and span.status == SPANSTATUS.ERROR:
+                with capture_internal_exceptions():
+                    span.__exit__(None, None, None)
 
     return _sentry_patched_create_async

sentry_sdk/integrations/anthropic.py

@@ -7,6 +7,8 @@
 
 from typing import TYPE_CHECKING
 
+from sentry_sdk.tracing_utils import set_span_errored
+
 if TYPE_CHECKING:
     from typing import Any, Callable, Iterator
     from sentry_sdk.tracing import Span
@@ -84,6 +86,8 @@ def setup_once():
 
 def _capture_exception(exc):
     # type: (Any) -> None
+    set_span_errored()
+
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,

sentry_sdk/integrations/cohere.py

@@ -7,6 +7,7 @@
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
@@ -52,6 +53,8 @@ def setup_once():
 
 def _capture_exception(exc):
     # type: (Any) -> None
+    set_span_errored()
+
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,
@@ -127,8 +130,6 @@ def new_huggingface_task(*args, **kwargs):
         try:
             res = f(*args, **kwargs)
         except Exception as e:
-            # Error Handling
-            span.set_status("error")
             _capture_exception(e)
             span.__exit__(None, None, None)
             raise e from None

sentry_sdk/integrations/huggingface_hub.py

@@ -8,8 +8,7 @@
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.tracing import Span
-from sentry_sdk.tracing_utils import _get_value
+from sentry_sdk.tracing_utils import _get_value, set_span_errored
 from sentry_sdk.utils import logger, capture_internal_exceptions
 
 from typing import TYPE_CHECKING
@@ -26,6 +25,7 @@
         Union,
     )
     from uuid import UUID
+    from sentry_sdk.tracing import Span
 
 
 try:
@@ -116,7 +116,7 @@ def _handle_error(self, run_id, error):
 
             span_data = self.span_map[run_id]
             span = span_data.span
-            span.set_status("unknown")
+            set_span_errored(span)
 
             sentry_sdk.capture_exception(error, span.scope)
 

sentry_sdk/integrations/langchain.py

@@ -7,6 +7,7 @@
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
@@ -83,6 +84,8 @@ def _capture_exception(exc, manual_span_cleanup=True):
     # Close an eventually open span
     # We need to do this by hand because we are not using the start_span context manager
     current_span = sentry_sdk.get_current_span()
+    set_span_errored(current_span)
+
     if manual_span_cleanup and current_span is not None:
         current_span.__exit__(None, None, None)
 

sentry_sdk/integrations/openai.py

@@ -42,7 +42,7 @@ def update_execute_tool_span(span, agent, tool, result):
     if isinstance(result, str) and result.startswith(
         "An error occurred while running the tool"
     ):
-        span.set_status(SPANSTATUS.INTERNAL_ERROR)
+        span.set_status(SPANSTATUS.ERROR)
 
     if should_send_default_pii():
         span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, result)

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

@@ -3,6 +3,7 @@
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import event_from_exception, safe_serialize
 
 from typing import TYPE_CHECKING
@@ -20,6 +21,8 @@
 
 def _capture_exception(exc):
     # type: (Any) -> None
+    set_span_errored()
+
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,