Add Ability to Specify Gemini Safety Settings

[dAIsySHEng1]

fix #373

Add ability to specify Gemini safety settings when doing an agent run. This addition allows you to specify safety categories and thresholds via the model_settings field. Valid settings are based on the Gemini API docs. This should hopefully add flexibility to deal with Gemini's safety guardrails as needed, aligning with what one is able to currently do in the vertexai package.

Example Code:

from pydantic_ai import Agent
from pydantic_ai.models.gemini import GeminiModel, GeminiModelSettings

GEMINI_API_KEY = '<api key here>'

model = GeminiModel('gemini-1.5-flash', api_key=GEMINI_API_KEY)
agent = Agent(model, system_prompt='Be concise, reply with one sentence.')

result = agent.run_sync(
    'Write a list of 3 mean things that I might say to Paddington after he steals my marmalade:',
    model_settings=GeminiModelSettings(
        gemini_safety_settings=[
            {'category': 'HARM_CATEGORY_HARASSMENT', 'threshold': 'BLOCK_LOW_AND_ABOVE'},
            {'category': 'HARM_CATEGORY_HATE_SPEECH', 'threshold': 'BLOCK_LOW_AND_ABOVE'},
            {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'threshold': 'BLOCK_LOW_AND_ABOVE'},
        ],
    ),
)

To Do:

• Fix validation errors when running live test
  • Add blocked field to _GeminiSafetyRating
  • Revise code to raise an UnexpectedModelBehavior error when finish_reason is 'SAFETY'
• Figure out way to deal with potential duplicates (deferred)
• Do more testing and add tests
• Update docs (Gemini, VertexAI?)
• Add more detail to PR description

[sydney-runkle]

This looks fantastic, @dAIsySHEng1! Great work.

In your PR description, you can add "fix #373", and when this PR is merged, it'll automatically close that issue. If you could add a bit more of a detailed description (just a few sentences), that would be great as well.

The test looks great. Given that we'll probably want to test this sort of pattern often, we should generalize this further (given the similar pattern in the above test). That lies outside of the scope of this PR, though.

[sydney-runkle]

Ah one other thing - the example you sent me on slack is hilarious, let's add this to the gemini model docs :)

[zaidalyafeai]

Will this also work for VertexAIModel?

[Finndersen]

Would it perhaps be better to make these StrEnum? e.g.

class HarmCategory(StrEnum):
    UNSPECIFIED = auto()
    DEROGATORY = auto()
    ....

This actually already exists in the SDK

[Finndersen]

And perhaps even make GeminiSafetySettings a dataclass instead?

[sydney-runkle]

I think Literal is fine here - it's consistent with other patterns we're using in this file. If we were using the sdk, I definitely think we should pull from there.

[sydney-runkle]

Why make GeminiSafetySettings a dataclass? Does it need any methods?

[Finndersen]

Enums provide better developer experience with stronger type safety, IDE autocomplete, enumeration, avoiding typos etc

[Finndersen]

And dataclasses often don't have any methods, that's not their purpose. It's for having a date structure similar to Pydantic models but without overhead of validation. Again provides better developer experience of attribute instead of key access, with IDE autocomplete, stronger type safety, defaults, etc

[Finndersen]

I prefer using enums & dataclasses over str literals and typeddicts for more explicit typing, But I realise that opinion is somewhat subjective

[hyperlint-ai]

The style guide flagged several spelling errors that seemed like false positives. We skipped posting inline suggestions for the following words:

• Chinoy
• Hussain

[sydney-runkle]

You can remove my previous example, this one is better.

Just preface with:

If you wish to further customize model behavior, you can use a subclass of [`ModelSettings`][pydantic_ai.settings.ModelSettings], like [`GeminiModelSettings`][pydantic_ai.models.gemini.GeminiModelSettings], associated with your model of choice.
For example:

[sydney-runkle]

Maybe wrap this in a try/except block and print the error response?

[sydney-runkle]

I'm going to chat with the team tomorrow morning to see if this is the right way to go about this. I think we might just want to catch the validation issue earlier.

[sydney-runkle]

Ok confirmed, we can move forward with this approach for now 👍