arize-phoenix 3.16.1__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/session/session.py

@@ -1,26 +1,19 @@
 import json
 import logging
 import os
+import shutil
 import warnings
 from abc import ABC, abstractmethod
 from collections import UserList
+from collections.abc import Iterable, Mapping
 from datetime import datetime
 from enum import Enum
 from importlib.util import find_spec
+from itertools import chain
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from threading import Thread
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Iterable,
-    List,
-    Mapping,
-    NamedTuple,
-    Optional,
-    Set,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union
+from urllib.parse import urljoin
 
 import pandas as pd
 
@@ -29,26 +22,32 @@ from phoenix.config import (
     ENV_PHOENIX_COLLECTOR_ENDPOINT,
     ENV_PHOENIX_HOST,
     ENV_PHOENIX_PORT,
+    ensure_working_dir,
+    get_env_database_connection_str,
+    get_env_enable_websockets,
     get_env_host,
     get_env_port,
-    get_env_project_name,
     get_exported_files,
+    get_working_dir,
 )
-from phoenix.core.model_schema_adapter import create_model_from_datasets
-from phoenix.core.traces import Traces
-from phoenix.datasets.dataset import EMPTY_DATASET, Dataset
+from phoenix.core.model_schema_adapter import create_model_from_inferences
+from phoenix.inferences.inferences import EMPTY_INFERENCES, Inferences
 from phoenix.pointcloud.umap_parameters import get_umap_parameters
-from phoenix.server.app import create_app
+from phoenix.server.app import (
+    _db,
+    create_app,
+    create_engine_and_run_migrations,
+    instrument_engine_if_enabled,
+)
 from phoenix.server.thread_server import ThreadServer
+from phoenix.server.types import DbSessionFactory
 from phoenix.services import AppService
 from phoenix.session.client import Client
-from phoenix.session.data_extractor import TraceDataExtractor
+from phoenix.session.data_extractor import DEFAULT_SPAN_LIMIT, TraceDataExtractor
 from phoenix.session.evaluation import encode_evaluations
 from phoenix.trace import Evaluations
 from phoenix.trace.dsl.query import SpanQuery
 from phoenix.trace.trace_dataset import TraceDataset
-from phoenix.utilities import query_spans
-from phoenix.utilities.span_store import get_span_store, load_traces_data_from_store
 
 try:
     from IPython.display import IFrame  # type: ignore
@@ -64,6 +63,13 @@ if TYPE_CHECKING:
 else:
     _BaseList = UserList
 
+# Temporary directory for the duration of the session
+global _session_working_dir
+_session_working_dir: Optional["TemporaryDirectory[str]"] = None
+
+
+DEFAULT_TIMEOUT_IN_SECONDS = 5
+
 
 class NotebookEnvironment(Enum):
     COLAB = "colab"
@@ -74,8 +80,8 @@ class NotebookEnvironment(Enum):
 
 class ExportedData(_BaseList):
     def __init__(self) -> None:
-        self.paths: Set[Path] = set()
-        self.names: List[str] = []
+        self.paths: set[Path] = set()
+        self.names: list[str] = []
         super().__init__()
 
     def __repr__(self) -> str:
@@ -95,50 +101,30 @@ class Session(TraceDataExtractor, ABC):
     """Session that maintains a 1-1 shared state with the Phoenix App."""
 
     trace_dataset: Optional[TraceDataset]
-    traces: Optional[Traces]
     notebook_env: NotebookEnvironment
     """The notebook environment that the session is running in."""
 
-    def __dir__(self) -> List[str]:
+    def __dir__(self) -> list[str]:
         return ["exports", "view", "url"]
 
     def __init__(
         self,
-        primary_dataset: Dataset,
-        reference_dataset: Optional[Dataset] = None,
-        corpus_dataset: Optional[Dataset] = None,
+        database_url: str,
+        primary_inferences: Inferences,
+        reference_inferences: Optional[Inferences] = None,
+        corpus_inferences: Optional[Inferences] = None,
         trace_dataset: Optional[TraceDataset] = None,
         default_umap_parameters: Optional[Mapping[str, Any]] = None,
         host: Optional[str] = None,
         port: Optional[int] = None,
         notebook_env: Optional[NotebookEnvironment] = None,
     ):
