feat(ui): artifact ref links

[jamie-rasmussen]

Description

Pulling frontend changes out of #2920 and updating them to resolve merge conflicts

Before:

After:

Testing

Now that the changes in #3484 have been deployed, these changes can be tested against the prod trace server backend. I used a tweaked version of the script from the original PR.

from typing import List
import weave

import wandb

# Configuration Constants
PROJECT = "2025-01-23_artref"
TEAM_ENTITY = "jamie-team"  # MAKE SURE THIS IS NOT A PERSONAL ENTITY
PERSONAL_ENTITY = "jamie-rasmussen" # Most likely your username
ORG_ENTITY = "wandb-jamie"  # Fetch this from the UI banner {BASE_URL}/orgs/wandb/registry/model

artifact_name = "test"

# log an empty personal entity artifact
run = wandb.init(project=PROJECT, entity=PERSONAL_ENTITY)
model_artifact = wandb.Artifact(name=artifact_name, type="model")
run.log_artifact(model_artifact).wait()
ARTIFACT = f"wandb-artifact:///{model_artifact.qualified_name}"
print(f"Logged personal entity artifact: {ARTIFACT}")

# link to model registry artifact
run.link_artifact(model_artifact, f"model-registry/{artifact_name}")
MODEL_REG_URL = f"wandb-artifact:///{PERSONAL_ENTITY}/model-registry/{artifact_name}:latest"
run.finish()

# Log a team entity artifact
run = wandb.init(project=PROJECT, entity=TEAM_ENTITY)
model_artifact = wandb.Artifact(name=artifact_name, type="model")
run.log_artifact(model_artifact).wait()
ARTIFACT = f"wandb-artifact:///{model_artifact.qualified_name}"
print(f"Logged team entity artifact: {ARTIFACT}")


# link to global registry
run.link_artifact(model_artifact, f"wandb-registry-model/{artifact_name}")
GLOBAL_REGISTRY_URL = (
    f"wandb-artifact:///{ORG_ENTITY}/wandb-registry-model/{artifact_name}:latest"
)
print(f"Linked team entity artifact to global registry: {GLOBAL_REGISTRY_URL}")
run.finish()


NON_EXISTENT_ARTIFACT = "wandb-artifact:///im-not-valid/im-not-real/testArtifact:v1"


class UnslothLoRAChatModel(weave.Model):
    """
    ChatModel class to store and version additional parameters beyond just the model name.
    Especially useful for fine-tuning scenarios.
    """

    model_registry: str
    wandb_registry_model: str
    regular_artifact: str
    invalid_artifact: str
    cm_temperature: float
    device: str

    @weave.op()
    def predict(self, query: List[str]) -> str:
        return "Berlin"


def main():
    # Instantiate the Chat Model
    new_chat_model = UnslothLoRAChatModel(
        name="UnslothLoRAChatModelRag",
        model_registry=MODEL_REG_URL,
        wandb_registry_model=GLOBAL_REGISTRY_URL,
        regular_artifact=ARTIFACT,
        invalid_artifact=NON_EXISTENT_ARTIFACT,
        cm_temperature=1.0,
        device="cpu",
    )
    print(new_chat_model)

    weave.init(f"{PERSONAL_ENTITY}/{PROJECT}")

    result = new_chat_model.predict(["What is the capital of Germany?"])
    print(result)


if __name__ == "__main__":
    main()

[circle-job-mirror]

Preview this PR with FeatureBee: https://beta.wandb.ai/?betaVersion=c723d6d6144e6ca41795e17f931b6d6cc4bc2293