AI Representation

[mlucool]

In this JEP we propose:

1. A way for objects to represent themselves at runtime for AI
2. A registry for users to define representations for objects that do not have them
3. A new messaging protocol to query for this data

JEP for #128

[mlucool]

CC @krassowski @ellisonbg @dlqqq @jtpio @SylvainCorlay @Zsailer @brichet @Carreau

[dlqqq]

@mlucool @govinda18 Thank you for opening this JEP! Really excited to see this moving forward. 💪

Left some feedback & typo fixes below.

[dlqqq]

Comment: Consumers of this protocol (e.g. Jupyter AI) will likely not be able to support every MIME type, nor do they need to. Jupyter AI should document which MIME types we read from. That way, people implementing these methods know how to define _ai_repr_() to provide usable reprs for Jupyter AI. Other consuming extensions should do the same.

[dlqqq]

Is it worth defining a type for the dictionary returned by _ai_repr_() instead of using Dict[str, Any]? This may be important in the case of image reprs, since the same MIME type may be encoded in different ways. For example, it's ambiguous as to whether image/png may return a bytes object or a base64-encoded string.

MIME types do allow encodings/parameters to also be specified, so image/png;base64 could refer to a string object while image/png could refer to a bytes object. It may be worth defining this in the same proposal to reduce ambiguity.

[dlqqq]

Hmm, one problem is that using a TypedDict doesn't allow extra keys to be defined. So if we define _ai_repr_() -> AiRepr, implementers can only use the keys we define in the AiRepr type. Another issue that this custom type would have to be provided by some dependency, which means every consumer & implementer needs 1 extra dependency.

We should continue to think of ways to provide better type safety guarantees on the values in the returned dictionary.

[dlqqq]

It may be better to have the async version live in a different method, e.g. _async_ai_repr_().

[dlqqq]

It seems like the kwargs structure needs to be documented somewhere. Would this be standardized by the kernel implementation when handling ai_repr_request?

[Carreau]

my understanding is that they should be some consensus, but some free parameters that could depends on which repr understand the need of which models.

[mlucool]

My intent was not to be prescriptive here and let that evolve naturally. That is, given jupyter-ai popularity, I expect it to implicitly set some kwargs as a standard. No kwarg should be required is maybe a better statement here, but plugins are free to document things they support.

[dlqqq]

A registry can define AI representations for objects that lack a _ai_repr_(). It seems natural to also a registry to define a "fallback" AI representation, i.e. the method to be used when neither a _ai_repr_() method or a registry entry exist for an object.

[Carreau]

My guess is the registry is already the fallback, because the object does not have a _ai_repr_, Unless you see the registry as an override of _ai_repr_ ?

[Carreau]

Do we want to allow dots in names ? it may be property and trigger side effects, but I think that is fine.

[mlucool]

I think that's ok too

[Carreau]

For passed kwargs, I would at least standardise 1 which is the list/set of mimetype accepted; I think it should be fairly standard to select between or text images, and having a standard would be good to start.

I think beyond the protocol, nothing is jupyter specific, and I'm happy to add some integration with IPython – even if likely not necessary.

I think we should likely have introspection facilities, like do we want to be able to list the various mimetype an object can provide ? And do we want the user to be able to ask for these ?

There is also the slight technical question that _ai_repr_(**kwargs) -> Dict[str, T], requires T to be serialisable, no really be Any.

But in general I'm +1, I'll see if I can prototype and push that to the Jupyter EC/SSC.

[mlucool]

There is also the slight technical question that ai_repr(**kwargs) -> Dict[str, T], requires T to be serialisable, no really be Any.

Agreed

[Carreau]

I think there was a bug in github (or my laptop), I submitted my review yesterday, but there was no internet, and it resubmitted today when I reopened the tab and internet was back. Sorry if there are crossed wires.

[krassowski]

This was discussed in the jupyter-ai call today.

We just wanted to ping CC @jupyter/software-steering-council on this :)

[minrk]

This seems generally sensible