-        self.primary_dataset = primary_dataset
-        self.reference_dataset = reference_dataset
-        self.corpus_dataset = corpus_dataset
+        self._database_url = database_url
+        self.primary_inferences = primary_inferences
+        self.reference_inferences = reference_inferences
+        self.corpus_inferences = corpus_inferences
         self.trace_dataset = trace_dataset
         self.umap_parameters = get_umap_parameters(default_umap_parameters)
-        self.model = create_model_from_datasets(
-            primary_dataset,
-            reference_dataset,
-        )
-
-        self.corpus = (
-            create_model_from_datasets(
-                corpus_dataset,
-            )
-            if corpus_dataset is not None
-            else None
-        )
-
-        self.traces = Traces()
-        if trace_dataset:
-            for span in trace_dataset.to_spans():
-                self.traces.put(span)
-            for evaluations in trace_dataset.evaluations:
-                for pb_evaluation in encode_evaluations(evaluations):
-                    self.traces.put(pb_evaluation)
-
         self.host = host or get_env_host()
         self.port = port or get_env_port()
         self.temp_dir = TemporaryDirectory()
@@ -147,6 +133,88 @@ class Session(TraceDataExtractor, ABC):
         self.exported_data = ExportedData()
         self.notebook_env = notebook_env or _get_notebook_environment()
         self.root_path = _get_root_path(self.notebook_env, self.port)
+        host = "127.0.0.1" if self.host == "0.0.0.0" else self.host
+        self._client = Client(
+            endpoint=f"http://{host}:{self.port}", warn_if_server_not_running=False
+        )
+
+    def query_spans(
+        self,
+        *queries: SpanQuery,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+        limit: Optional[int] = DEFAULT_SPAN_LIMIT,
+        root_spans_only: Optional[bool] = None,
+        project_name: Optional[str] = None,
+        # Deprecated fields
+        stop_time: Optional[datetime] = None,
+        timeout: Optional[int] = DEFAULT_TIMEOUT_IN_SECONDS,
+    ) -> Optional[Union[pd.DataFrame, list[pd.DataFrame]]]:
+        """
+        Queries the spans in the project based on the provided parameters.
+
+        Parameters
+        ----------
+            queries : *SpanQuery
+                Variable-length argument list of SpanQuery objects representing
+                the queries to be executed.
+
+            start_time : datetime, optional
+                 datetime representing the start time of the query.
+
+            end_time : datetime, optional
+                datetime representing the end time of the query.
+
+            root_spans_only : boolean, optional
+                whether to include only root spans in the results.
+
+            project_name : string, optional
+                name of the project to query. Defaults to the project name set
+                in the environment variable `PHOENIX_PROJECT_NAME` or 'default' if not set.
+
+        Returns:
+            results : DataFrame
+                DataFrame or list of DataFrames containing the query results.
+        """
+        if stop_time is not None:
+            warnings.warn(
+                "The `stop_time` parameter is deprecated and will be removed in a future release. "
+                "Please use `end_time` instead.",
+                DeprecationWarning,
+            )
+            end_time = stop_time
+        return self._client.query_spans(
+            *queries,
+            start_time=start_time,
+            end_time=end_time,
+            limit=limit,
+            root_spans_only=root_spans_only,
+            project_name=project_name,
+        )
+
+    def get_evaluations(
+        self,
+        project_name: Optional[str] = None,
+    ) -> list[Evaluations]:
+        """
+        Get the evaluations for a project.
+
+        Parameters
+        ----------
+            project_name :  str, optional
+                The name of the project. If not provided, the project name set
+                in the environment variable `PHOENIX_PROJECT_NAME` will be used.
+                Otherwise, 'default' will be used.
+
+        Returns
+        -------
+            evaluations : list[Evaluations]
+                A list of evaluations for the specified project.
+
+        """
+        return self._client.get_evaluations(
+            project_name=project_name,
+        )
 
     @abstractmethod
     def end(self) -> None:
@@ -170,23 +238,29 @@ class Session(TraceDataExtractor, ABC):
         self.exported_data.add(files)
         return self.exported_data
 
-    def view(self, height: int = 1000) -> "IFrame":
-        """
-        Returns an IFrame that can be displayed in a notebook to view the app.
+    def view(self, *, height: int = 1000, slug: str = "") -> "IFrame":
+        """View the session in a notebook embedded iFrame.
 
