opik 1.6.4__py3-none-any.whl → 1.9.71__py3-none-any.whl

opik/file_upload/upload_client.py

@@ -0,0 +1,114 @@
+import dataclasses
+from typing import List, Optional, Iterable
+
+import httpx
+
+from ..rest_api import client as rest_api_client
+from ..rest_api import types as rest_api_types
+from ..rest_api.core import api_error as rest_api_error
+from ..rest_client_configurator import retry_decorator
+from . import file_upload_monitor
+from . import upload_options as file_upload_options
+
+LOCAL_UPLOAD_MAGIC_ID = "BEMinIO"
+
+
+@dataclasses.dataclass
+class MultipartUploadMetadata:
+    upload_id: str
+    urls: List[str]
+
+    def should_use_s3_uploader(self) -> bool:
+        """Allows to check if upload should go directly to S3 or use local backend endpoint."""
+        return self.upload_id is not None and self.upload_id != LOCAL_UPLOAD_MAGIC_ID
+
+
+class RestFileUploadClient:
+    """Defines the file upload REST API client wrapper that is used for communication with backend in order
+    to start and complete S3 file upload operation as well as multipart upload against local backend.
+    Args:
+        rest_client: The REST API client to communicate with the backend.
+    """
+
+    def __init__(
+        self, rest_client: rest_api_client.OpikApi, httpx_client: httpx.Client
+    ) -> None:
+        self.rest_client = rest_client
+        self.httpx_client = httpx_client
+
+    def start_upload(
+        self,
+        upload_options: file_upload_options.FileUploadOptions,
+        num_of_file_parts: int,
+        base_url_path: str,
+    ) -> MultipartUploadMetadata:
+        """Starts upload by sending request to the backend and receiving upload metadata. The upload metadata
+        will include the list of pre-signed URLs for direct S3 upload and upload ID assigned to this file
+        upload operation. If backend decides to force upload to then local endpoint then list of URLs will
+        include only one URL and upload ID will have magic value 'BEMinIO'."""
+        response = self.rest_client.attachments.start_multi_part_upload(
+            file_name=upload_options.file_name,
+            num_of_file_parts=num_of_file_parts,
+            entity_type=upload_options.entity_type,
+            entity_id=upload_options.entity_id,
+            path=base_url_path,
+            mime_type=upload_options.mime_type,
+            project_name=upload_options.project_name,
+        )
+        return MultipartUploadMetadata(
+            upload_id=response.upload_id, urls=response.pre_sign_urls
+        )
+
+    def s3_upload_completed(
+        self,
+        upload_options: file_upload_options.FileUploadOptions,
+        upload_metadata: MultipartUploadMetadata,
+        file_parts: List[rest_api_types.MultipartUploadPart],
+    ) -> None:
+        """Invoked to finalize direct S3 file upload operation on the backend. It is invoked after all file parts
+        was successfully uploaded to S3."""
+        self.rest_client.attachments.complete_multi_part_upload(
+            file_name=upload_options.file_name,
+            entity_type=upload_options.entity_type,
+            entity_id=upload_options.entity_id,
+            file_size=upload_options.file_size,
+            upload_id=upload_metadata.upload_id,
+            uploaded_file_parts=file_parts,
+            project_name=upload_options.project_name,
+            mime_type=upload_options.mime_type,
+        )
+
+    @retry_decorator.opik_rest_retry
+    def upload_file_local(
+        self,
+        upload_url: str,
+        file_path: str,
+        monitor: Optional[file_upload_monitor.FileUploadMonitor],
+        chunk_size: int = -1,
+    ) -> None:
+        """Invoked to upload a file to the local backend using httpx client configured with necessary authorization
+        headers. Raises the `HTTPStatusError` if one occurred."""
+        response = self.httpx_client.put(
+            url=upload_url,
+            content=_data_generator(file_path, chunk_size=chunk_size, monitor=monitor),
+        )
+        if response.status_code >= 400:
+            # reset monitor to avoid multiplication of sent bytes count on retries
+            if monitor is not None:
+                monitor.reset()
+            # raise rest_api_error.ApiError to trigger retry by our REST retry logic (see: retry_decorator.py)
+            raise rest_api_error.ApiError(
+                status_code=response.status_code, body=response.content
+            )
+
+
+def _data_generator(
+    file_path: str,
+    monitor: Optional[file_upload_monitor.FileUploadMonitor],
+    chunk_size: int = -1,
+) -> Iterable[bytes]:
+    with open(file_path, "rb") as file:
+        while data := file.read(chunk_size):
+            yield data
+            if monitor is not None:
+                monitor.update(len(data))

