Merge pull request #957 from parea-ai/fix-nested-cattrs

Avoid nested cattrs check on TraceLogTree

Author: jalexanderII

parea/client.py

@@ -32,6 +32,7 @@
     ListExperimentUUIDsFilters,
     ProjectSchema,
     TestCaseCollection,
+    TraceLog,
     TraceLogFilters,
     TraceLogTree,
     UseDeployedPrompt,
@@ -402,32 +403,32 @@ def _update_data_and_trace(self, data: Completion) -> Completion:
 
         return data
 
-    def get_trace_log(self, trace_id: str) -> TraceLogTree:
+    def get_trace_log(self, trace_id: str, return_children: bool = True) -> Union[TraceLogTree, TraceLog]:
         response = self._client.request("GET", GET_TRACE_LOG_ENDPOINT.format(trace_id=trace_id))
-        return structure_trace_log_from_api(response.json())
+        return structure_trace_log_from_api(response.json(), return_children)
 
-    def get_trace_log_scores(self, trace_id: str, check_context: bool = True) -> List[EvaluationResult]:
+    def get_trace_log_scores(self, trace_id: str, check_context: bool = True, return_children: bool = True) -> List[EvaluationResult]:
         """
         Get the scores from the trace log. If the scores are not present in the trace log, fetch them from the DB.
         Args:
             trace_id: The trace id to get the scores for.
             check_context: If True, will check the context for the scores first before fetching from the DB.
+            return_children: If True, will return the children logs in the tree structure.
 
         Returns: A list of EvaluationResult objects.
         """
         # try to get trace_id scores from context
         if check_context:
             if scores := (trace_data.get()[trace_id].scores or []):
-                print("Scores from context", scores)
                 return scores
 
         response = self._client.request("GET", GET_TRACE_LOG_ENDPOINT.format(trace_id=trace_id))
-        tree: TraceLogTree = structure_trace_log_from_api(response.json())
+        tree = structure_trace_log_from_api(response.json(), return_children)
         return extract_scores(tree)
 
-    async def aget_trace_log(self, trace_id: str) -> TraceLogTree:
+    async def aget_trace_log(self, trace_id: str, return_children: bool = True) -> Union[TraceLogTree, TraceLog]:
         response = await self._client.request_async("GET", GET_TRACE_LOG_ENDPOINT.format(trace_id=trace_id))
-        return structure_trace_log_from_api(response.json())
+        return structure_trace_log_from_api(response.json(), return_children)
 
     def list_experiments(self, filter_conditions: Optional[ListExperimentUUIDsFilters] = ListExperimentUUIDsFilters()) -> List[ExperimentWithPinnedStatsSchema]:
         response = self._client.request("POST", LIST_EXPERIMENTS_ENDPOINT, data=asdict(filter_conditions))
@@ -437,13 +438,15 @@ async def alist_experiments(self, filter_conditions: Optional[ListExperimentUUID
         response = await self._client.request_async("POST", LIST_EXPERIMENTS_ENDPOINT, data=asdict(filter_conditions))
         return structure(response.json(), List[ExperimentWithPinnedStatsSchema])
 
-    def get_experiment_trace_logs(self, experiment_uuid: str, filters: TraceLogFilters = TraceLogFilters()) -> List[TraceLogTree]:
+    def get_experiment_trace_logs(self, experiment_uuid: str, filters: TraceLogFilters = TraceLogFilters(), return_children: bool = False) -> List[Union[TraceLogTree, TraceLog]]:
         response = self._client.request("POST", GET_EXPERIMENT_LOGS_ENDPOINT.format(experiment_uuid=experiment_uuid), data=asdict(filters))
-        return structure_trace_logs_from_api(response.json())
+        return structure_trace_logs_from_api(response.json(), return_children)
 
-    async def aget_experiment_trace_logs(self, experiment_uuid: str, filters: TraceLogFilters = TraceLogFilters()) -> List[TraceLogTree]:
+    async def aget_experiment_trace_logs(
+        self, experiment_uuid: str, filters: TraceLogFilters = TraceLogFilters(), return_children: bool = False
+    ) -> List[Union[TraceLogTree, TraceLog]]:
         response = await self._client.request_async("POST", GET_EXPERIMENT_LOGS_ENDPOINT.format(experiment_uuid=experiment_uuid), data=asdict(filters))
-        return structure_trace_logs_from_api(response.json())
+        return structure_trace_logs_from_api(response.json(), return_children)
 
     def get_experiment(self, experiment_uuid: str) -> Optional[ExperimentWithPinnedStatsSchema]:
         filter_conditions = ListExperimentUUIDsFilters(experiment_uuids=[experiment_uuid])
@@ -472,14 +475,15 @@ def new_init(self, *args, **kwargs):
     return subclass
 
 
-def extract_scores(tree: TraceLogTree) -> List[EvaluationResult]:
+def extract_scores(tree: Union[TraceLogTree, TraceLog]) -> List[EvaluationResult]:
     scores: List[EvaluationResult] = []
 
-    def traverse(node: TraceLogTree):
+    def traverse(node: Union[TraceLogTree, TraceLog]):
         if node.scores:
             scores.extend(node.scores or [])
-        for child in node.children_logs:
-            traverse(child)
+        if isinstance(node, TraceLogTree):
+            for child in node.children_logs or []:
+                traverse(child)
 
     traverse(tree)
     return scores

parea/helpers.py

@@ -81,7 +81,7 @@ def timezone_aware_now() -> datetime:
     return datetime.now(pytz.utc)
 
 
-def structure_trace_log_from_api(d: dict) -> TraceLogTree:
+def structure_trace_log_from_api(d: dict, return_children: bool = False) -> Union[TraceLogTree, TraceLog]:
     def structure_union_type(obj: Any, cl: type) -> Any:
         if isinstance(obj, str):
             return obj
@@ -92,11 +92,13 @@ def structure_union_type(obj: Any, cl: type) -> Any:
 
     converter = GenConverter()
     converter.register_structure_hook(Union[str, Dict[str, str], None], structure_union_type)
-    return converter.structure(d, TraceLogTree)
+    if return_children:
+        return converter.structure(d, TraceLogTree)
+    return converter.structure(d, TraceLog)
 
 
-def structure_trace_logs_from_api(data: List[dict]) -> List[TraceLogTree]:
-    return [structure_trace_log_from_api(d) for d in data]
+def structure_trace_logs_from_api(data: List[dict], return_children: bool = False) -> List[Union[TraceLogTree, TraceLog]]:
+    return [structure_trace_log_from_api(d, return_children) for d in data]
 
 
 PAREA_LOGGING_DISABLED = False