Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Gemini integration #1243

Draft
wants to merge 5 commits into
base: v1.0
Choose a base branch
from
Draft

Add Gemini integration #1243

wants to merge 5 commits into from

Conversation

rlouf
Copy link
Member

@rlouf rlouf commented Nov 5, 2024

This PR provides a thin wrapper around the Gemini client. Gemini provides structured generation by accepting the following to define output types:

  • A pydantic BaseModel. The output type is assumed to be JSON.
  • A TypedDict. The output type is assumed to be JSON.
  • A List type with a BaseModel or TypedDict argument. The output type is assumed to be a list of JSON objects.
  • An Enum. The output type is assumed to be one of the fields in the enum.

This integration is already fairly complex:

  1. We need to define a new Choice output type to handle the Enum output type;
  2. The Gemini SDK does not accept JSON Schema string. Json thus needs to keep track of the original definition object. We should thus add a method to translate the definition object to JSON Schema instead of translating during initialization.
  3. The Gemini SDK also accepts TypedDicts
  4. The Gemini SDK accepts List type with arguments.

TODO

  • Handle list and List more elegantly
  • Use Json to infer the type of the definition object / reject invalid definition objects
  • Add Json method to convert the definition object to JSON Schema
  • Use Choice to infer the type of the definition object / reject invalid definition objects

Builds on #1240.

@rlouf rlouf added this to the 1.0 milestone Nov 5, 2024
@rlouf rlouf force-pushed the add-gemini branch 2 times, most recently from 71adef9 to 7edb6bd Compare November 5, 2024 11:13

@format_output_type.register(Json)
def format_json_output_type(self, output_type):
if issubclass(output_type.original_definition, BaseModel):
Copy link
Member Author

@rlouf rlouf Nov 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could let the Json object do type inference for the definition (and thus be the judge of what counts as a valid definition for a JSON object), and use a simple match statement here.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1, a good indication of this is that this part stays the same:

            return {
                "response_mime_type": "application/json",
                "response_schema": output_type.original_definition,
            }

Keeping this translation clean of any other extra logic is desirable to clearly separate concerns.


@singledispatchmethod
def format_output_type(self, output_type):
if output_type.__origin__ == list:
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Having to handle List types here is not ideal.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is interesting, seems that any model might have such custom types or even named similarly, but maybe defined a bit differently. We wouldn't want to limit these expressions too, but at the same time offer it only locally within this models?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you mean define our own List type?

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, maybe it might worth clearly defining such custom types separately for the model. And since they're specific to the model they won't be under common ones:
from outlines.types import Choice, Json

But instead from somewhere like:
from outlines.gemini.types import List?

Considering the variations of the models, it seems it might be quite common situation to manage the difference between them all. What do you think?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that

from outlines.types import Choice, Json, List

would make sense as we want to support this type for open source models down the line as well.

We may want to support both list and typing.List down the line, but this is a quick and easy solution for now.

Copy link

@torymur torymur Nov 12, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, extending the common when we are sure it's a good addition is ok, I wonder how to handle purely custom ones, which we wouldn't want to support for any other model but this particular one. For that we'd need some good examples, if List is out to be such custom example.

@rlouf rlouf self-assigned this Nov 5, 2024
@rlouf rlouf linked an issue Nov 5, 2024 that may be closed by this pull request
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add Gemini integration
2 participants