tests, lint, some type

Author: CalebCourier

guardrails/api_client.py

@@ -1,6 +1,6 @@
 import json
 import os
-from typing import Any, Iterable, Optional
+from typing import Any, Iterator, Optional
 
 import requests
 from guardrails_api_client.configuration import Configuration
@@ -80,7 +80,7 @@ def stream_validate(
         guard: Guard,
         payload: ValidatePayload,
         openai_api_key: Optional[str] = None,
-    ) -> Iterable[Any]:
+    ) -> Iterator[Any]:
         _openai_api_key = (
             openai_api_key
             if openai_api_key is not None

guardrails/async_guard.py

@@ -4,7 +4,7 @@
 from opentelemetry import context as otel_context
 from typing import (
     Any,
-    AsyncIterable,
+    AsyncIterator,
     Awaitable,
     Callable,
     Dict,
@@ -188,7 +188,7 @@ async def _execute(
     ) -> Union[
         ValidationOutcome[OT],
         Awaitable[ValidationOutcome[OT]],
-        AsyncIterable[ValidationOutcome[OT]],
+        AsyncIterator[ValidationOutcome[OT]],
     ]:
         self._fill_validator_map()
         self._fill_validators()
@@ -220,7 +220,7 @@ async def __exec(
         ) -> Union[
             ValidationOutcome[OT],
             Awaitable[ValidationOutcome[OT]],
-            AsyncIterable[ValidationOutcome[OT]],
+            AsyncIterator[ValidationOutcome[OT]],
         ]:
             prompt_params = prompt_params or {}
             metadata = metadata or {}
@@ -334,7 +334,7 @@ async def _exec(
     ) -> Union[
         ValidationOutcome[OT],
         Awaitable[ValidationOutcome[OT]],
-        AsyncIterable[ValidationOutcome[OT]],
+        AsyncIterator[ValidationOutcome[OT]],
     ]:
         """Call the LLM asynchronously and validate the output.
 
@@ -416,7 +416,7 @@ async def __call__(
     ) -> Union[
         ValidationOutcome[OT],
         Awaitable[ValidationOutcome[OT]],
-        AsyncIterable[ValidationOutcome[OT]],
+        AsyncIterator[ValidationOutcome[OT]],
     ]:
         """Call the LLM and validate the output. Pass an async LLM API to
         return a coroutine.
@@ -534,7 +534,7 @@ async def parse(
 
     async def _stream_server_call(
         self, *, payload: Dict[str, Any]
-    ) -> AsyncIterable[ValidationOutcome[OT]]:
+    ) -> AsyncIterator[ValidationOutcome[OT]]:
         # TODO: Once server side supports async streaming, this function will need to
         # yield async generators, not generators
         if self._api_client:

guardrails/classes/llm/llm_response.py

@@ -1,6 +1,6 @@
 import asyncio
 from itertools import tee
-from typing import Any, Dict, Iterable, Optional, AsyncIterable
+from typing import Any, Dict, Iterator, Optional, AsyncIterator
 
 from guardrails_api_client import LLMResponse as ILLMResponse
 from pydantic.config import ConfigDict
@@ -19,9 +19,9 @@ class LLMResponse(ILLMResponse):
 
     Attributes:
         output (str): The output from the LLM.
-        stream_output (Optional[Iterable]): A stream of output from the LLM.
+        stream_output (Optional[Iterator]): A stream of output from the LLM.
             Default None.
-        async_stream_output (Optional[AsyncIterable]): An async stream of output
+        async_stream_output (Optional[AsyncIterator]): An async stream of output
             from the LLM.  Default None.
         prompt_token_count (Optional[int]): The number of tokens in the prompt.
             Default None.
@@ -35,8 +35,8 @@ class LLMResponse(ILLMResponse):
     prompt_token_count: Optional[int] = None
     response_token_count: Optional[int] = None
     output: str
-    stream_output: Optional[Iterable] = None
-    async_stream_output: Optional[AsyncIterable] = None
+    stream_output: Optional[Iterator] = None
+    async_stream_output: Optional[AsyncIterator] = None
 
     def to_interface(self) -> ILLMResponse:
         stream_output = None

guardrails/cli/telemetry.py

@@ -6,7 +6,8 @@
 
 def trace_if_enabled(command_name: str):
     if settings.rc.enable_metrics is True:
-        telemetry = HubTelemetry(enabled=True)
+        telemetry = HubTelemetry()
+        telemetry._enabled = True
         telemetry.create_new_span(
             f"guardrails-cli/{command_name}",
             [

guardrails/guard.py

@@ -7,7 +7,7 @@
     Callable,
     Dict,
     Generic,
-    Iterable,
+    Iterator,
     List,
     Optional,
     Sequence,
@@ -302,11 +302,13 @@ def _configure_hub_telemtry(
 
         self._allow_metrics_collection = allow_metrics_collection
 
+        # Initialize Hub Telemetry singleton and get the tracer
+        self._hub_telemetry = HubTelemetry()
+        self._hub_telemetry._enabled = allow_metrics_collection
+
         if allow_metrics_collection is True:
             # Get unique id of user from rc file
             self._user_id = settings.rc.id or ""
-            # Initialize Hub Telemetry singleton and get the tracer
-            self._hub_telemetry = HubTelemetry(enabled=True)
 
     def _fill_validator_map(self):
         # dont init validators if were going to call the server
@@ -696,7 +698,7 @@ def _execute(
         metadata: Optional[Dict],
         full_schema_reask: Optional[bool] = None,
         **kwargs,
-    ) -> Union[ValidationOutcome[OT], Iterable[ValidationOutcome[OT]]]:
+    ) -> Union[ValidationOutcome[OT], Iterator[ValidationOutcome[OT]]]:
         self._fill_validator_map()
         self._fill_validators()
         self._fill_exec_opts(
@@ -837,7 +839,7 @@ def _exec(
         instructions: Optional[str] = None,
         msg_history: Optional[List[Dict]] = None,
         **kwargs,
-    ) -> Union[ValidationOutcome[OT], Iterable[ValidationOutcome[OT]]]:
+    ) -> Union[ValidationOutcome[OT], Iterator[ValidationOutcome[OT]]]:
         api = None
 
         if llm_api is not None or kwargs.get("model") is not None:
@@ -901,7 +903,7 @@ def __call__(
         metadata: Optional[Dict] = None,
         full_schema_reask: Optional[bool] = None,
         **kwargs,
-    ) -> Union[ValidationOutcome[OT], Iterable[ValidationOutcome[OT]]]:
+    ) -> Union[ValidationOutcome[OT], Iterator[ValidationOutcome[OT]]]:
         """Call the LLM and validate the output.
 
         Args:
@@ -1187,7 +1189,7 @@ def _stream_server_call(
         self,
         *,
         payload: Dict[str, Any],
-    ) -> Iterable[ValidationOutcome[OT]]:
+    ) -> Iterator[ValidationOutcome[OT]]:
         if settings.use_server and self._api_client:
             validation_output: Optional[IValidationOutcome] = None
             response = self._api_client.stream_validate(
@@ -1237,7 +1239,7 @@ def _call_server(
         metadata: Optional[Dict] = {},
         full_schema_reask: Optional[bool] = True,
         **kwargs,
-    ) -> Union[ValidationOutcome[OT], Iterable[ValidationOutcome[OT]]]:
+    ) -> Union[ValidationOutcome[OT], Iterator[ValidationOutcome[OT]]]:
         if settings.use_server and self._api_client:
             payload: Dict[str, Any] = {
                 "args": list(args),

guardrails/hub_telemetry/hub_tracing.py

@@ -1,5 +1,14 @@
 from functools import wraps
-from typing import Any, Awaitable, Callable, Dict, Optional
+from typing import (
+    Any,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterator,
+    Optional,
+    TypeVar,
+)
 
 from opentelemetry.trace import Span
 
@@ -9,6 +18,8 @@
 from guardrails.utils.safe_get import safe_get
 from guardrails.utils.hub_telemetry_utils import HubTelemetry
 
+R = TypeVar("R", covariant=True)
+
 
 def get_guard_call_attributes(
     attrs: Dict[str, Any], origin: str, *args, **kwargs
@@ -76,7 +87,9 @@ def get_validator_inference_attributes(
 def get_validator_usage_attributes(
     attrs: Dict[str, Any], response, *args, **kwargs
 ) -> Dict[str, Any]:
-    validator_self = safe_get(args, 0)
+    # We're wrapping a wrapped function,
+    #   so the first arg is the validator service
+    validator_self = safe_get(args, 1)
     if validator_self is not None:
         attrs["validator_name"] = validator_self.rail_alias
         attrs["validator_on_fail"] = validator_self.on_fail_descriptor
@@ -90,11 +103,17 @@ def get_validator_usage_attributes(
 
 
 def add_attributes(
-    span: Span, attrs: Dict[str, Any], name: str, origin: str, response, *args, **kwargs
+    span: Span,
+    attrs: Dict[str, Any],
+    name: str,
+    origin: str,
+    *args,
+    response=None,
+    **kwargs,
 ):
     attrs["origin"] = origin
     if name == "/guard_call":
-        attrs = get_guard_call_attributes(attrs, *args, **kwargs)
+        attrs = get_guard_call_attributes(attrs, origin, *args, **kwargs)
     elif name == "/reasks":
         if response is not None and hasattr(response, "iterations"):
             attrs["reask_count"] = len(response.iterations) - 1
@@ -103,7 +122,7 @@ def add_attributes(
     elif name == "/validator_inference":
         attrs = get_validator_inference_attributes(attrs, *args, **kwargs)
     elif name == "/validator_usage":
-        attrs = get_validator_usage_attributes(attrs, response * args, **kwargs)
+        attrs = get_validator_usage_attributes(attrs, response, *args, **kwargs)
 
     for key, value in attrs.items():
         if value is not None:
@@ -117,9 +136,9 @@ def trace(
     is_parent: Optional[bool] = False,
     **attrs,
 ):
-    def decorator(fn: Callable[..., Any]):
+    def decorator(fn: Callable[..., R]):
         @wraps(fn)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args, **kwargs) -> R:
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
                 context = (
@@ -137,7 +156,9 @@ def wrapper(*args, **kwargs):
                     origin = origin if origin is not None else name
 
                     resp = fn(*args, **kwargs)
-                    add_attributes(span, attrs, origin, resp, *args, **kwargs)
+                    add_attributes(
+                        span, attrs, name, origin, *args, response=resp, **kwargs
+                    )
                     return resp
             else:
                 return fn(*args, **kwargs)
@@ -153,9 +174,9 @@ def async_trace(
     origin: Optional[str] = None,
     is_parent: Optional[bool] = False,
 ):
-    def decorator(fn: Callable[..., Awaitable[Any]]):
+    def decorator(fn: Callable[..., Awaitable[R]]):
         @wraps(fn)
-        async def async_wrapper(*args, **kwargs):
+        async def async_wrapper(*args, **kwargs) -> R:
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
                 context = (
@@ -170,7 +191,7 @@ async def async_wrapper(*args, **kwargs):
 
                     nonlocal origin
                     origin = origin if origin is not None else name
-                    add_attributes(span, {"async": True}, origin, *args, **kwargs)
+                    add_attributes(span, {"async": True}, name, origin, *args, **kwargs)
                     return await fn(*args, **kwargs)
             else:
                 return await fn(*args, **kwargs)
@@ -193,9 +214,9 @@ def trace_stream(
     is_parent: Optional[bool] = False,
     **attrs,
 ):
-    def decorator(fn: Callable[..., Any]):
+    def decorator(fn: Callable[..., Iterator[R]]):
         @wraps(fn)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args, **kwargs) -> Iterator[R]:
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
                 context = (
@@ -212,7 +233,7 @@ def wrapper(*args, **kwargs):
 
                     nonlocal origin
                     origin = origin if origin is not None else name
-                    add_attributes(span, attrs, name, origin, None, *args, **kwargs)
+                    add_attributes(span, attrs, name, origin, *args, **kwargs)
                     return _run_gen(fn, *args, **kwargs)
             else:
                 return fn(*args, **kwargs)
@@ -235,9 +256,9 @@ def async_trace_stream(
     is_parent: Optional[bool] = False,
     **attrs,
 ):
-    def decorator(fn: Callable[..., Awaitable[Any]]):
+    def decorator(fn: Callable[..., AsyncIterator[R]]):
         @wraps(fn)
-        async def wrapper(*args, **kwargs):
+        async def wrapper(*args, **kwargs) -> AsyncIterator[R]:
             hub_telemetry = HubTelemetry()
             if hub_telemetry._enabled and hub_telemetry._tracer is not None:
                 context = (
@@ -254,7 +275,7 @@ async def wrapper(*args, **kwargs):
 
                     nonlocal origin
                     origin = origin if origin is not None else name
-                    add_attributes(span, attrs, name, origin, None, *args, **kwargs)
+                    add_attributes(span, attrs, name, origin, *args, **kwargs)
                     return _run_async_gen(fn, *args, **kwargs)
             else:
                 return fn(*args, **kwargs)