opik/file_upload/upload_manager.py

@@ -0,0 +1,255 @@
+import dataclasses
+import logging
+import math
+import time
+from concurrent.futures import Future, CancelledError
+from typing import Callable, List, Optional
+
+import httpx
+
+from . import (
+    base_upload_manager,
+    upload_options,
+    file_upload_monitor,
+    thread_pool,
+    file_uploader,
+)
+from .. import format_helpers, synchronization
+from ..message_processing import messages
+from ..rest_api import client as rest_api_client
+
+LOGGER = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class UploadResult:
+    future: Future
+    monitor: file_upload_monitor.FileUploadMonitor
+    upload_options: upload_options.FileUploadOptions
+
+    def ready(self) -> bool:
+        """Allows to check if wrapped Future successfully finished"""
+        return self.future.done()
+
+    def successful(self, timeout: Optional[float] = None) -> bool:
+        """Allows to check if wrapped Future completed without raising an exception"""
+        try:
+            exception = self.future.exception(timeout)
+            if exception is None:
+                return True
+
+            raise exception
+        except (CancelledError, TimeoutError) as e:
+            LOGGER.warning(
+                "Timeout while waiting for the result of file '%s' upload. Error: %s",
+                self.upload_options.file_name,
+                e,
+            )
+        except Exception as exception:
+            LOGGER.error(
+                "Failed to upload file with name '%s' from path [%s] with size [%s]. Error: %s",
+                self.upload_options.file_name,
+                self.upload_options.file_path,
+                format_helpers.format_bytes(self.upload_options.file_size),
+                exception,
+                exc_info=True,
+            )
+
+        return False
+
+
+class FileUploadManagerMonitor:
+    def __init__(
+        self, file_upload_manager: base_upload_manager.BaseFileUploadManager
+    ) -> None:
+        self.file_upload_manager = file_upload_manager
+        self.last_remaining_bytes = 0
+        self.last_remaining_uploads_display: Optional[float] = None
+
+    def log_remaining_uploads(self) -> None:
+        remaining = self.file_upload_manager.remaining_data()
+
+        current_time = time.monotonic()
+
+        if remaining.bytes == 0:
+            LOGGER.info(
+                "All assets have been sent, waiting for delivery confirmation",
+            )
+        elif self.last_remaining_uploads_display is None:
+            LOGGER.info(
+                "Still uploading %d file(s), remaining %s/%s",
+                remaining.uploads,
+                format_helpers.format_bytes(remaining.bytes),
+                format_helpers.format_bytes(remaining.total_size),
+            )
+        else:
+            uploaded_bytes = self.last_remaining_bytes - remaining.bytes
+            time_elapsed = current_time - self.last_remaining_uploads_display
+            upload_speed = uploaded_bytes / time_elapsed
+
+            # Avoid 0 division if no bytes were uploaded in the last period
+            if uploaded_bytes <= 0:
+                # avoid negative upload speed
+                if upload_speed < 0:
+                    upload_speed = 0
+
+                LOGGER.info(
+                    "Still uploading %d file(s), remaining %s/%s, Throughput %s/s, ETA unknown",
+                    remaining.uploads,
+                    format_helpers.format_bytes(remaining.bytes),
+                    format_helpers.format_bytes(remaining.total_size),
+                    format_helpers.format_bytes(upload_speed),
+                )
+
+            else:
+                remaining_time = str(int(math.ceil(remaining.bytes / upload_speed)))
+
+                LOGGER.info(
+                    "Still uploading %d file(s), remaining %s/%s, Throughput %s/s, ETA ~%ss",
+                    remaining.uploads,
+                    format_helpers.format_bytes(remaining.bytes),
+                    format_helpers.format_bytes(remaining.total_size),
+                    format_helpers.format_bytes(upload_speed),
+                    remaining_time,
+                )
+
+        self.last_remaining_bytes = remaining.bytes
+        self.last_remaining_uploads_display = current_time
+
+    def all_done(self) -> bool:
+        return self.file_upload_manager.all_done()
+
+
+class FileUploadManager(base_upload_manager.BaseFileUploadManager):
+    """Manages concurrent file uploads."""
+
+    def __init__(
+        self,
+        rest_client: rest_api_client.OpikApi,
+        httpx_client: httpx.Client,
+        worker_count: int,
+    ) -> None:
+        self._httpx_client = httpx_client
+        self._rest_client = rest_client
+
+        self._executor = thread_pool.get_thread_pool(worker_count=worker_count)
+        self._upload_results: List[UploadResult] = []
+        self.closed = False
+
+    def upload(self, message: messages.BaseMessage) -> None:
+        if isinstance(message, messages.CreateAttachmentMessage):
+            self.upload_attachment(message)
+        else:
+            raise ValueError(f"Message {message} is not supported for file upload.")
+
+    def upload_attachment(self, attachment: messages.CreateAttachmentMessage) -> None:
+        assert isinstance(
+            attachment, messages.CreateAttachmentMessage
+        ), "Wrong attachment message type"
+
+        options = upload_options.file_upload_options_from_attachment(attachment)
+        self._submit_upload(
+            uploader=file_uploader.upload_attachment,
+            options=options,
+        )
+
+    def _submit_upload(
+        self, options: upload_options.FileUploadOptions, uploader: Callable
+    ) -> None:
+        if self.closed:
+            LOGGER.warning(
+                "The file upload manager has been already closed. No more files can be submitted for upload. (%s)",
+                options.file_name,
+            )
+            return
+
+        monitor = file_upload_monitor.FileUploadMonitor()
+        if options.file_size > 0:
+            monitor.total_size = options.file_size
+
+        kwargs = {
+            "monitor": monitor,
+            "upload_options": options,
+            "rest_client": self._rest_client,
+            "upload_httpx_client": self._httpx_client,
+        }
+        future = self._executor.submit(uploader, **kwargs)
+        self._upload_results.append(
+            UploadResult(future, monitor=monitor, upload_options=options)
+        )
+
+    def all_done(self) -> bool:
+        return all(result.ready() for result in self._upload_results)
+
+    def remaining_data(self) -> base_upload_manager.RemainingUploadData:
+        remaining_uploads = 0
+        remaining_bytes_to_upload = 0
+        total_size = 0
+        for result in self._upload_results:
+            if result.ready() is True:
+                continue
+
+            remaining_uploads += 1
+
+            monitor = result.monitor
+            if monitor.total_size is None or monitor.bytes_sent is None:
+                continue
+
+            total_size += monitor.total_size
+            remaining_bytes_to_upload += monitor.total_size - monitor.bytes_sent
+
+        return base_upload_manager.RemainingUploadData(
+            uploads=remaining_uploads,
+            bytes=remaining_bytes_to_upload,
+            total_size=total_size,
+        )
+
+    def remaining_uploads(self) -> int:
+        """Returns the number of remaining uploads. Non-blocking - can be called at any time."""
+        status_list = [result.ready() for result in self._upload_results]
+        return status_list.count(False)
+
+    def failed_uploads(self, timeout: Optional[float]) -> int:
+        """Important - this is blocking method waiting for all remaining uploads to complete or while
+        timeout is expired."""
+        failed = 0
+        for result in self._upload_results:
+            if not result.ready() or not result.successful(timeout):
+                failed += 1
+
+        return failed
+
+    def flush(self, timeout: Optional[float], sleep_time: int = 5) -> bool:
+        """Flushes all pending uploads. This is a blocking method that waits for all remaining uploads to complete,
+        either until they finish or the specified timeout expires. If no timeout is set, it waits indefinitely.
+        Args:
+             timeout: Timeout in seconds to wait for all remaining uploads to complete.
+                If None is provided, it will wait for all remaining uploads to complete.
+            sleep_time: The sleep interval between checks and printing progress.
+        Returns:
+            The flag to indicate whether all remaining uploads are completed or not within the provided timeout.
+        """
+        upload_monitor = FileUploadManagerMonitor(self)
+
+        synchronization.wait_for_done(
+            check_function=lambda: self.all_done(),
+            progress_callback=upload_monitor.log_remaining_uploads,
+            timeout=timeout,
+            sleep_time=sleep_time,
+        )
+
+        # check failed uploads number only if all upload operations completed to avoid blocking
+        if self.all_done():
+            failed_uploads = self.failed_uploads(timeout)
+            if failed_uploads > 0:
+                LOGGER.warning(
+                    "Failed to upload %d file(s). Check logs for details.",
+                    failed_uploads,
+                )
+            return True
+
+        return False
+
+    def close(self) -> None:
+        self._executor.shutdown(wait=True)
+        self.closed = True