My main design question is at the protocol level, why does this require a distinct message handler, as opposed to defining a schema for one (or more) custom x-jupyter/ai-repr mimetype via the existing display protocol, e.g. an inspect_request which already returns a mimebundle representing an object.

I don't really understand why the existing mimebundle messages have "performance concerns and their inability to adapt" while this doesn't, when it appears to do the same thing (returns a mimebundle representing an object), and has the same design (register reprs methods both via method name and explicit dispatch).

[minrk]

Absolutely minor detail, but is there a reason this inverts the standardized _repr_ai_ naming scheme to _ai_repr_? If not, then I'd suggest following the prefix precedent (using a prefix improves discoverability via tab completion, etc.).

[mlucool]

I believe this came from chatting with @Carreau and we wanted to avoid it ever being automatically called during display. (FWIW aside from this, if its not supposed to go into the general mime bundle, it feels strange to use the same pattern).

@Carreau do you recall the exact details?

[krassowski]

via the existing display protocol, e.g. an inspect_request which already returns a mimebundle representing an object.

inspect_request takes code, cursor_pos (position in given code) and detail_level (0 or 1). I proposed re-using this before, some of the counter arguments were:

• it does not allow to pass extra keyword arguments
• detail_level is poorly defined, and LLMs will often need a hard cutoff on number of tokens they want in representation as this it is proportional to cost/performance
• sending code and cursor_pos is a design deriving from "user has a cursor in a cell", not from "user asks about a variable in the chat". cursor_pos could be always set to 0 and code could represent the variable so it is not a strong argument

Now, another message already in the protocol that could be considered is richInspectVariables which is an addition to the Debugger Adopter Protocol. It is defined as:

{
    'type' : 'request',
    'command' : 'richInspectVariables',
    'arguments' : {
        'variableName' : str,
        # The frameId is used when the debugger hit a breakpoint only.
        'frameId' : int
    }
}

and returns:

{
    'type' : 'response',
    'success' : bool,
    'body' : {
        # Dictionary of rich representations of the variable
        'data' : dict,
        'metadata' : dict
    }
}

The advantage of using frame IDs is that local variables can be distinguished from global variables. Now, the problem is that debugger is opt-in, has side-effects and performance overhead so it is not necessarily a go-to solution. The basic DAP uses all three: name, variablesReference, and frameId to determine the variable being requested.

I wonder if when debugger is active, should we have a way to pass frameId and maybe variablesReference to allow user to ask about contents of a local vs non-local or global variable?

[krassowski]

A broader edge-case is definition disambiguation:

• Julia and R use multiple dispatch natively
• Python has definition overloads with @overload decorator
• the same name in TypeScript can refer to both a namespace and a class

While it is fine to return description of all cases, it can grow long. In these cases to get a definition of a function unambiguously we would need to know both the function and the argument(s); This is where inspect_request's way of doing things shines. Of course if would require the user to provide more context than just the function name.

This JEP does not say how user would specify @var. My assumption is that UI could use complete_request request to get a list of variables as user starts typing. If that's the case one solution to "which definition is it referring to" could be augmenting complete_reply to return a unique variablesReference for each completion. Imagine user asks:

What are the arguments for @predict

In R if they loaded a few statistical packages, they could be presented with a list dozens of different variants; a simple case (for brevity):

# S3 method for glm
predict(object, newdata = NULL,
            type = c("link", "response", "terms"),
            se.fit = FALSE, dispersion = NULL, terms = NULL,
            na.action = na.pass, …)

# S3 method for lm
predict(object, newdata, se.fit = FALSE, scale = NULL, df = Inf,
        interval = c("none", "confidence", "prediction"),
        level = 0.95, type = c("response", "terms"),
        terms = NULL, na.action = na.pass,
        pred.var = res.var/weights, weights = 1, …)

Now, we may want to present these as separate items in complete_reply for a few reasons:

• each has a different signature and documentation, concatenating all of these would be too long to display in completion panel
• each signature defines a different completion snippet (think template for filling in arguments, see Snippet Syntax for LSP way of doing it)

