humalab 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

humalab/humalab_api_client.py

@@ -1,9 +1,28 @@
 """HTTP client for accessing HumaLab service APIs with API key authentication."""
 
+from enum import Enum
 import os
 import requests
 from typing import Dict, Any, Optional, List
 from urllib.parse import urljoin
+from humalab.humalab_config import HumalabConfig
+
+
+class RunStatus(Enum):
+    """Status of runs"""
+    RUNNING = "running"
+    CANCELED = "canceled"
+    ERRORED = "errored"
+    FINISHED = "finished"
+
+
+class EpisodeStatus(Enum):
+    """Status of validation episodes"""
+    RUNNING = "running"
+    CANCELED = "canceled"
+    ERRORED = "errored"
+    SUCCESS = "success"
+    FAILED = "failed"
 
 
 class HumaLabApiClient:
@@ -19,13 +38,14 @@ class HumaLabApiClient:
         Initialize the HumaLab API client.
         
         Args:
-            base_url: Base URL for the HumaLab service (defaults to localhost:8000)
+            base_url: Base URL for the HumaLab service (defaults to https://api.humalab.ai)
             api_key: API key for authentication (defaults to HUMALAB_API_KEY env var)
             timeout: Request timeout in seconds
         """
-        self.base_url = base_url or os.getenv("HUMALAB_SERVICE_URL", "http://localhost:8000")
-        self.api_key = api_key or os.getenv("HUMALAB_API_KEY")
-        self.timeout = timeout or 30.0  # Default timeout of 30 seconds
+        humalab_config = HumalabConfig()
+        self.base_url = base_url or humalab_config.base_url or os.getenv("HUMALAB_SERVICE_URL", "https://api.humalab.ai")
+        self.api_key = api_key or humalab_config.api_key or os.getenv("HUMALAB_API_KEY")
+        self.timeout = timeout or humalab_config.timeout or 30.0  # Default timeout of 30 seconds
         
         # Ensure base_url ends without trailing slash
         self.base_url = self.base_url.rstrip('/')
@@ -77,16 +97,35 @@ class HumaLabApiClient:
         if files:
             headers.pop("Content-Type", None)
         
-        response = requests.request(
-            method=method,
-            url=url,
-            json=data,
-            params=params,
-            files=files,
-            headers=headers,
-            timeout=self.timeout,
-            **kwargs
-        )
+        # Determine if we should send form data or JSON
+        # Form data endpoints: /artifacts/code, /artifacts/blob/upload, /artifacts/python
+        is_form_endpoint = any(form_path in endpoint for form_path in ['/artifacts/code', '/artifacts/blob', '/artifacts/python'])
+        
+        if is_form_endpoint or files:
+            # Send as form data
+            headers.pop("Content-Type", None)  # Let requests set multipart/form-data
+            response = requests.request(
+                method=method,
+                url=url,
+                data=data,
+                params=params,
+                files=files,
+                headers=headers,
+                timeout=self.timeout,
+                **kwargs
+            )
+        else:
+            # Send as JSON (default behavior)
+            response = requests.request(
+                method=method,
+                url=url,
+                json=data,
+                params=params,
+                files=files,
+                headers=headers,
+                timeout=self.timeout,
+                **kwargs
+            )
         
         # Raise an exception for HTTP error responses
         response.raise_for_status()
@@ -123,7 +162,7 @@ class HumaLabApiClient:
         Returns:
             User information from the validated token
         """
-        response = self.get("/api-key/validate")
+        response = self.get("/auth/validate")
         return response.json()
 
     # Convenience methods for common API operations
@@ -134,7 +173,7 @@ class HumaLabApiClient:
         resource_types: Optional[str] = None,
         limit: int = 20,
         offset: int = 0,
-        latest_only: bool = False
+        latest_only: bool = True
     ) -> Dict[str, Any]:
         """
         Get list of all resources.
@@ -258,8 +297,8 @@ class HumaLabApiClient:
     def get_scenarios(
         self,
         project_name: str,
-        skip: int = 0,
-        limit: int = 10,
+        limit: int = 20,
+        offset: int = 0,     
         include_inactive: bool = False,
         search: Optional[str] = None,
         status_filter: Optional[str] = None
@@ -269,8 +308,8 @@ class HumaLabApiClient:
 
         Args:
             project_name: Project name (required)
-            skip: Number of scenarios to skip for pagination
             limit: Maximum number of scenarios to return (1-100)
+            offset: Number of scenarios to skip
             include_inactive: Include inactive scenarios in results
             search: Search term to filter by name, description, or UUID
             status_filter: Filter by specific status
@@ -280,7 +319,7 @@ class HumaLabApiClient:
         """
         params = {
             "project_name": project_name,
-            "skip": skip,
+            "skip": offset,
             "limit": limit,
             "include_inactive": include_inactive
         }