opik/file_upload/upload_options.py

@@ -0,0 +1,37 @@
+import dataclasses
+import os
+from typing import Optional
+
+from ..message_processing import messages
+from ..types import AttachmentEntityType
+
+
+@dataclasses.dataclass
+class FileUploadOptions:
+    file_path: str
+    file_name: str
+    file_size: int
+    mime_type: Optional[str]
+    entity_type: AttachmentEntityType
+    entity_id: str
+    project_name: str
+    encoded_url_override: str
+    delete_after_upload: bool
+
+
+def file_upload_options_from_attachment(
+    attachment: messages.CreateAttachmentMessage,
+) -> FileUploadOptions:
+    file_size = os.path.getsize(attachment.file_path)
+
+    return FileUploadOptions(
+        file_path=attachment.file_path,
+        file_name=attachment.file_name,
+        file_size=file_size,
+        mime_type=attachment.mime_type,
+        entity_type=attachment.entity_type,
+        entity_id=attachment.entity_id,
+        project_name=attachment.project_name,
+        encoded_url_override=attachment.encoded_url_override,
+        delete_after_upload=attachment.delete_after_upload,
+    )

opik/format_helpers.py

@@ -0,0 +1,17 @@
+ONE_KBYTE = float(1024)
+ONE_MBYTE = float(1024 * 1024)
+ONE_GBYTE = float(1024 * 1024 * 1024)
+
+
+def format_bytes(size: float) -> str:
+    """
+    Given a size in bytes, return a sort string representation.
+    """
+    if size >= ONE_GBYTE:
+        return "%.2f %s" % (size / ONE_GBYTE, "GB")
+    elif size >= ONE_MBYTE:
+        return "%.2f %s" % (size / ONE_MBYTE, "MB")
+    elif size >= ONE_KBYTE:
+        return "%.2f %s" % (size / ONE_KBYTE, "KB")
+    else:
+        return "%d %s" % (size, "bytes")

