[ENH] - Add support for AI21labs assistants

[smokestacklightnin]

This PR addresses Issue #226 by adding an assistant for AI21 Labs

Please leave opinions and change requests in comments

[smokestacklightnin]

@pmeier The first draft of this PR is finished. I'm testing it now.

Opinions and requests to modify are welcome

[smokestacklightnin]

I'm trying to test this using basic documents and prompts, but I keep getting an error like

[...]/_api.py:28: RuntimeWarning: coroutine 'AI21LabsAssistant._call_api' was never awaited
  async for chunk in self._call_api(  # type: ignore[attr-defined, misc]

I'm trying to debug this, including adding a breakpoint(), but the code has been rather opaque to peer into.
I've tried making a basic AI21Labs API request using a simple standalone script, and that seemed to work.

I'm looking for ideas and input, please

[smokestacklightnin]

An educated guess is that streaming is not specified. I will test this hypothesis

Edit: Simply adding the stream option to the http request didn't fix the error, although stream isn't part of the API parameters

[smokestacklightnin]

A later part of the same error gives

File "/home/[...]/ragna/ragna/assistants/_api.py", line 28, in answer
    |     async for chunk in self._call_api(  # type: ignore[attr-defined, misc]
    | TypeError: 'async for' requires an object with __aiter__ method, got coroutine

What's confusing to me is that RagnaException is not getting raised, leading me to believe that the request is (hopefully) coming back successfully. On the other hand, if _call_api isn't being awaited, then something is being held up.

[pmeier]

Thanks @smokestacklightnin for the PR. This already looks pretty good. The issue you are facing has a simple resolution. I've suggested it below along two other minor things.

[pmeier]

Almost there just one minor thing below. Apart from that, I'm unable to use the Jurassic2Mid assistant. The server always returns a 500 / Internal Server error for me. Did it work for you? If so, let's wait for them to fix this so I get at least one clean run.

[smokestacklightnin]

Almost there just one minor thing below. Apart from that, I'm unable to use the Jurassic2Mid assistant. The server always returns a 500 / Internal Server error for me. Did it work for you? If so, let's wait for them to fix this so I get at least one clean run.

I get the same 500 http response code for Jurassic2Mid, but Jurassic2Ultra works for me. Is it an option to comment out the Mid model and leave only the Ultra?

[smokestacklightnin]

It's also unclear to me why there's a new import error

[pmeier]

I get the same 500 http response code for Jurassic2Mid,

Are you saying you are also getting the 500 now or never got a proper response so far? If its the former, I say we wait a few days before we merge this and see if they figure it out on their side. No worries, this PR will make it into the 0.2.0 release

It's also unclear to me why there's a new import error

Well, lancedb==5.0.2 was released an hour ago and our env in CI is likely still cached and does not contain polars. Nothing to worry about on your side.

[smokestacklightnin]

Are you saying you are also getting the 500 now or never got a proper response so far? If its the former, I say we wait a few days before we merge this and see if they figure it out on their side. No worries, this PR will make it into the 0.2.0 release

The former

[pmeier]

In that case let's revisit this next week and see if they have fixed their server.

[pmeier]

I've tried again today and I'm still getting the 500 error. Since the release is still at least a week out I'm ok with waiting. However, also fine to just remove JurrassicMid and only add it again later.

[smokestacklightnin]

I've tried again today and I'm still getting the 500 error. Since the release is still at least a week out I'm ok with waiting. However, also fine to just remove JurrassicMid and only add it again later.

I'm mostly fine with either option.

I tested JurassicUltra with a text file that reads

Roses are red
Violets are blue

and then asked the system the question

How many legs does a spider have?

to which the system responded

A spider has eight legs.

While this is true, it doesn't feel right that the information provided didn't come from the included documents, even though the system is only supposed to answer the questions from what's included. I double checked the documentation (this and this), and everything seems to be set up correctly in the code from this PR. What are your thoughts?

[pmeier]

I'm mostly fine with either option.

In that case let's wait a little longer. If we get closer to release and still have no green light, we can always pull the trigger then.

While this is true, it doesn't feel right that the information provided didn't come from the included documents, even though the system is only supposed to answer the questions from what's included.

Unfortunately, this is simply how LLMs work. There is no way to force them into any behavior. The system prompt just gives an instruction to it. Whether or not this is respected is a completely different story. IMO this is a quality metric and it seems that the Jurassic model is not on par with the other LLMs that we have built in.

Unless it actually can't answer a question even if given the right context, I still think it makes sense to include it.

[smokestacklightnin]

The 500 internal service error is still present from AI21labs, but I've rebased this branch to get rid of the conflict in the FAQ documentation file

[pmeier]

Up to you how you want to handle this: keep waiting for them to fix their server or remove the "mid" model from the PR and we can merge as is.

[smokestacklightnin]

Up to you how you want to handle this: keep waiting for them to fix their server or remove the "mid" model from the PR and we can merge as is.

I'd love to have another PR merge on the books. If you don't see any reason not to comment out the Mid model, I'm happy to go ahead and do so to make this PR ready to merge

[pmeier]

Please leave a comment on the "mid" model on why this is commented out linking to this thread.

[smokestacklightnin]

Please leave a comment on the "mid" model on why this is commented out linking to this thread.

I added a comment. I'll add a link to this thread

[smokestacklightnin]

Please leave a comment on the "mid" model on why this is commented out linking to this thread.

Done

[pmeier]

Thanks @smokestacklightnin!