Merge pull request #2345 from Agenta-AI/fix-pycryptodome-vulnerability

Fix pycryptodome vulnerability
Agenta-AI · Dec 6, 2024 · 634c2d1 · 634c2d1
2 parents b31bb85 + f6a50dc
commit 634c2d1
Show file tree

Hide file tree

Showing 5 changed files with 625 additions and 660 deletions.
diff --git a/agenta-backend/agenta_backend/models/api/evaluation_model.py b/agenta-backend/agenta_backend/models/api/evaluation_model.py
@@ -2,7 +2,7 @@
 from datetime import datetime
 from typing import Optional, List, Dict, Any, Union
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, model_validator, field_validator
 
 from agenta_backend.utils import traces
 from agenta_backend.models.api.api_models import Result
@@ -293,9 +293,15 @@ class NewEvaluation(BaseModel):
     evaluators_configs: List[str]
     testset_id: str
     rate_limit: LLMRunRateLimit
-    lm_providers_keys: Optional[Dict[LMProvidersEnum, str]] = None
+    lm_providers_keys: Optional[Dict[str, str]] = None
     correct_answer_column: Optional[str] = None
 
+    @field_validator("lm_providers_keys", mode="after")
+    def validate_lm_providers_keys(cls, value):
+        if value is not None:
+            return {LMProvidersEnum(key): v for key, v in value.items()}
+        return value
+
 
 class NewEvaluatorConfig(BaseModel):
     app_id: str

diff --git a/agenta-backend/agenta_backend/services/llm_apps_service.py b/agenta-backend/agenta_backend/services/llm_apps_service.py
@@ -58,12 +58,24 @@ def extract_result_from_response(response: dict):
             if "tree" in response:
                 trace_tree = response.get("tree", {}).get("nodes", [])[0]
 
-                latency = (
-                    get_nested_value(
-                        trace_tree, ["metrics", "acc", "duration", "total"]
-                    )
-                    / 1000
+                duration_ms = get_nested_value(
+                    trace_tree, ["metrics", "acc", "duration", "total"]
                 )
+                if duration_ms:
+                    duration_seconds = duration_ms / 1000
+                else:
+                    start_time = get_nested_value(trace_tree, ["time", "start"])
+                    end_time = get_nested_value(trace_tree, ["time", "end"])
+
+                    if start_time and end_time:
+                        duration_seconds = (
+                            datetime.fromisoformat(end_time)
+                            - datetime.fromisoformat(start_time)
+                        ).total_seconds()
+                    else:
+                        duration_seconds = None
+
+                latency = duration_seconds
                 cost = get_nested_value(
                     trace_tree, ["metrics", "acc", "costs", "total"]
                 )

diff --git a/agenta-backend/agenta_backend/tasks/evaluations.py b/agenta-backend/agenta_backend/tasks/evaluations.py
@@ -283,6 +283,8 @@ def evaluate(
             ]
 
             # 4. We save the result of the eval scenario in the db
+            print("============ App Output ============: ", app_output.result.value)
+
             loop.run_until_complete(
                 create_new_evaluation_scenario(
                     project_id=project_id,