opik/guardrails/__init__.py

@@ -0,0 +1,4 @@
+from .guardrail import Guardrail
+from .guards import Topic, PII
+
+__all__ = ["Guardrail", "Topic", "PII"]

opik/guardrails/guardrail.py

@@ -0,0 +1,157 @@
+from typing import (
+    List,
+    Optional,
+)
+
+import httpx
+
+import opik.exceptions as exceptions
+import opik.config as config
+from opik.api_objects import opik_client
+from opik.message_processing.messages import (
+    GuardrailBatchItemMessage,
+    GuardrailBatchMessage,
+)
+from opik.opik_context import get_current_span_data, get_current_trace_data
+
+from . import guards, rest_api_client, schemas, tracing
+
+
+GUARDRAIL_DECORATOR = tracing.GuardrailsTrackDecorator()
+
+
+class Guardrail:
+    """
+    Client for the Opik Guardrails API.
+
+    This class provides a way to validate text against a set of guardrails.
+    """
+
+    def __init__(
+        self,
+        guards: List[guards.Guard],
+        guardrail_timeout: Optional[int] = None,
+    ) -> None:
+        """
+        Initialize a Guardrail client.
+
+        Args:
+            guards: List of Guard objects to validate text against
+
+        Example:
+
+        ```python
+        from opik.guardrails import Guardrail, PII, Topic
+        from opik import exceptions
+        guard = Guardrail(
+            guards=[
+                Topic(restricted_topics=["finance"], threshold=0.8),
+                PII(blocked_entities=["CREDIT_CARD", "PERSON"], threshold=0.4),
+            ]
+        )
+
+        result = guard.validate("How can I start with evaluation in Opik platform?")
+        # Guardrail passes
+
+        try:
+            result = guard.validate("Where should I invest my money?")
+        except exceptions.GuardrailValidationFailed as e:
+            print("Guardrail failed:", e)
+
+        try:
+            result = guard.validate("John Doe, here is my card number 4111111111111111 how can I use it in Opik platform?.")
+        except exceptions.GuardrailValidationFailed as e:
+            print("Guardrail failed:", e)
+        ```
+
+        """
+        self.guards = guards
+        self._client = opik_client.get_client_cached()
+
+        self.config_ = config.get_from_user_inputs(
+            guardrail_timeout=guardrail_timeout,
+        )
+
+        self._initialize_api_client(
+            host_url=self._client.config.guardrails_backend_host,
+        )
+
+    def _initialize_api_client(self, host_url: str) -> None:
+        self._api_client = rest_api_client.GuardrailsApiClient(
+            httpx_client=httpx.Client(timeout=self.config_.guardrail_timeout),
+            host_url=host_url,
+        )
+
+    def validate(self, text: str) -> schemas.ValidationResponse:
+        """
+        Validate text against all configured guardrails.
+
+        Args:
+            text: Text to validate
+
+        Returns:
+            ValidationResponse: API response containing validation results
+
+        Raises:
+            opik.exceptions.GuardrailValidationFailed: If validation fails
+            httpx.HTTPStatusError: If the API returns an error status code
+        """
+        result: schemas.ValidationResponse = self._validate(generation=text)  # type: ignore
+
+        return self._parse_result(result)
+
+    @GUARDRAIL_DECORATOR.track
+    def _validate(self, generation: str) -> schemas.ValidationResponse:
+        validations = []
+
+        for guard in self.guards:
+            validations.extend(guard.get_validation_configs())
+
+        result = self._api_client.validate(generation, validations)
+
+        if not result.validation_passed:
+            result.guardrail_result = "failed"
+        else:
+            result.guardrail_result = "passed"
+
+        batch = []
+
+        # Makes mypy happy that a current span and trace exists
+        current_span = get_current_span_data()
+        current_trace = get_current_trace_data()
+        assert current_span is not None
+        assert current_trace is not None
+
+        for validation in result.validations:
+            guardrail_batch_item_message = GuardrailBatchItemMessage(
+                project_name=self._client._project_name,
+                entity_id=current_trace.id,
+                secondary_id=current_span.id,
+                name=validation.type,
+                result="passed" if validation.validation_passed else "failed",
+                config=validation.validation_config,
+                details=validation.validation_details,
+            )
+            batch.append(guardrail_batch_item_message)
+
+        message = GuardrailBatchMessage(batch=batch)
+        self._client._streamer.put(message)
+
+        return result
+
+    def _parse_result(
+        self, result: schemas.ValidationResponse
+    ) -> schemas.ValidationResponse:
+        if not result.validation_passed:
+            failed_validations = []
+            for validation in result.validations:
+                if not validation.validation_passed:
+                    failed_validations.append(validation)
+
+            raise exceptions.GuardrailValidationFailed(
+                "Guardrail validation failed",
+                validation_results=result,
+                failed_validations=failed_validations,
+            )
+
+        return result

