chore(weave): Implement enhaced feedback structure and mvp filter/query layer

[tssweeney]

This PR lays the groundwork for the next leg of feedback types in our system. Specifically, we have two "classes" of feedback: runnables and annotations. "Runnables" are feedbacks that are generated by running a program (think: Op, Configured Action, Scorer), while "Annotations" are feedbacks created by humans with specific types (aka human in the loop, aka custom columns, etc...).

There were three problems to solve with this emerging data model:

1. How do we store additional metadata about these feedbacks linking them to other objects in our system.
2. How do we group / collect feedbacks that belong to the same logical "column" or concept?
3. Given the above, how can we filter/sort without reading large JSON dumped columns?

After much iteration and discussion, the solution that seemed most suitable is as follows:

1. feedback_type how has 2 special prefixes: wandb.runnable and wandb.annotation, where the total type should be wandb.runnable.RUNNABLE_NAME or wandb.annotation.ANNOTATION_NAME. here, RUNNABLE_NAME or ANNOTATION_NAME are the name (aka object_id) components of the backing Object or Op. This is the most common group key and indexed in Clickhouse already.
2. I have added 4 new columns to the feedback table. Note, I originally had these as fields in the payload itself, but this would result in more complex, heavy lookups & more ridged structure over the payload itself. This approach allows us to put our foreign keys in columns that can be indexed in the future if needed.:
   • annotation_ref: The ref pointing to the annotation definition for this feedback.
   • runnable_ref: The ref pointing to the runnable definition for this feedback.
   • call_ref: The ref pointing to the resulting call associated with generating this feedback.
   • trigger_ref: The ref pointing to the trigger definition which resulted in this feedback.
3. When these types of feedback are entered into the DB, our server now will enforce that these ref values are filled out when required and matched the correct format. Moreover, the payloads themselves conform to a very simple structure:

class AnnotationPayloadSchema(BaseModel):
    value: Any


class RunnablePayloadSchema(BaseModel):
    output: Any

4. Finally, I implement basic (not optimized yet) support for filter and sort on the calls query using the notation feedback[feedback_type].payload.json.selector. This allows us to specify the feedback type (while supporting dots) and match our other field access patterns.

---

With all of this together, we can have code like:

@weave.op
def my_scorer(x: int, output: int) -> int:
    expected = ["a", "b", "c", "d"][x]
    return {
        "model_output": output,
        "match": output == expected,
    }

@weave.op
def my_model(x: int) -> str:
    return [
        "a",
        "x",  # intentional "mistake"
        "c",
        "y",  # intentional "mistake"
    ][x]

ids = []
for x in range(4):
    _, c = my_model.call(x)
    ids.append(c.id)
    # Note: `_apply_scorer` is not user-facing (yet!) but will be made public during the eval api project.
    c._apply_scorer(my_scorer)

... then query ...

calls = client.server.calls_query_stream(
    tsi.CallsQueryReq(
        project_id=client._project_id(),
        filter=tsi.CallsFilter(op_names=[get_ref(my_model).uri()]),
        # Filter down to just correct matches
        query={
            "$expr": {
                "$eq": [
                    {
                        "$getField": "feedback.[wandb.runnable.my_scorer].payload.output.match"
                    },
                    {"$literal": "true"},
                ]
            }
        },
        # Sort by the model output desc
        sort_by=[
            {
                "field": "feedback.[wandb.runnable.my_scorer].payload.output.model_output",
                "direction": "desc",
            }
        ],
    )
)

This can easily be extended to support different aggregation logic and specific version selectors.

[circle-job-mirror]

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

[tssweeney]

This is ok to change as the UI/query layer does not consume it yet.

[tssweeney]

again, safe to change now that we have a good format

[tssweeney]

we learned from this - no longer needed

[gtarpenning]

This new feedback query is going to be spicy in big projects, but looks good. The calls query builder is also feeling... clunky. Generally this makes sense, I wonder how much of the implementation we can abstract away from the user when adding, but still create an intuitive way for them to get the data out. It's possible that we might want to have some way of auto-constructing queries client-side, i'm imagining users not finding the following easy to use....
"$getField": "feedback.[wandb.runnable.my_scorer].payload.output.match"

[gtarpenning]

any reason to do this concat in the query vs outside and pass it in?

[tssweeney]

i don't think so? but not sure. We have to do a concat either way since the last part is dynamic

[gtarpenning]

it would be nice if we could type this to a kind of ref, like objectRef, with a pydantic validator and then check its construction in the client.

[tssweeney]

Agreed