judgeval 0.0.32__py3-none-any.whl → 0.0.33__py3-none-any.whl

judgeval/common/s3_storage.py

@@ -0,0 +1,93 @@
+import os
+import json
+import boto3
+from typing import Optional
+from datetime import datetime, UTC
+from botocore.exceptions import ClientError
+from judgeval.common.logger import warning, info
+
+class S3Storage:
+    """Utility class for storing and retrieving trace data from S3."""
+    
+    def __init__(
+        self,
+        bucket_name: str,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+        region_name: Optional[str] = None
+    ):
+        """Initialize S3 storage with credentials and bucket name.
+        
+        Args:
+            bucket_name: Name of the S3 bucket to store traces in
+            aws_access_key_id: AWS access key ID (optional, will use environment variables if not provided)
+            aws_secret_access_key: AWS secret access key (optional, will use environment variables if not provided)
+            region_name: AWS region name (optional, will use environment variables if not provided)
+        """
+        self.bucket_name = bucket_name
+        self.s3_client = boto3.client(
+            's3',
+            aws_access_key_id=aws_access_key_id or os.getenv('AWS_ACCESS_KEY_ID'),
+            aws_secret_access_key=aws_secret_access_key or os.getenv('AWS_SECRET_ACCESS_KEY'),
+            region_name=region_name or os.getenv('AWS_REGION', 'us-west-1')
+        )
+        
+    def _ensure_bucket_exists(self):
+        """Ensure the S3 bucket exists, creating it if necessary."""
+        try:
+            self.s3_client.head_bucket(Bucket=self.bucket_name)
+        except ClientError as e:
+            error_code = e.response['Error']['Code']
+            if error_code == '404':
+                # Bucket doesn't exist, create it
+                info(f"Bucket {self.bucket_name} doesn't exist, creating it ...")
+                try:
+                    self.s3_client.create_bucket(
+                        Bucket=self.bucket_name,
+                        CreateBucketConfiguration={
+                            'LocationConstraint': self.s3_client.meta.region_name
+                        }
+                    )
+                    info(f"Created S3 bucket: {self.bucket_name}")
+                except ClientError as create_error:
+                    if create_error.response['Error']['Code'] == 'BucketAlreadyOwnedByYou':
+                        # Bucket was just created by another process
+                        warning(f"Bucket {self.bucket_name} was just created by another process")
+                        pass
+                    else:
+                        raise create_error
+            else:
+                # Some other error occurred
+                raise e
+        
+    def save_trace(self, trace_data: dict, trace_id: str, project_name: str) -> str:
+        """Save trace data to S3.
+        
+        Args:
+            trace_data: The trace data to save
+            trace_id: Unique identifier for the trace
+            project_name: Name of the project the trace belongs to
+            
+        Returns:
+            str: S3 key where the trace was saved
+        """
+        # Ensure bucket exists before saving
+        self._ensure_bucket_exists()
+        
+        # Create a timestamped key for the trace
+        timestamp = datetime.now(UTC).strftime('%Y%m%d_%H%M%S')
+        s3_key = f"traces/{project_name}/{trace_id}_{timestamp}.json"
+        
+        # Convert trace data to JSON string
+        trace_json = json.dumps(trace_data)
+        
+        # Upload to S3
+        info(f"Uploading trace to S3 at key {s3_key}, in bucket {self.bucket_name} ...")
+        self.s3_client.put_object(
+            Bucket=self.bucket_name,
+            Key=s3_key,
+            Body=trace_json,
+            ContentType='application/json'
+        )
+        
+        return s3_key