@@ -441,7 +480,8 @@ class HumaLabApiClient:
             "name": name,
             "project_name": project_name,
             "arguments": arguments or [],
-            "tags": tags or []
+            "tags": tags or [],
+            "status": RunStatus.RUNNING.value
         }
         if description:
             data["description"] = description
@@ -452,7 +492,7 @@ class HumaLabApiClient:
     def get_runs(
         self,
         project_name: Optional[str],
-        status: Optional[str] = None,
+        status: Optional[RunStatus] = None,
         tags: Optional[List[str]] = None,
         limit: int = 20,
         offset: int = 0
@@ -475,7 +515,7 @@ class HumaLabApiClient:
             raise ValueError("project_name is required to get runs.")
         params["project_name"] = project_name
         if status:
-            params["status"] = status
+            params["status"] = status.value
         if tags:
             params["tags"] = ",".join(tags)
             
@@ -500,7 +540,8 @@ class HumaLabApiClient:
         run_id: str,
         name: Optional[str] = None,
         description: Optional[str] = None,
-        status: Optional[str] = None,
+        status: Optional[RunStatus] = None,
+        err_msg: Optional[str] = None,
         arguments: Optional[List[Dict[str, str]]] = None,
         tags: Optional[List[str]] = None
     ) -> Dict[str, Any]:
@@ -512,6 +553,7 @@ class HumaLabApiClient:
             name: Optional new name
             description: Optional new description
             status: Optional new status
+            err_msg: Optional error message
             arguments: Optional new arguments
             tags: Optional new tags
 
@@ -524,7 +566,9 @@ class HumaLabApiClient:
         if description is not None:
             data["description"] = description
         if status is not None:
-            data["status"] = status
+            data["status"] = status.value
+        if err_msg is not None:
+            data["err_msg"] = err_msg
         if arguments is not None:
             data["arguments"] = arguments
         if tags is not None:
@@ -536,26 +580,26 @@ class HumaLabApiClient:
     def create_episode(
         self, 
         run_id: str, 
-        episode_name: str,
-        status: Optional[str] = None
+        episode_id: str,
+        status: Optional[EpisodeStatus] = None
     ) -> Dict[str, Any]:
         """
         Create a new episode.
         
         Args:
             run_id: Run ID
-            episode_name: Episode name
+            episode_id: Episode name
             status: Optional episode status
             
         Returns:
             Created episode data
         """
         data = {
-            "episode_name": episode_name,
+            "episode_id": episode_id,
             "run_id": run_id
         }
         if status:
-            data["status"] = status
+            data["status"] = status.value
             
         response = self.post("/episodes", data=data)
         return response.json()
@@ -563,7 +607,7 @@ class HumaLabApiClient:
     def get_episodes(
         self,
         run_id: Optional[str] = None,
-        status: Optional[str] = None,
+        status: Optional[EpisodeStatus] = None,
         limit: int = 20,
         offset: int = 0
     ) -> Dict[str, Any]:
@@ -583,182 +627,340 @@ class HumaLabApiClient:
         if run_id:
             params["run_id"] = run_id
         if status:
-            params["status"] = status
+            params["status"] = status.value
             
         response = self.get("/episodes", params=params)
         return response.json()
     
-    def get_episode(self, run_id: str, episode_name: str) -> Dict[str, Any]:
+    def get_episode(self, run_id: str, episode_id: str) -> Dict[str, Any]:
         """
         Get a specific episode.
         
         Args:
             run_id: Run ID
-            episode_name: Episode name
+            episode_id: Episode name
             
         Returns:
             Episode data
         """
-        response = self.get(f"/episodes/{run_id}/{episode_name}")
+        response = self.get(f"/episodes/{run_id}/{episode_id}")
         return response.json()
     
     def update_episode(
         self,
         run_id: str,
-        episode_name: str,
-        status: Optional[str] = None
+        episode_id: str,
+        status: Optional[EpisodeStatus] = None,
+        err_msg: Optional[str] = None
     ) -> Dict[str, Any]:
         """
         Update an episode.
         
         Args:
             run_id: Run ID
-            episode_name: Episode name
+            episode_id: Episode name
             status: Optional new status
-            
+            err_msg: Optional error message
+
         Returns:
             Updated episode data
         """
         data = {}
         if status is not None:
-            data["status"] = status
-            
-        response = self.put(f"/episodes/{run_id}/{episode_name}", data=data)
+            data["status"] = status.value
+        if err_msg is not None:
+            data["err_msg"] = err_msg
+        response = self.put(f"/episodes/{run_id}/{episode_id}", data=data)
         return response.json()
     