-        Parameters
-        ----------
-        height : int, optional
-            The height of the IFrame in pixels. Defaults to 1000.
+        Args:
+            slug (str, optional): the path of the app to view
+            height (int, optional): the height of the iFrame in px. Defaults to 1000.
+
+        Returns:
+            IFrame: the iFrame will be rendered in the notebook
         """
+        url_to_view = urljoin(self.url, slug)
         print(f"📺 Opening a view to the Phoenix app. The app is running at {self.url}")
-        return IFrame(src=self.url, width="100%", height=height)
+        return IFrame(src=url_to_view, width="100%", height=height)
 
     @property
     def url(self) -> str:
         """Returns the url for the phoenix app"""
         return _get_url(self.host, self.port, self.notebook_env)
 
+    @property
+    def database_url(self) -> str:
+        return self._database_url
+
 
 _session: Optional[Session] = None
 
@@ -194,9 +268,11 @@ _session: Optional[Session] = None
 class ProcessSession(Session):
     def __init__(
         self,
-        primary_dataset: Dataset,
-        reference_dataset: Optional[Dataset] = None,
-        corpus_dataset: Optional[Dataset] = None,
+        database_url: str,
+        primary_inferences: Inferences,
+        enable_websockets: bool,
+        reference_inferences: Optional[Inferences] = None,
+        corpus_inferences: Optional[Inferences] = None,
         trace_dataset: Optional[TraceDataset] = None,
         default_umap_parameters: Optional[Mapping[str, Any]] = None,
         host: Optional[str] = None,
@@ -205,20 +281,21 @@ class ProcessSession(Session):
         notebook_env: Optional[NotebookEnvironment] = None,
     ) -> None:
         super().__init__(
-            primary_dataset=primary_dataset,
-            reference_dataset=reference_dataset,
-            corpus_dataset=corpus_dataset,
+            database_url=database_url,
+            primary_inferences=primary_inferences,
+            reference_inferences=reference_inferences,
+            corpus_inferences=corpus_inferences,
             trace_dataset=trace_dataset,
             default_umap_parameters=default_umap_parameters,
             host=host,
             port=port,
             notebook_env=notebook_env,
         )
-        primary_dataset.to_disc()
-        if isinstance(reference_dataset, Dataset):
-            reference_dataset.to_disc()
-        if isinstance(corpus_dataset, Dataset):
-            corpus_dataset.to_disc()
+        primary_inferences.to_disc()
+        if isinstance(reference_inferences, Inferences):
+            reference_inferences.to_disc()
+        if isinstance(corpus_inferences, Inferences):
+            corpus_inferences.to_disc()
         if isinstance(trace_dataset, TraceDataset):
             trace_dataset.to_disc()
         umap_params_str = (
@@ -228,26 +305,23 @@ class ProcessSession(Session):
         )
         # Initialize an app service that keeps the server running
         self.app_service = AppService(
-            self.export_path,
-            self.host,
-            self.port,
-            self.root_path,
-            self.primary_dataset.name,
-            umap_params_str,
-            reference_dataset_name=(
-                self.reference_dataset.name if self.reference_dataset is not None else None
+            database_url=database_url,
+            export_path=self.export_path,
+            host=self.host,
+            port=self.port,
+            root_path=self.root_path,
+            primary_inferences_name=self.primary_inferences.name,
+            umap_params=umap_params_str,
+            reference_inferences_name=(
+                self.reference_inferences.name if self.reference_inferences is not None else None
             ),
-            corpus_dataset_name=(
-                self.corpus_dataset.name if self.corpus_dataset is not None else None
+            corpus_inferences_name=(
+                self.corpus_inferences.name if self.corpus_inferences is not None else None
             ),
             trace_dataset_name=(
                 self.trace_dataset.name if self.trace_dataset is not None else None
             ),
-        )
-        host = "127.0.0.1" if self.host == "0.0.0.0" else self.host
-        self._client = Client(
-            endpoint=f"http://{host}:{self.port}",
-            use_active_session_if_available=False,
+            enable_websockets=enable_websockets,
         )
 
     @property
@@ -258,35 +332,15 @@ class ProcessSession(Session):
         self.app_service.stop()
         self.temp_dir.cleanup()
 
-    def query_spans(
-        self,
-        *queries: SpanQuery,
-        start_time: Optional[datetime] = None,
-        stop_time: Optional[datetime] = None,
-        root_spans_only: Optional[bool] = None,
-        project_name: Optional[str] = None,
-    ) -> Optional[Union[pd.DataFrame, List[pd.DataFrame]]]:
-        return self._client.query_spans(
-            *queries,
-            start_time=start_time,
-            stop_time=stop_time,
-            root_spans_only=root_spans_only,
-            project_name=project_name,
-        )
-
-    def get_evaluations(
-        self,
-        project_name: Optional[str] = None,
-    ) -> List[Evaluations]:
-        return self._client.get_evaluations()
-
 
 class ThreadSession(Session):
     def __init__(
         self,
-        primary_dataset: Dataset,
-        reference_dataset: Optional[Dataset] = None,
-        corpus_dataset: Optional[Dataset] = None,
+        database_url: str,
+        primary_inferences: Inferences,
+        enable_websockets: bool,
+        reference_inferences: Optional[Inferences] = None,
+        corpus_inferences: Optional[Inferences] = None,
         trace_dataset: Optional[TraceDataset] = None,
         default_umap_parameters: Optional[Mapping[str, Any]] = None,
         host: Optional[str] = None,
@@ -295,29 +349,46 @@ class ThreadSession(Session):
         notebook_env: Optional[NotebookEnvironment] = None,
     ):
         super().__init__(
-            primary_dataset=primary_dataset,
-            reference_dataset=reference_dataset,
-            corpus_dataset=corpus_dataset,
+            database_url=database_url,
+            primary_inferences=primary_inferences,
+            reference_inferences=reference_inferences,
+            corpus_inferences=corpus_inferences,
             trace_dataset=trace_dataset,
             default_umap_parameters=default_umap_parameters,
             host=host,
             port=port,
             notebook_env=notebook_env,
         )
-        if span_store := get_span_store():
-            Thread(
-                target=load_traces_data_from_store,
-                args=(self.traces, span_store),
-                daemon=True,
-            ).start()
+        self.model = create_model_from_inferences(
+            primary_inferences,
+            reference_inferences,
+        )
+        self.corpus = (
+            create_model_from_inferences(
+                corpus_inferences,
+            )
+            if corpus_inferences is not None
+            else None
+        )
         # Initialize an app service that keeps the server running
+        engine = create_engine_and_run_migrations(database_url)
+        instrumentation_cleanups = instrument_engine_if_enabled(engine)
+        factory = DbSessionFactory(db=_db(engine), dialect=engine.dialect.name)
         self.app = create_app(
+            db=factory,
             export_path=self.export_path,
             model=self.model,
+            authentication_enabled=False,
+            enable_websockets=enable_websockets,
             corpus=self.corpus,
-            traces=self.traces,
             umap_params=self.umap_parameters,
-            span_store=span_store,
+            initial_spans=trace_dataset.to_spans() if trace_dataset else None,
+            initial_evaluations=(
+                chain.from_iterable(map(encode_evaluations, initial_evaluations))
+                if (trace_dataset and (initial_evaluations := trace_dataset.evaluations))
+                else None
+            ),
+            shutdown_callbacks=instrumentation_cleanups,
         )
         self.server = ThreadServer(
             app=self.app,
@@ -336,103 +407,43 @@ class ThreadSession(Session):
         self.server.close()
         self.temp_dir.cleanup()
 
-    def query_spans(
-        self,
-        *queries: SpanQuery,
-        start_time: Optional[datetime] = None,
-        stop_time: Optional[datetime] = None,
-        root_spans_only: Optional[bool] = None,
-        project_name: Optional[str] = None,
-    ) -> Optional[Union[pd.DataFrame, List[pd.DataFrame]]]:
-        """
-        Queries the spans in the project based on the provided parameters.
-
-        Parameters
-        ----------
-            queries : *SpanQuery
-                Variable-length argument list of SpanQuery objects representing
-                the queries to be executed.
-
-            start_time : datetime, optional
-                 datetime representing the start time of the query.
-
-            stop_time : datetime, optional
-                datetime representing the stop time of the query.
 
