don't type trace wrappers

Author: CalebCourier

guardrails/classes/llm/llm_response.py

@@ -73,7 +73,7 @@ def to_dict(self) -> Dict[str, Any]:
     def from_interface(cls, i_llm_response: ILLMResponse) -> "LLMResponse":
         stream_output = None
         if i_llm_response.stream_output:
-            stream_output = [so for so in i_llm_response.stream_output]
+            stream_output = iter([so for so in i_llm_response.stream_output])
 
         async_stream_output = None
         if i_llm_response.async_stream_output:

guardrails/hub_telemetry/hub_tracing.py

@@ -1,11 +1,7 @@
 from functools import wraps
 from typing import (
     Any,
-    AsyncIterator,
-    Awaitable,
-    Callable,
     Dict,
-    Iterator,
     Optional,
     TypeVar,
 )
@@ -136,9 +132,11 @@ def trace(
     is_parent: Optional[bool] = False,
     **attrs,
 ):
-    def decorator(fn: Callable[..., R]):
+    # def decorator(fn: Callable[..., R]):
+    def decorator(fn):
         @wraps(fn)
-        def wrapper(*args, **kwargs) -> R:
+        # def wrapper(*args, **kwargs) -> R:
+        def wrapper(*args, **kwargs):
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
                 context = (
@@ -174,9 +172,11 @@ def async_trace(
     origin: Optional[str] = None,
     is_parent: Optional[bool] = False,
 ):
-    def decorator(fn: Callable[..., Awaitable[R]]):
+    # def decorator(fn: Callable[..., Awaitable[R]]):
+    def decorator(fn):
         @wraps(fn)
-        async def async_wrapper(*args, **kwargs) -> R:
+        # 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 = (
@@ -214,9 +214,11 @@ def trace_stream(
     is_parent: Optional[bool] = False,
     **attrs,
 ):
-    def decorator(fn: Callable[..., Iterator[R]]):
+    # def decorator(fn: Callable[..., Iterator[R]]):
+    def decorator(fn):
         @wraps(fn)
-        def wrapper(*args, **kwargs) -> Iterator[R]:
+        # 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 = (
@@ -256,9 +258,10 @@ def async_trace_stream(
     is_parent: Optional[bool] = False,
     **attrs,
 ):
-    def decorator(fn: Callable[..., AsyncIterator[R]]):
+    # def decorator(fn: Callable[..., AsyncIterator[R]]):
+    def decorator(fn):
         @wraps(fn)
-        async def wrapper(*args, **kwargs) -> AsyncIterator[R]:
+        async def wrapper(*args, **kwargs):
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
                 context = (

guardrails/integrations/databricks/ml_flow_instrumentor.py

@@ -7,7 +7,6 @@
     Awaitable,
     Callable,
     Coroutine,
-    Generator,
     Iterator,
     Union,
 )
@@ -220,12 +219,12 @@ def trace_step_wrapper(*args, **kwargs) -> Iteration:
         return trace_step_wrapper
 
     def _instrument_stream_runner_step(
-        self, runner_step: Callable[..., Generator[ValidationOutcome[OT], None, None]]
+        self, runner_step: Callable[..., Iterator[ValidationOutcome[OT]]]
     ):
         @wraps(runner_step)
         def trace_stream_step_wrapper(
             *args, **kwargs
-        ) -> Generator[ValidationOutcome[OT], None, None]:
+        ) -> Iterator[ValidationOutcome[OT]]:
             with mlflow.start_span(
                 name="guardrails/guard/step",
                 span_type="step",

guardrails/validator_service/__init__.py

@@ -12,6 +12,9 @@
 )
 from guardrails.types import ValidatorMap
 from guardrails.telemetry.legacy_validator_tracing import trace_validation_result
+
+# Keep this imported for backwards compatibility
+from guardrails.validator_service.validator_service_base import ValidatorServiceBase  # noqa
 from guardrails.validator_service.async_validator_service import AsyncValidatorService
 from guardrails.validator_service.sequential_validator_service import (
     SequentialValidatorService,