-    def delete_episode(self, run_id: str, episode_name: str) -> None:
+    def delete_episode(self, run_id: str, episode_id: str) -> None:
         """
         Delete an episode.
         
         Args:
             run_id: Run ID
-            episode_name: Episode name
+            episode_id: Episode name
         """
-        self.delete(f"/episodes/{run_id}/{episode_name}")
+        self.delete(f"/episodes/{run_id}/{episode_id}")
     
     def upload_blob(
         self,
         artifact_key: str,
         run_id: str,
-        file_path: str,
         artifact_type: str,
-        episode_name: Optional[str] = None,
-        description: Optional[str] = None
+        file_content: bytes | None = None,
+        file_path: str | None = None,
+        episode_id: Optional[str] = None,
+        filename: Optional[str] = None,
+        content_type: Optional[str] = None
     ) -> Dict[str, Any]:
         """
         Upload a blob artifact (image/video).
-        
+
         Args:
             artifact_key: Artifact key identifier
             run_id: Run ID
-            file_path: Path to file to upload
             artifact_type: Type of artifact ('image' or 'video')
-            episode_name: Optional episode name (None for run-level artifacts)
-            description: Optional description
-            
+            file_content: File content as bytes
+            file_path: Path to file to upload
+            episode_id: Optional episode ID (None for run-level artifacts)
+            filename: Optional filename to use for the uploaded file
+            content_type: Optional content type (e.g., 'image/png', 'video/mp4')
+
         Returns:
             Created artifact data
         """
-        with open(file_path, 'rb') as f:
-            files = {'file': f}
-            form_data = {
-                'artifact_key': artifact_key,
-                'run_id': run_id,
-                'artifact_type': artifact_type
-            }
-            if episode_name:
-                form_data['episode_name'] = episode_name
-            if description:
-                form_data['description'] = description
-                
+        form_data = {
+            'artifact_key': artifact_key,
+            'run_id': run_id,
+            'artifact_type': artifact_type
+        }
+        if episode_id:
+            form_data['episode_id'] = episode_id
+        if filename:
+            form_data['filename'] = filename
+        if content_type:
+            form_data['content_type'] = content_type
+
+        if file_path:
+            with open(file_path, 'rb') as f:
+                files = {'file': f}
+                response = self.post("/artifacts/blob/upload", files=files, data=form_data)
+        elif file_content:
+            files = {'file': ('blob', file_content)}
             response = self.post("/artifacts/blob/upload", files=files, data=form_data)
-            return response.json()
+        else:
+            raise ValueError("Either file_path or file_content must be provided for blob upload.")
+        return response.json()
     
     def upsert_metrics(
         self,
         artifact_key: str,
         run_id: str,
         metric_type: str,
-        metric_data: List[Dict[str, Any]],
-        episode_name: Optional[str] = None,
-        description: Optional[str] = None
+        metric_data: Optional[List[Dict[str, Any]]] = None,
+        episode_id: Optional[str] = None
     ) -> Dict[str, Any]:
         """
         Upsert metrics artifact (create or append).
-        
+
         Args:
             artifact_key: Artifact key identifier
             run_id: Run ID
             metric_type: Type of metric display ('line', 'bar', 'scatter', 'gauge', 'counter')
             metric_data: List of metric data points with 'key', 'values', 'timestamp'
-            episode_name: Optional episode name (None for run-level artifacts)
-            description: Optional description
-            
+            episode_id: Optional episode ID (None for run-level artifacts)
+
         Returns:
             Created/updated artifact data
         """
         data = {
             "artifact_key": artifact_key,
             "run_id": run_id,
-            "metric_type": metric_type,
-            "metric_data": metric_data
+            "metric_type": metric_type
         }
-        if episode_name:
-            data["episode_name"] = episode_name
-        if description:
-            data["description"] = description
-            
+        if episode_id:
+            data["episode_id"] = episode_id
+        if metric_data:
+            data["metric_data"] = metric_data
+
         response = self.post("/artifacts/metrics", data=data)
         return response.json()
     
     def get_artifacts(
         self,
         run_id: Optional[str] = None,
-        episode_name: Optional[str] = None,
+        episode_id: Optional[str] = None,
         artifact_type: Optional[str] = None,
         limit: int = 20,
         offset: int = 0
     ) -> Dict[str, Any]:
         """
         Get list of artifacts.
-        
+
         Args:
             run_id: Filter by run ID
-            episode_name: Filter by episode name
+            episode_id: Filter by episode ID
             artifact_type: Filter by artifact type
-            limit: Maximum number of artifacts to return
+            limit: Maximum number of artifacts to return (0 for no limit)
             offset: Number of artifacts to skip
-            
+
         Returns:
             Artifact list with pagination info
         """
         params = {"limit": limit, "offset": offset}
         if run_id:
             params["run_id"] = run_id