Then if a user chooses one or another, a different snippet would show up. These separate items could be annotated with a variable reference, so when it is used for AI request proposed in this JEP the user's choice would be remembered.

I think defining a variable reference in complete_reply and as an argument for ai_repr_request does not need to be a part of this JEP, but I think it highlights the difference from inspect_request which could not achieve the same easily, unless it's signature is changed in ways which I would suspect would be called breaking/requiring a major release.

[ivanov]

AI is such a nebulous term that whatever description or formulation this representation would hold could both be useful to humans, and also useless to past, present, and future applications that would still nevertheless be reasonably called AI. So I came here to chime in that this can already be done as a set of conventions around mimebundle keys, and hence not require a protocol revision.

Fortunately, @minrk is already doing the same

I don't really understand why the existing mimebundle messages have "performance concerns and their inability to adapt" while this doesn't, when it appears to do the same thing (returns a mimebundle representing an object), and has the same design (register reprs methods both via method name and explicit dispatch).

We need to be very deliberate and conservative with protocol revision, I think this would grow the surface area and complexity of interactions for the jupyter protocol. For example, if you end up with objects that declare "ai repr" but not other rich repr, do you end up with users who aren't using "ai" having to inspect objects twice in separate ways, hoping to get something better than <Foo object at 0x7f600ecf8990> ?

[mlucool]

AI is such a nebulous term that whatever description or formulation this representation would hold could both be useful to humans, and also useless to past, present, and future applications that would still nevertheless be reasonably called AI.

I agree it is nebulous, do you have a better term we should use? I started with calling it repr_llm, but llm seemed too specific.

I don't really understand why the existing mimebundle messages have "performance concerns and their inability to adapt" while this doesn't, when it appears to do the same thing (returns a mimebundle representing an object), and has the same design (register reprs methods both via method name and explicit dispatch).

This is a important question and I believe there are a number of reasons

1. If this is included in the mimebundle, then all displays of an object must also spend time computing an AI centric display of the data. This information is rarely needed to render an object in a notebook, but is useful for representing an object in an context window. This makes things slower to render in notebook and sends potentially a large amount of extra data.
2. I think an async version is useful. Imagine you want to render a JS chart in notebook (e.g. plotly). In the mime bundle you'll want to pass datapoints etc., and the front end has the chart library render. Now if I want to use this chart with an LLM, I'd want an image. This is a costly process to compute and likely happens outside the kernel (e.g. send it to playwright).
3. I believe you still need a new protocol for the reasons @krassowski outlines above (even if we go with this is just a mimebundle).

That all being said, 1/2 are a bit unrelated to the protocol change (we still need a way to get this data from UIs). If we came up with other solutions here (e.g. an async mimebundle, where there was a clean way to pass args to), I don't think it needs to be "ai".

[minrk]

I don't want to get off on too much of a tangent, but being able to select supported mimetype representations in the request is something that's seemed useful for some time, to avoid computing representations that won't be understood (like an http Accept header). Would a custom mime-type in existing messages be satisfactory if it could be opt-in via the request, and not computed by default? Then the same request that asks for the "genai repi" could explicitly exclude all the other reprs, too. That appears to address all of the performance/efficiency questions, IIUC.

[mlucool]

Would a custom mime-type in existing messages be satisfactory if it could be opt-in via the request, and not computed by default? Then the same request that asks for the "genai repi" could explicitly exclude all the other reprs, too. That appears to address all of the performance/efficiency questions, IIUC.

Not a tangent at all. It's a great question. Here's a first pass at what I think would have to change:

1. We still need a new message to reach into a kernel and call this method for the reasons described above.
2. It's not just about opting in/out but also letting kwargs flow through to the bundle. We'd want to really start using the optional kwargs, which seems fine, but is a departure to how this is called. I'm not sure the impact on existing code
3. The merits of an async version should still be discussed. Maybe I am over indexing on Chart->image, but I suspect it'll be really common. It's also possible that subshells is a cleaner way to allow for compatibility.