-            root_spans_only : boolean, optional
-                whether to include only root spans in the results.
-
-            project_name : string, optional
-                name of the project to query. Defaults to the project name set
-                in the environment variable `PHOENIX_PROJECT_NAME` or 'default' if not set.
-
-        Returns:
-            results : DataFrame
-                DataFrame or list of DataFrames containing the query results.
-        """
-        if not (traces := self.traces) or not (
-            project := traces.get_project(project_name or get_env_project_name())
-        ):
-            return None
-        if not queries:
-            queries = (SpanQuery(),)
-        valid_eval_names = project.get_span_evaluation_names() if project else ()
-        queries = tuple(
-            SpanQuery.from_dict(
-                query.to_dict(),
-                evals=project,
-                valid_eval_names=valid_eval_names,
+def delete_all(prompt_before_delete: Optional[bool] = True) -> None:
+    """
+    Deletes the entire contents of the working directory. This will delete, traces, evaluations,
+    and any other data stored in the working directory.
+    """
+    global _session_working_dir
+    working_dir = get_working_dir()
+    directories_to_delete = []
+    if working_dir.exists():
+        directories_to_delete.append(working_dir)
+    if _session_working_dir is not None:
+        directories_to_delete.append(Path(_session_working_dir.name))
+
+    # Loop through directories to delete
+    for directory in directories_to_delete:
+        if prompt_before_delete:
+            input(
+                f"You have data at {directory}. Are you sure you want to delete?"
+                + " This cannot be undone. Press Enter to delete, Escape to cancel."
             )
-            for query in queries
-        )
-        results = query_spans(
-            project,
-            *queries,
-            start_time=start_time,
-            stop_time=stop_time,
-            root_spans_only=root_spans_only,
-        )
-        if len(results) == 1:
-            df = results[0]
-            return None if df.shape == (0, 0) else df
-        return results
-
-    def get_evaluations(
-        self,
-        project_name: Optional[str] = None,
-    ) -> List[Evaluations]:
-        """
-        Get the evaluations for a project.
-
-        Parameters
-        ----------
-            project_name :  str, optional
-                The name of the project. If not provided, the project name set
-                in the environment variable `PHOENIX_PROJECT_NAME` will be used.
-                Otherwise, 'default' will be used.
-
-        Returns
-        -------
-            evaluations : List[Evaluations]
-                A list of evaluations for the specified project.
-
-        """
-        project_name = project_name or get_env_project_name()
-        if not (traces := self.traces) or not (project := traces.get_project(project_name)):
-            return []
-        return project.export_evaluations()
+        shutil.rmtree(directory)
+    _session_working_dir = None
 
 
 def launch_app(
-    primary: Optional[Dataset] = None,
-    reference: Optional[Dataset] = None,
-    corpus: Optional[Dataset] = None,
+    primary: Optional[Inferences] = None,
+    reference: Optional[Inferences] = None,
+    corpus: Optional[Inferences] = None,
     trace: Optional[TraceDataset] = None,
     default_umap_parameters: Optional[Mapping[str, Any]] = None,
     host: Optional[str] = None,
     port: Optional[int] = None,
     run_in_thread: bool = True,
     notebook_environment: Optional[Union[NotebookEnvironment, str]] = None,
+    use_temp_dir: bool = True,
+    enable_websockets: Optional[bool] = None,
 ) -> Optional[Session]:
     """
     Launches the phoenix application and returns a session to interact with.