-        if episode_name:
-            params["episode_name"] = episode_name
+        if episode_id:
+            params["episode_id"] = episode_id
         if artifact_type:
             params["artifact_type"] = artifact_type
-            
+
         response = self.get("/artifacts", params=params)
         return response.json()
     
     def get_artifact(
-        self, 
-        run_id: str, 
-        episode_name: str, 
+        self,
+        run_id: str,
+        episode_id: str,
         artifact_key: str
     ) -> Dict[str, Any]:
         """
         Get a specific artifact.
-        
+
         Args:
             run_id: Run ID
-            episode_name: Episode name
+            episode_id: Episode ID
             artifact_key: Artifact key
-            
+
         Returns:
             Artifact data
         """
-        response = self.get(f"/artifacts/{run_id}/{episode_name}/{artifact_key}")
+        response = self.get(f"/artifacts/{run_id}/{episode_id}/{artifact_key}")
+        return response.json()
+
+    def upload_code(
+        self,
+        artifact_key: str,
+        run_id: str,
+        code_content: str,
+        episode_id: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Upload code artifact (YAML/string content).
+
+        Args:
+            artifact_key: Artifact key identifier
+            run_id: Run ID
+            code_content: Code/text content to upload
+            episode_id: Optional episode ID (None for run-level artifacts)
+
+        Returns:
+            Created artifact data
+        """
+        data = {
+            'artifact_key': artifact_key,
+            'run_id': run_id,
+            'code_content': code_content
+        }
+        if episode_id:
+            data['episode_id'] = episode_id
+
+        response = self.post("/artifacts/code", data=data)
+        return response.json()
+
+    def upload_python(
+        self,
+        artifact_key: str,
+        run_id: str,
+        pickled_bytes: bytes,
+        episode_id: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Upload pickled Python object as artifact.
+
+        Args:
+            artifact_key: Artifact key identifier
+            run_id: Run ID
+            pickled_bytes: Pickled Python object as bytes
+            episode_id: Optional episode ID (None for run-level artifacts)
+
+        Returns:
+            Created artifact data
+        """
+        data = {
+            'artifact_key': artifact_key,
+            'run_id': run_id
+        }
+        if episode_id:
+            data['episode_id'] = episode_id
+
+        files = {'file': pickled_bytes}
+        response = self.post("/artifacts/python", files=files, data=data)
+        return response.json()
+
+    def upload_scenario_stats_artifact(
+        self,
+        artifact_key: str,
+        run_id: str,
+        pickled_bytes: bytes,
+        graph_type: str,
+    ) -> Dict[str, Any]:
+        """
+        Upload scenario stats artifact (pickled Python dict data).
+        This is an upsert operation - creates if doesn't exist, appends if it does.
+        Run-level only (no episode_id support).
+
+        Args:
+            artifact_key: Artifact key identifier
+            run_id: Run ID
+            pickled_bytes: Pickled Python dict as bytes containing scenario stats
+            graph_type: Graph display type - one of: 'line', 'bar', 'scatter',
+                       'histogram', 'gaussian', 'heatmap', '3d_map'
+
+        Returns:
+            Created/updated artifact data
+        """
+        data = {
+            'artifact_key': artifact_key,
+            'run_id': run_id,
+            'graph_type': graph_type
+        }
+
+        files = {'file': pickled_bytes}
+        response = self.post("/artifacts/scenario_stats", files=files, data=data)
+        return response.json()
+
+    def download_artifact(
+        self,
+        run_id: str,
+        episode_id: str,
+        artifact_key: str
+    ) -> bytes:
+        """
+        Download a blob artifact file.
+
+        Args:
+            run_id: Run ID
+            episode_id: Episode ID
+            artifact_key: Artifact key
+
+        Returns:
+            Artifact file content as bytes
+        """
+        endpoint = f"/artifacts/{run_id}/{episode_id}/{artifact_key}/download"
+        response = self.get(endpoint)
+        return response.content
+
+    def upload_metrics(
+        self,
+        run_id: str,
+        artifact_key: str,
+        pickled_bytes: bytes,
+        graph_type: str,
+        episode_id: str | None = None,
+    ) -> Dict[str, Any]:
+        """
+        Upload metrics artifact.
+
+        Args:
+            run_id: Run ID
+            artifact_key: Artifact key
+            pickled_bytes: Pickled metrics data as bytes
+            graph_type: Optional new graph type
+            episode_id: Optional new episode ID
+
+        Returns:
+            Updated artifact data
+        """
+        data = {
+            "run_id": run_id,
+            "artifact_key": artifact_key,
+            'graph_type': graph_type
+        }
+        files = {'file': pickled_bytes}
+        if episode_id:
+            data["episode_id"] = episode_id
+
+        response = self.post("/artifacts/metrics", files=files, data=data)
         return response.json()