opik/guardrails/guards/__init__.py

@@ -0,0 +1,5 @@
+from .guard import Guard
+from .topic import Topic
+from .pii import PII
+
+__all__ = ["Guard", "Topic", "PII"]

opik/guardrails/guards/guard.py

@@ -0,0 +1,17 @@
+import abc
+from typing import List, Dict, Any
+
+
+class Guard(abc.ABC):
+    # Is this guard running locally or remotely on the Guardrail API
+    remote = False
+
+    @abc.abstractmethod
+    def get_validation_configs(self) -> List[Dict[str, Any]]:
+        """
+        Get the validation configuration for this guard.
+
+        Returns:
+            List of validation configurations to be sent to the API
+        """
+        pass

opik/guardrails/guards/pii.py

@@ -0,0 +1,47 @@
+from typing import List, Dict, Any, Optional
+from . import guard
+from .. import schemas
+import functools
+
+
+class PII(guard.Guard):
+    """
+    Guard that validates text for personally identifiable information (PII).
+    """
+
+    def __init__(
+        self,
+        blocked_entities: Optional[List[str]] = None,
+        language: str = "en",
+        threshold: float = 0.5,
+    ) -> None:
+        """
+        Initialize a PII guard.
+
+        Args:
+            blocked_entities: List of PII entity types to block.
+                Default entities include: "IP_ADDRESS", "PHONE_NUMBER", "PERSON", "MEDICAL_LICENSE", "URL", "EMAIL_ADDRESS", "IBAN_CODE".
+            Supported entities can be found here: https://microsoft.github.io/presidio/supported_entities/
+        """
+        self._blocked_entities = blocked_entities
+        self._language = language
+        self._threshold = threshold
+
+    @functools.lru_cache()
+    def get_validation_configs(self) -> List[Dict[str, Any]]:
+        """
+        Get the validation configuration for PII detection.
+
+        Returns:
+            List containing the PII validation configuration
+        """
+        return [
+            {
+                "type": schemas.ValidationType.PII,
+                "config": {
+                    "entities": self._blocked_entities,
+                    "language": self._language,
+                    "threshold": self._threshold,
+                },
+            }
+        ]