@@ -447,7 +458,7 @@ def launch_app(
     corpus : Dataset, optional
         The dataset containing corpus for LLM context retrieval.
     trace: TraceDataset, optional
-        **Experimental** The trace dataset containing the trace data.
+        The trace dataset containing the trace data.
     host: str, optional
         The host on which the server runs. It can also be set using environment
         variable `PHOENIX_HOST`, otherwise it defaults to `127.0.0.1`.
@@ -457,13 +468,18 @@ def launch_app(
         Defaults to 6006.
     run_in_thread: bool, optional, default=True
         Whether the server should run in a Thread or Process.
-    default_umap_parameters: Dict[str, Union[int, float]], optional, default=None
+    default_umap_parameters: dict[str, Union[int, float]], optional, default=None
         User specified default UMAP parameters
         eg: {"n_neighbors": 10, "n_samples": 5, "min_dist": 0.5}
     notebook_environment: str, optional, default=None
         The environment the notebook is running in. This is either 'local', 'colab', or 'sagemaker'.
         If not provided, phoenix will try to infer the environment. This is only needed if
         there is a failure to infer the environment.
+    use_temp_dir: bool, optional, default=True
+        Whether to use a temporary directory to store the data. If set to False, the data will be
+        stored in the directory specified by PHOENIX_WORKING_DIR environment variable via SQLite.
+    enable_websockets: bool, optional, default=False
+        Whether to enable websockets.
 
     Returns
     -------
@@ -473,17 +489,21 @@ def launch_app(
     Examples
     --------
     >>> import phoenix as px
-    >>> # construct a dataset to analyze
-    >>> dataset = px.Dataset(...)
-    >>> session = px.launch_app(dataset)
+    >>> # construct an inference set to analyze
+    >>> inferences = px.Inferences(...)
+    >>> session = px.launch_app(inferences)
     """
     global _session
 
+    # First we must ensure that the working directory is setup
+    # NB: this is because the working directory can be deleted by the user
+    ensure_working_dir()
+
     # Stopgap solution to allow the app to run without a primary dataset
     if primary is None:
-        # Dummy dataset
-        # TODO: pass through the lack of a primary dataset to the app
-        primary = EMPTY_DATASET
+        # Dummy inferences
+        # TODO: pass through the lack of a primary inferences to the app
+        primary = EMPTY_INFERENCES
 
     if _session is not None and _session.active:
         logger.warning(
@@ -533,10 +553,23 @@ def launch_app(
 
     host = host or get_env_host()
     port = port or get_env_port()
+    if use_temp_dir:
+        global _session_working_dir
+        _session_working_dir = _session_working_dir or TemporaryDirectory()
+        database_url = f"sqlite:///{_session_working_dir.name}/phoenix.db"
+    else:
+        database_url = get_env_database_connection_str()
+
+    enable_websockets_env = get_env_enable_websockets() or False
+    enable_websockets = (
+        enable_websockets if enable_websockets is not None else enable_websockets_env
+    )
 
     if run_in_thread:
         _session = ThreadSession(
+            database_url,
             primary,
+            enable_websockets,
             reference,
             corpus,
             trace,
@@ -548,7 +581,9 @@ def launch_app(
         # TODO: catch exceptions from thread
     else:
         _session = ProcessSession(
+            database_url,
             primary,
+            enable_websockets,
             reference,
             corpus,
             trace,
@@ -568,7 +603,8 @@ def launch_app(
         return None
 
     print(f"🌍 To view the Phoenix app in your browser, visit {_session.url}")
-    print("📺 To view the Phoenix app in a notebook, run `px.active_session().view()`")
+    if not use_temp_dir:
+        print(f"💽 Your data is being persisted to {database_url}")
     print("📖 For more information on how to use Phoenix, check out https://docs.arize.com/phoenix")
     return _session
 
@@ -582,10 +618,15 @@ def active_session() -> Optional[Session]:
     return None
 
 
-def close_app() -> None:
+def close_app(delete_data: bool = False) -> None:
     """
     Closes the phoenix application.
     The application server is shut down and will no longer be accessible.
+
+    Parameters
+    ----------
+    delete_data : bool, optional
+        If set to true, all stored phoenix data, including traces and evaluations. Default False.
     """
     global _session
     if _session is None:
@@ -594,12 +635,15 @@ def close_app() -> None:
     _session.end()
     _session = None
     logger.info("Session closed")
+    if delete_data:
+        logger.info("Deleting all data")
+        delete_all(prompt_before_delete=False)
 
 
 def _get_url(host: str, port: int, notebook_env: NotebookEnvironment) -> str:
     """Determines the IFrame URL based on whether this is in a Colab or in a local notebook"""
     if notebook_env == NotebookEnvironment.COLAB:
-        from google.colab.output import eval_js  # type: ignore
+        from google.colab.output import eval_js
 
         return str(eval_js(f"google.colab.kernel.proxyPort({port}, {{'cache': true}})"))
     if notebook_env == NotebookEnvironment.SAGEMAKER:
@@ -617,7 +661,7 @@ def _get_url(host: str, port: int, notebook_env: NotebookEnvironment) -> str:
 def _is_colab() -> bool:
     """Determines whether this is in a Colab"""
     try:
-        import google.colab  # type: ignore # noqa: F401
+        import google.colab  # noqa: F401
     except ImportError:
         return False
     try: