fix context issues

Author: CalebCourier

guardrails/cli/create.py

@@ -15,8 +15,8 @@
 console = Console()
 
 
-@trace(name="guardrails-cli/create", is_parent=True)
 @gr_cli.command(name="create")
+@trace(name="guardrails-cli/create")
 def create_command(
     validators: Optional[str] = typer.Option(
         default="",

guardrails/cli/hub/create_validator.py

@@ -147,8 +147,8 @@ def test_failure_case(self):
 )
 
 
-@trace(name="guardrails-cli/hub/create-validator", is_parent=True)
 @hub_command.command(name="create-validator")
+@trace(name="guardrails-cli/hub/create-validator")
 def create_validator(
     name: str = typer.Argument(help="The name for your validator."),
     filepath: str = typer.Argument(

guardrails/cli/hub/install.py

@@ -10,8 +10,8 @@
 from guardrails.cli.version import version_warnings_if_applicable
 
 
-@trace(name="guardrails-cli/hub/install", is_parent=True)
 @hub_command.command()
+@trace(name="guardrails-cli/hub/install")
 def install(
     package_uris: List[str] = typer.Argument(
         ...,

guardrails/cli/hub/list.py

@@ -7,8 +7,8 @@
 from .console import console
 
 
-@trace(name="guardrails-cli/hub/list", is_parent=True)
 @hub_command.command(name="list")
+@trace(name="guardrails-cli/hub/list")
 def list():
     """List all installed validators."""
     site_packages = get_site_packages_location()

guardrails/cli/hub/submit.py

@@ -11,8 +11,8 @@
 from guardrails.hub_telemetry.hub_tracing import trace
 
 
-@trace(name="guardrails-cli/hub/submit", is_parent=True)
 @hub_command.command(name="submit")
+@trace(name="guardrails-cli/hub/submit")
 def submit(
     package_name: str = typer.Argument(help="The package name for your validator."),
     filepath: str = typer.Argument(

guardrails/cli/hub/uninstall.py

@@ -74,8 +74,8 @@ def uninstall_hub_module(manifest: Manifest, site_packages: str):
         sys.exit(1)
 
 
-@trace(name="guardrails-cli/hub/uninstall", is_parent=True)
 @hub_command.command()
+@trace(name="guardrails-cli/hub/uninstall")
 def uninstall(
     package_uri: str = typer.Argument(
         help="URI to the package to uninstall. Example: hub://guardrails/regex_match."

guardrails/cli/telemetry.py

@@ -19,6 +19,4 @@ def trace_if_enabled(command_name: str):
                 ("machine", platform.machine()),
                 ("processor", platform.processor()),
             ],
-            True,
-            False,
         )

guardrails/cli/validate.py

@@ -15,8 +15,8 @@ def validate_llm_output(rail: str, llm_output: str) -> Union[str, Dict, List, No
     return result.validated_output
 
 
-@trace(name="guardrails-cli/validate", is_parent=True)
 @guardrails.command()
+@trace(name="guardrails-cli/validate")
 def validate(
     rail: str = typer.Argument(
         ..., help="Path to the rail spec.", exists=True, file_okay=True, dir_okay=False

guardrails/hub_telemetry/hub_tracing.py

@@ -7,6 +7,7 @@
 )
 
 from opentelemetry.trace import Span
+from opentelemetry.trace.propagation import set_span_in_context
 
 from guardrails.classes.validation.validation_result import ValidationResult
 from guardrails.hub_token.token import VALIDATOR_HUB_SERVICE
@@ -129,27 +130,20 @@ def trace(
     *,
     name: str,
     origin: Optional[str] = None,
-    is_parent: Optional[bool] = False,
     **attrs,
 ):
-    # def decorator(fn: Callable[..., R]):
     def decorator(fn):
         @wraps(fn)
-        # def wrapper(*args, **kwargs) -> R:
         def wrapper(*args, **kwargs):
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
-                context = (
-                    hub_telemetry.extract_current_context() if not is_parent else None
-                )
                 with hub_telemetry._tracer.start_span(
                     name,
-                    context=context,
+                    context=hub_telemetry.extract_current_context(),
                     set_status_on_exception=True,
                 ) as span:  # noqa
-                    if is_parent:
-                        # Inject the current context
-                        hub_telemetry.inject_current_context()
+                    context = set_span_in_context(span)
+                    hub_telemetry.inject_current_context(context=context)
                     nonlocal origin
                     origin = origin if origin is not None else name
 
@@ -170,24 +164,19 @@ def async_trace(
     *,
     name: str,
     origin: Optional[str] = None,
-    is_parent: Optional[bool] = False,
 ):
-    # def decorator(fn: Callable[..., Awaitable[R]]):
     def decorator(fn):
         @wraps(fn)
-        # async def async_wrapper(*args, **kwargs) -> R:
         async def async_wrapper(*args, **kwargs):
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
-                context = (
-                    hub_telemetry.extract_current_context() if not is_parent else None
-                )
-                with hub_telemetry._tracer.start_as_current_span(
-                    name, context=context
+                with hub_telemetry._tracer.start_span(
+                    name,
+                    context=hub_telemetry.extract_current_context(),
+                    set_status_on_exception=True,
                 ) as span:  # noqa
-                    if is_parent:
-                        # Inject the current context
-                        hub_telemetry.inject_current_context()
+                    context = set_span_in_context(span)
+                    hub_telemetry.inject_current_context(context=context)
 
                     nonlocal origin
                     origin = origin if origin is not None else name
@@ -211,27 +200,20 @@ def trace_stream(
     *,
     name: str,
     origin: Optional[str] = None,
-    is_parent: Optional[bool] = False,
     **attrs,
 ):
-    # def decorator(fn: Callable[..., Iterator[R]]):
     def decorator(fn):
         @wraps(fn)
-        # def wrapper(*args, **kwargs) -> Iterator[R]:
         def wrapper(*args, **kwargs):
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
-                context = (
-                    hub_telemetry.extract_current_context() if not is_parent else None
-                )
                 with hub_telemetry._tracer.start_span(
                     name,
-                    context=context,
+                    context=hub_telemetry.extract_current_context(),
                     set_status_on_exception=True,
                 ) as span:  # noqa
-                    if is_parent:
-                        # Inject the current context
-                        hub_telemetry.inject_current_context()
+                    context = set_span_in_context(span)
+                    hub_telemetry.inject_current_context(context=context)
 
                     nonlocal origin
                     origin = origin if origin is not None else name
@@ -255,26 +237,20 @@ def async_trace_stream(
     *,
     name: str,
     origin: Optional[str] = None,
-    is_parent: Optional[bool] = False,
     **attrs,
 ):
-    # def decorator(fn: Callable[..., AsyncIterator[R]]):
     def decorator(fn):
         @wraps(fn)
         async def wrapper(*args, **kwargs):
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
-                context = (
-                    hub_telemetry.extract_current_context() if not is_parent else None
-                )
                 with hub_telemetry._tracer.start_span(
                     name,
-                    context=context,
+                    context=hub_telemetry.extract_current_context(),
                     set_status_on_exception=True,
                 ) as span:  # noqa
-                    if is_parent:
-                        # Inject the current context
-                        hub_telemetry.inject_current_context()
+                    context = set_span_in_context(span)
+                    hub_telemetry.inject_current_context(context=context)
 
                     nonlocal origin
                     origin = origin if origin is not None else name

guardrails/utils/hub_telemetry_utils.py

@@ -9,6 +9,7 @@
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from opentelemetry.trace.propagation import set_span_in_context
 
 
 class HubTelemetry:
@@ -57,10 +58,10 @@ def initialize_tracer(
         self._enabled = enabled
         self._carrier = {}
         self._service_name = service_name
-        # self._endpoint = "http://localhost:4318/v1/traces"
-        self._endpoint = (
-            "https://hty0gc1ok3.execute-api.us-east-1.amazonaws.com/v1/traces"
-        )
+        self._endpoint = "http://localhost:5318/v1/traces"
+        # self._endpoint = (
+        #     "https://hty0gc1ok3.execute-api.us-east-1.amazonaws.com/v1/traces"
+        # )
         self._tracer_name = tracer_name
 
         # Create a resource
@@ -85,11 +86,14 @@ def initialize_tracer(
 
         self._prop = TraceContextTextMapPropagator()
 
-    def inject_current_context(self) -> None:
+    def inject_current_context(self, context=None) -> None:
         """Injects the current context into the carrier."""
         if not self._prop:
             return
-        self._prop.inject(carrier=self._carrier)
+        if context is not None:
+            self._prop.inject(carrier=self._carrier, context=context)
+        else:
+            self._prop.inject(carrier=self._carrier)
 
     def extract_current_context(self):
         """Extracts the current context from the carrier."""
@@ -98,13 +102,7 @@ def extract_current_context(self):
         context = self._prop.extract(carrier=self._carrier)
         return context
 
-    def create_new_span(
-        self,
-        span_name: str,
-        attributes: list,
-        is_parent: bool,  # Inject current context if IS a parent span
-        has_parent: bool,  # Extract current context if HAS a parent span
-    ):
+    def create_new_span(self, span_name: str, attributes: list):
         """Creates a new span within the tracer with the given name and
         attributes.
 
@@ -121,13 +119,12 @@ def create_new_span(
         """
         if self._tracer is None:
             return
-        with self._tracer.start_as_current_span(
+        with self._tracer.start_span(
             span_name,  # type: ignore (Fails in Python 3.9 for invalid reason)
-            context=self.extract_current_context() if has_parent else None,
+            context=self.extract_current_context(),
         ) as span:
-            if is_parent:
-                # Inject the current context
-                self.inject_current_context()
+            context = set_span_in_context(span)
+            self.inject_current_context(context=context)
 
             for attribute in attributes:
                 span.set_attribute(attribute[0], attribute[1])