arize-phoenix 2.7.0__py3-none-any.whl → 2.8.0__py3-none-any.whl

phoenix/experimental/evals/models/vertex.py

@@ -46,12 +46,6 @@ class GeminiModel(BaseEvalModel):
     def __post_init__(self) -> None:
         self._init_client()
         self._init_rate_limiter()
-        self.retry = self._retry(
-            error_types=[],  # default to catching all errors
-            min_seconds=self.retry_min_seconds,
-            max_seconds=self.retry_max_seconds,
-            max_retries=self.max_retries,
-        )
 
     def reload_client(self) -> None:
         self._init_client()
@@ -115,30 +109,17 @@ class GeminiModel(BaseEvalModel):
         # instruction is an invalid input to Gemini models, it is passed in by
         # BaseEvalModel.__call__ and needs to be removed
         kwargs.pop("instruction", None)
-        response = self._generate_with_retry(
-            prompt=prompt,
-            generation_config=self.generation_config,
-            **kwargs,
-        )
 
-        return str(response)
-
-    def _generate_with_retry(
-        self, prompt: str, generation_config: Dict[str, Any], **kwargs: Any
-    ) -> Any:
-        @self.retry
         @self._rate_limiter.limit
-        def _completion_with_retry(**kwargs: Any) -> Any:
+        def _rate_limited_completion(
+            prompt: str, generation_config: Dict[str, Any], **kwargs: Any
+        ) -> Any:
             response = self._model.generate_content(
                 contents=prompt, generation_config=generation_config, **kwargs
             )
             return self._parse_response_candidates(response)
 
-        return _completion_with_retry(**kwargs)
-
-    async def _async_generate(self, prompt: str, **kwargs: Dict[str, Any]) -> str:
-        kwargs.pop("instruction", None)
-        response = await self._async_generate_with_retry(
+        response = _rate_limited_completion(
             prompt=prompt,
             generation_config=self.generation_config,
             **kwargs,
@@ -146,18 +127,27 @@ class GeminiModel(BaseEvalModel):
 
         return str(response)
 
-    async def _async_generate_with_retry(
-        self, prompt: str, generation_config: Dict[str, Any], **kwargs: Any
-    ) -> Any:
-        @self.retry
+    async def _async_generate(self, prompt: str, **kwargs: Dict[str, Any]) -> str:
+        # instruction is an invalid input to Gemini models, it is passed in by
+        # BaseEvalModel.__call__ and needs to be removed
+        kwargs.pop("instruction", None)
+
         @self._rate_limiter.alimit
-        async def _completion_with_retry(**kwargs: Any) -> Any:
+        async def _rate_limited_completion(
+            prompt: str, generation_config: Dict[str, Any], **kwargs: Any
+        ) -> Any:
             response = await self._model.generate_content_async(
                 contents=prompt, generation_config=generation_config, **kwargs
             )
             return self._parse_response_candidates(response)
 
-        return await _completion_with_retry(**kwargs)
+        response = await _rate_limited_completion(
+            prompt=prompt,
+            generation_config=self.generation_config,
+            **kwargs,
+        )
+
+        return str(response)
 
     def _parse_response_candidates(self, response: Any) -> Any:
         if hasattr(response, "candidates"):

phoenix/experimental/evals/models/vertexai.py

@@ -52,18 +52,6 @@ class VertexAIModel(BaseEvalModel):
 
             self._vertexai = vertexai
             self._google_exceptions = google_exceptions
-            self._google_api_retry_errors = [
-                self._google_exceptions.ResourceExhausted,
-                self._google_exceptions.ServiceUnavailable,
-                self._google_exceptions.Aborted,
-                self._google_exceptions.DeadlineExceeded,
-            ]
-            self.retry = self._retry(
-                error_types=self._google_api_retry_errors,
-                min_seconds=self.retry_min_seconds,
-                max_seconds=self.retry_max_seconds,
-                max_retries=self.max_retries,
-            )
         except ImportError:
             self._raise_import_error(
                 package_display_name="VertexAI",
@@ -97,19 +85,12 @@ class VertexAIModel(BaseEvalModel):
 
     def _generate(self, prompt: str, **kwargs: Dict[str, Any]) -> str:
         invoke_params = self.invocation_params
-        response = self._generate_with_retry(
+        response = self._model.predict(
             prompt=prompt,
             **invoke_params,
         )
         return str(response.text)
 
-    def _generate_with_retry(self, **kwargs: Any) -> Any:
-        @self.retry
-        def _completion_with_retry(**kwargs: Any) -> Any:
-            return self._model.predict(**kwargs)
-
-        return _completion_with_retry(**kwargs)
-
     @property
     def is_codey_model(self) -> bool:
         return is_codey_model(self.tuned_model_name or self.model_name)

phoenix/server/api/schema.py

@@ -2,7 +2,6 @@ from collections import defaultdict
 from datetime import datetime
 from itertools import chain
 from typing import Dict, List, Optional, Set, Tuple, Union, cast
-from uuid import UUID
 
 import numpy as np
 import numpy.typing as npt
@@ -22,7 +21,7 @@ from phoenix.server.api.input_types.Coordinates import (
 from phoenix.server.api.input_types.SpanSort import SpanSort
 from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
 from phoenix.trace.dsl import SpanFilter
-from phoenix.trace.schemas import SpanID
+from phoenix.trace.schemas import SpanID, TraceID
 
 from .context import Context
 from .input_types.TimeRange import TimeRange
@@ -264,7 +263,7 @@ class Query:
                 root_spans_only=root_spans_only,
             )
         else:
-            spans = chain.from_iterable(map(traces.get_trace, map(UUID, trace_ids)))
+            spans = chain.from_iterable(map(traces.get_trace, map(TraceID, trace_ids)))
         if predicate:
             spans = filter(predicate, spans)
         if sort: