feat(weave): link artifact references in weave trace

[estellazx]

Description

https://wandb.atlassian.net/browse/WB-21724
Building on @levinandrew's PR

New loom with feedback:
https://www.loom.com/share/e25073c52f5c4f73bb523e043b1672ec?sid=6b3d6655-da3f-4645-9543-a04d54d0b224

old loom:
https://www.loom.com/share/056ae7450d064784b6c8b3a79dfcbc97?sid=1ed27f9c-d9fe-42a8-b2bc-0a93cb676a41

Testing

locally. Note, you need to have a

1. a team entity
2. registries set up (all this requires is you to go to http://localhost:9000/orgs/gr-org/registry/ for example and it will auto set up Registries of type model)
3. a personal entity to have model registries set up (this is not a must you can just comment out the part where we test model registry code)

Test code example:

from typing import List
import weave

import wandb
import os

os.environ["WANDB_BASE_URL"] = "http://localhost:9001"
os.environ["WF_TRACE_SERVER_URL"] = "http://127.0.01:6345"

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

artifact_name = "test"

# log a 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}"

# 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}"


# 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"
)

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",
    )

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

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


if __name__ == "__main__":
    main()

TODO: Test on QA

[circle-job-mirror]

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

[ibindlish]

Reviewed the registry/model-registry url changes. Recommended some minor changes

[jamie-rasmussen]

Trace server changes pulled out (with minor tweaks) as #3484

[jamie-rasmussen]

FE changes were updated and merged as #3500