arize-phoenix 12.0.0__py3-none-any.whl → 12.2.0__py3-none-any.whl

phoenix/server/api/mutations/dataset_label_mutations.py

@@ -0,0 +1,291 @@
+from typing import Optional
+
+import sqlalchemy
+import strawberry
+from sqlalchemy import delete, select
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from strawberry import UNSET
+from strawberry.relay.types import GlobalID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.auth import IsLocked, IsNotReadOnly
+from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
+from phoenix.server.api.queries import Query
+from phoenix.server.api.types.Dataset import Dataset
+from phoenix.server.api.types.DatasetLabel import DatasetLabel, to_gql_dataset_label
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+
+
+@strawberry.input
+class CreateDatasetLabelInput:
+    name: str
+    description: Optional[str] = UNSET
+    color: str
+
+
+@strawberry.type
+class CreateDatasetLabelMutationPayload:
+    dataset_label: DatasetLabel
+
+
+@strawberry.input
+class DeleteDatasetLabelsInput:
+    dataset_label_ids: list[GlobalID]
+
+
+@strawberry.type
+class DeleteDatasetLabelsMutationPayload:
+    dataset_labels: list[DatasetLabel]
+
+
+@strawberry.input
+class UpdateDatasetLabelInput:
+    dataset_label_id: GlobalID
+    name: str
+    description: Optional[str] = None
+    color: str
+
+
+@strawberry.type
+class UpdateDatasetLabelMutationPayload:
+    dataset_label: DatasetLabel
+
+
+@strawberry.input
+class SetDatasetLabelsInput:
+    dataset_label_ids: list[GlobalID]
+    dataset_ids: list[GlobalID]
+
+
+@strawberry.type
+class SetDatasetLabelsMutationPayload:
+    query: "Query"
+
+
+@strawberry.input
+class UnsetDatasetLabelsInput:
+    dataset_label_ids: list[GlobalID]
+    dataset_ids: list[GlobalID]
+
+
+@strawberry.type
+class UnsetDatasetLabelsMutationPayload:
+    query: "Query"
+
+
+@strawberry.type
+class DatasetLabelMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def create_dataset_label(
+        self,
+        info: Info[Context, None],
+        input: CreateDatasetLabelInput,
+    ) -> CreateDatasetLabelMutationPayload:
+        name = input.name
+        description = input.description
+        color = input.color
+        async with info.context.db() as session:
+            dataset_label_orm = models.DatasetLabel(name=name, description=description, color=color)
+            session.add(dataset_label_orm)
+            try:
+                await session.commit()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"A dataset label named '{name}' already exists")
+            except sqlalchemy.exc.StatementError as error:
+                raise BadRequest(str(error.orig))
+        return CreateDatasetLabelMutationPayload(
+            dataset_label=to_gql_dataset_label(dataset_label_orm)
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def update_dataset_label(
+        self, info: Info[Context, None], input: UpdateDatasetLabelInput
+    ) -> UpdateDatasetLabelMutationPayload:
+        if not input.name or not input.name.strip():
+            raise BadRequest("Dataset label name cannot be empty")
+
+        try:
+            dataset_label_id = from_global_id_with_expected_type(
+                input.dataset_label_id, DatasetLabel.__name__
+            )
+        except ValueError:
+            raise BadRequest(f"Invalid dataset label ID: {input.dataset_label_id}")
+
+        async with info.context.db() as session:
+            dataset_label_orm = await session.get(models.DatasetLabel, dataset_label_id)
+            if not dataset_label_orm:
+                raise NotFound(f"DatasetLabel with ID {input.dataset_label_id} not found")
+
+            dataset_label_orm.name = input.name.strip()
+            dataset_label_orm.description = input.description
+            dataset_label_orm.color = input.color.strip()
+
+            try:
+                await session.commit()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"A dataset label named '{input.name}' already exists")
+            except sqlalchemy.exc.StatementError as error:
+                raise BadRequest(str(error.orig))
+        return UpdateDatasetLabelMutationPayload(
+            dataset_label=to_gql_dataset_label(dataset_label_orm)
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def delete_dataset_labels(
+        self, info: Info[Context, None], input: DeleteDatasetLabelsInput
+    ) -> DeleteDatasetLabelsMutationPayload:
+        dataset_label_row_ids: dict[int, None] = {}
+        for dataset_label_node_id in input.dataset_label_ids:
+            try:
+                dataset_label_row_id = from_global_id_with_expected_type(
+                    dataset_label_node_id, DatasetLabel.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Unknown dataset label: {dataset_label_node_id}")
+            dataset_label_row_ids[dataset_label_row_id] = None
+        async with info.context.db() as session:
+            stmt = (
+                delete(models.DatasetLabel)
+                .where(models.DatasetLabel.id.in_(dataset_label_row_ids.keys()))
+                .returning(models.DatasetLabel)
+            )
+            deleted_dataset_labels = (await session.scalars(stmt)).all()
+            if len(deleted_dataset_labels) < len(dataset_label_row_ids):
+                await session.rollback()
+                raise NotFound("Could not find one or more dataset labels with given IDs")
+        deleted_dataset_labels_by_id = {
+            dataset_label.id: dataset_label for dataset_label in deleted_dataset_labels
+        }
+        return DeleteDatasetLabelsMutationPayload(
+            dataset_labels=[
+                to_gql_dataset_label(deleted_dataset_labels_by_id[dataset_label_row_id])
+                for dataset_label_row_id in dataset_label_row_ids
+            ]
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def set_dataset_labels(
+        self, info: Info[Context, None], input: SetDatasetLabelsInput
+    ) -> SetDatasetLabelsMutationPayload:
+        if not input.dataset_ids:
+            raise BadRequest("No datasets provided.")
+        if not input.dataset_label_ids:
+            raise BadRequest("No dataset labels provided.")
+
+        unique_dataset_rowids: set[int] = set()
+        for dataset_gid in input.dataset_ids:
+            try:
+                dataset_rowid = from_global_id_with_expected_type(dataset_gid, Dataset.__name__)
+            except ValueError:
+                raise BadRequest(f"Invalid dataset ID: {dataset_gid}")
+            unique_dataset_rowids.add(dataset_rowid)
+        dataset_rowids = list(unique_dataset_rowids)
+
+        unique_dataset_label_rowids: set[int] = set()
+        for dataset_label_gid in input.dataset_label_ids:
+            try:
+                dataset_label_rowid = from_global_id_with_expected_type(
+                    dataset_label_gid, DatasetLabel.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Invalid dataset label ID: {dataset_label_gid}")
+            unique_dataset_label_rowids.add(dataset_label_rowid)
+        dataset_label_rowids = list(unique_dataset_label_rowids)
+
+        async with info.context.db() as session:
+            existing_dataset_ids = (
+                await session.scalars(
+                    select(models.Dataset.id).where(models.Dataset.id.in_(dataset_rowids))
+                )
+            ).all()
+            if len(existing_dataset_ids) != len(dataset_rowids):
+                raise NotFound("One or more datasets not found")
+
+            existing_dataset_label_ids = (
+                await session.scalars(
+                    select(models.DatasetLabel.id).where(
+                        models.DatasetLabel.id.in_(dataset_label_rowids)
+                    )
+                )
+            ).all()
+            if len(existing_dataset_label_ids) != len(dataset_label_rowids):
+                raise NotFound("One or more dataset labels not found")
+
+            existing_dataset_label_keys = await session.execute(
+                select(
+                    models.DatasetsDatasetLabel.dataset_id,
+                    models.DatasetsDatasetLabel.dataset_label_id,
+                ).where(
+                    models.DatasetsDatasetLabel.dataset_id.in_(dataset_rowids)
+                    & models.DatasetsDatasetLabel.dataset_label_id.in_(dataset_label_rowids)
+                )
+            )
+            unique_dataset_label_keys = set(existing_dataset_label_keys.all())
+
+            datasets_dataset_labels = []
+            for dataset_rowid in dataset_rowids:
+                for dataset_label_rowid in dataset_label_rowids:
+                    if (dataset_rowid, dataset_label_rowid) in unique_dataset_label_keys:
+                        continue
+                    datasets_dataset_labels.append(
+                        models.DatasetsDatasetLabel(
+                            dataset_id=dataset_rowid,
+                            dataset_label_id=dataset_label_rowid,
+                        )
+                    )
+            session.add_all(datasets_dataset_labels)
+
+            if datasets_dataset_labels:
+                try:
+                    await session.commit()
+                except (PostgreSQLIntegrityError, SQLiteIntegrityError) as e:
+                    raise Conflict("Failed to add dataset labels to datasets.") from e
+
+        return SetDatasetLabelsMutationPayload(
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def unset_dataset_labels(
+        self, info: Info[Context, None], input: UnsetDatasetLabelsInput
+    ) -> UnsetDatasetLabelsMutationPayload:
+        if not input.dataset_ids:
+            raise BadRequest("No datasets provided.")
+        if not input.dataset_label_ids:
+            raise BadRequest("No dataset labels provided.")
+
+        unique_dataset_rowids: set[int] = set()
+        for dataset_gid in input.dataset_ids:
+            try:
+                dataset_rowid = from_global_id_with_expected_type(dataset_gid, Dataset.__name__)
+            except ValueError:
+                raise BadRequest(f"Invalid dataset ID: {dataset_gid}")
+            unique_dataset_rowids.add(dataset_rowid)
+        dataset_rowids = list(unique_dataset_rowids)
+
+        unique_dataset_label_rowids: set[int] = set()
+        for dataset_label_gid in input.dataset_label_ids:
+            try:
+                dataset_label_rowid = from_global_id_with_expected_type(
+                    dataset_label_gid, DatasetLabel.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Invalid dataset label ID: {dataset_label_gid}")
+            unique_dataset_label_rowids.add(dataset_label_rowid)
+        dataset_label_rowids = list(unique_dataset_label_rowids)
+
+        async with info.context.db() as session:
+            await session.execute(
+                delete(models.DatasetsDatasetLabel).where(
+                    models.DatasetsDatasetLabel.dataset_id.in_(dataset_rowids)
+                    & models.DatasetsDatasetLabel.dataset_label_id.in_(dataset_label_rowids)
+                )
+            )
+            await session.commit()
+
+        return UnsetDatasetLabelsMutationPayload(
+            query=Query(),
+        )

phoenix/server/api/mutations/dataset_split_mutations.py

@@ -15,6 +15,7 @@ from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
 from phoenix.server.api.helpers.playground_users import get_user
 from phoenix.server.api.queries import Query
+from phoenix.server.api.types.DatasetExample import DatasetExample, to_gql_dataset_example
 from phoenix.server.api.types.DatasetSplit import DatasetSplit, to_gql_dataset_split
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 
@@ -68,6 +69,13 @@ class DatasetSplitMutationPayload:
     query: "Query"
 
 
+@strawberry.type
+class DatasetSplitMutationPayloadWithExamples:
+    dataset_split: DatasetSplit
+    query: "Query"
+    examples: list[DatasetExample]
+
+
 @strawberry.type
 class DeleteDatasetSplitsMutationPayload:
     dataset_splits: list[DatasetSplit]
@@ -77,11 +85,13 @@ class DeleteDatasetSplitsMutationPayload:
 @strawberry.type
 class AddDatasetExamplesToDatasetSplitsMutationPayload:
     query: "Query"
+    examples: list[DatasetExample]
 
 
 @strawberry.type
 class RemoveDatasetExamplesFromDatasetSplitsMutationPayload:
     query: "Query"
+    examples: list[DatasetExample]
 
 
 @strawberry.type
@@ -262,8 +272,16 @@ class DatasetSplitMutationMixin:
                 except (PostgreSQLIntegrityError, SQLiteIntegrityError) as e:
                     raise Conflict("Failed to add examples to dataset splits.") from e
 
+            examples = (
+                await session.scalars(
+                    select(models.DatasetExample).where(
+                        models.DatasetExample.id.in_(example_rowids)
+                    )
+                )
+            ).all()
         return AddDatasetExamplesToDatasetSplitsMutationPayload(
             query=Query(),
+            examples=[to_gql_dataset_example(example) for example in examples],
         )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
@@ -314,14 +332,23 @@ class DatasetSplitMutationMixin:
 
             await session.execute(stmt)
 
+            examples = (
+                await session.scalars(
+                    select(models.DatasetExample).where(
+                        models.DatasetExample.id.in_(example_rowids)
+                    )
+                )
+            ).all()
+
         return RemoveDatasetExamplesFromDatasetSplitsMutationPayload(
             query=Query(),
+            examples=[to_gql_dataset_example(example) for example in examples],
         )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
     async def create_dataset_split_with_examples(
         self, info: Info[Context, None], input: CreateDatasetSplitWithExamplesInput
-    ) -> DatasetSplitMutationPayload:
+    ) -> DatasetSplitMutationPayloadWithExamples:
         user_id = get_user(info)
         validated_name = _validated_name(input.name)
         unique_example_rowids: set[int] = set()
@@ -374,9 +401,18 @@ class DatasetSplitMutationMixin:
                         "Failed to associate examples with the new dataset split."
                     ) from e
 
-        return DatasetSplitMutationPayload(
+            examples = (
+                await session.scalars(
+                    select(models.DatasetExample).where(
+                        models.DatasetExample.id.in_(example_rowids)
+                    )
+                )
+            ).all()
+
+        return DatasetSplitMutationPayloadWithExamples(
             dataset_split=to_gql_dataset_split(dataset_split_orm),
             query=Query(),
+            examples=[to_gql_dataset_example(example) for example in examples],
         )
 
 

phoenix/server/api/queries.py

@@ -48,6 +48,7 @@ from phoenix.server.api.types.AnnotationConfig import AnnotationConfig, to_gql_a
 from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
 from phoenix.server.api.types.Dataset import Dataset, to_gql_dataset
 from phoenix.server.api.types.DatasetExample import DatasetExample
+from phoenix.server.api.types.DatasetLabel import DatasetLabel, to_gql_dataset_label
 from phoenix.server.api.types.DatasetSplit import DatasetSplit, to_gql_dataset_split
 from phoenix.server.api.types.Dimension import to_gql_dimension
 from phoenix.server.api.types.EmbeddingDimension import (
@@ -1149,6 +1150,26 @@ class Query:
                 args=args,
             )
 
+    @strawberry.field
+    async def dataset_labels(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[DatasetLabel]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        async with info.context.db() as session:
+            dataset_labels = await session.scalars(select(models.DatasetLabel))
+        data = [to_gql_dataset_label(dataset_label) for dataset_label in dataset_labels]
+        return connection_from_list(data=data, args=args)
+
     @strawberry.field
     async def dataset_splits(
         self,

phoenix/server/api/routers/auth.py

@@ -2,7 +2,6 @@ import asyncio
 import secrets
 from datetime import datetime, timedelta, timezone
 from functools import partial
-from pathlib import Path
 from urllib.parse import urlencode, urlparse, urlunparse
 
 from fastapi import APIRouter, Depends, HTTPException, Request, Response
@@ -38,7 +37,6 @@ from phoenix.config import (
     get_base_url,
     get_env_disable_basic_auth,
     get_env_disable_rate_limit,
-    get_env_host_root_path,
 )
 from phoenix.db import models
 from phoenix.server.bearer_auth import PhoenixUser, create_access_and_refresh_tokens
@@ -52,6 +50,7 @@ from phoenix.server.types import (
     TokenStore,
     UserId,
 )
+from phoenix.server.utils import prepend_root_path
 
 rate_limiter = ServerRateLimiter(
     per_second_rate_limit=0.2,
@@ -145,7 +144,8 @@ async def logout(
         user_id = subject
     if user_id:
         await token_store.log_out(user_id)
-    redirect_url = "/logout" if get_env_disable_basic_auth() else "/login"
+    redirect_path = "/logout" if get_env_disable_basic_auth() else "/login"
+    redirect_url = prepend_root_path(request.scope, redirect_path)
     response = Response(status_code=HTTP_302_FOUND, headers={"Location": redirect_url})
     response = delete_access_token_cookie(response)
     response = delete_refresh_token_cookie(response)
@@ -242,9 +242,9 @@ async def initiate_password_reset(request: Request) -> Response:
     )
     token, _ = await token_store.create_password_reset_token(password_reset_token_claims)
     url = urlparse(request.headers.get("referer") or get_base_url())
-    path = Path(get_env_host_root_path()) / "reset-password-with-token"
+    path = prepend_root_path(request.scope, "/reset-password-with-token")
     query_string = urlencode(dict(token=token))
-    components = (url.scheme, url.netloc, path.as_posix(), "", query_string, "")
+    components = (url.scheme, url.netloc, path, "", query_string, "")
     reset_url = urlunparse(components)
     await sender.send_password_reset_email(email, reset_url)
     return Response(status_code=HTTP_204_NO_CONTENT)

phoenix/server/api/routers/oauth2.py

@@ -46,6 +46,7 @@ from phoenix.server.rate_limiters import (
     fastapi_route_rate_limiter,
 )
 from phoenix.server.types import TokenStore
+from phoenix.server.utils import get_root_path, prepend_root_path
 
 _LOWERCASE_ALPHANUMS_AND_UNDERSCORES = r"[a-z0-9_]+"
 
@@ -191,7 +192,7 @@ async def create_tokens(
         access_token_expiry=access_token_expiry,
         refresh_token_expiry=refresh_token_expiry,
     )
-    redirect_path = _prepend_root_path_if_exists(request=request, path=return_url or "/")
+    redirect_path = prepend_root_path(request.scope, return_url or "/")
     response = RedirectResponse(
         url=redirect_path,
         status_code=HTTP_302_FOUND,
@@ -564,8 +565,8 @@ def _redirect_to_login(*, request: Request, error: str) -> RedirectResponse:
     Creates a RedirectResponse to the login page to display an error message.
     """
     # TODO: this needs some cleanup
-    login_path = _prepend_root_path_if_exists(
-        request=request, path="/login" if not get_env_disable_basic_auth() else "/logout"
+    login_path = prepend_root_path(
+        request.scope, "/login" if not get_env_disable_basic_auth() else "/logout"
     )
     url = URL(login_path).include_query_params(error=error)
     response = RedirectResponse(url=url)
@@ -574,34 +575,15 @@ def _redirect_to_login(*, request: Request, error: str) -> RedirectResponse:
     return response
 
 
-def _prepend_root_path_if_exists(*, request: Request, path: str) -> str:
-    """
-    If a root path is configured, prepends it to the input path.
-    """
-    if not path.startswith("/"):
-        raise ValueError("path must start with a forward slash")
-    root_path = _get_root_path(request=request)
-    if root_path.endswith("/"):
-        root_path = root_path.rstrip("/")
-    return root_path + path
-
-
 def _append_root_path_if_exists(*, request: Request, base_url: str) -> str:
     """
     If a root path is configured, appends it to the input base url.
     """
-    if not (root_path := _get_root_path(request=request)):
+    if not (root_path := get_root_path(request.scope)):
         return base_url
     return str(URLPath(root_path).make_absolute_url(base_url=base_url))
 
 
-def _get_root_path(*, request: Request) -> str:
-    """
-    Gets the root path from the request.
-    """
-    return str(request.scope.get("root_path", ""))
-
-
 def _get_create_tokens_endpoint(*, request: Request, origin_url: str, idp_name: str) -> str:
     """
     Gets the endpoint for create tokens route.

phoenix/server/api/types/Dataset.py

@@ -18,6 +18,7 @@ from phoenix.server.api.types.DatasetExample import DatasetExample
 from phoenix.server.api.types.DatasetExperimentAnnotationSummary import (
     DatasetExperimentAnnotationSummary,
 )
+from phoenix.server.api.types.DatasetLabel import DatasetLabel, to_gql_dataset_label
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.Experiment import Experiment, to_gql_experiment
 from phoenix.server.api.types.node import from_global_id_with_expected_type
@@ -303,6 +304,13 @@ class Dataset(Node):
                 async for scores_tuple in await session.stream(query)
             ]
 
+    @strawberry.field
+    async def labels(self, info: Info[Context, None]) -> list[DatasetLabel]:
+        return [
+            to_gql_dataset_label(label)
+            for label in await info.context.data_loaders.dataset_labels.load(self.id_attr)
+        ]
+
     @strawberry.field
     def last_updated_at(self, info: Info[Context, None]) -> Optional[datetime]:
         return info.context.last_updated_at.get(self._table, self.id_attr)

phoenix/server/api/types/DatasetExample.py

@@ -142,3 +142,10 @@ class DatasetExample(Node):
             to_gql_dataset_split(split)
             for split in await info.context.data_loaders.dataset_example_splits.load(self.id_attr)
         ]
+
+
+def to_gql_dataset_example(example: models.DatasetExample) -> DatasetExample:
+    return DatasetExample(
+        id_attr=example.id,
+        created_at=example.created_at,
+    )

phoenix/server/api/types/DatasetLabel.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+import strawberry
+from strawberry.relay import Node, NodeID
+
+from phoenix.db import models
+
+
+@strawberry.type
+class DatasetLabel(Node):
+    id_attr: NodeID[int]
+    name: str
+    description: Optional[str]
+    color: str
+
+
+def to_gql_dataset_label(dataset_label: models.DatasetLabel) -> DatasetLabel:
+    return DatasetLabel(
+        id_attr=dataset_label.id,
+        name=dataset_label.name,
+        description=dataset_label.description,
+        color=dataset_label.color,
+    )

phoenix/server/api/types/Prompt.py

@@ -9,6 +9,7 @@ from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.types import Info
 
 from phoenix.db import models
+from phoenix.db.types.identifier import Identifier as IdentifierModel
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import NotFound
 from phoenix.server.api.types.Identifier import Identifier
@@ -37,7 +38,10 @@ class Prompt(Node):
 
     @strawberry.field
     async def version(
-        self, info: Info[Context, None], version_id: Optional[GlobalID] = None
+        self,
+        info: Info[Context, None],
+        version_id: Optional[GlobalID] = None,
+        tag_name: Optional[Identifier] = None,
     ) -> PromptVersion:
         async with info.context.db() as session:
             if version_id:
@@ -50,6 +54,19 @@ class Prompt(Node):
                 )
                 if not version:
                     raise NotFound(f"Prompt version not found: {version_id}")
+            elif tag_name:
+                try:
+                    name = IdentifierModel(tag_name)
+                except ValueError:
+                    raise NotFound(f"Prompt version tag not found: {tag_name}")
+                version = await session.scalar(
+                    select(models.PromptVersion)
+                    .where(models.PromptVersion.prompt_id == self.id_attr)
+                    .join_from(models.PromptVersion, models.PromptVersionTag)
+                    .where(models.PromptVersionTag.name == name)
+                )
+                if not version:
+                    raise NotFound(f"This prompt has no associated versions by tag {tag_name}")
             else:
                 stmt = (
                     select(models.PromptVersion)