Skip to content

[MRG] feat(model): Update gemini.py to use new Gemini API#309

Merged
huangyz0918 merged 11 commits into
MLSysOps:mainfrom
syangx38:main
Jul 13, 2025
Merged

[MRG] feat(model): Update gemini.py to use new Gemini API#309
huangyz0918 merged 11 commits into
MLSysOps:mainfrom
syangx38:main

Conversation

@syangx38
Copy link
Copy Markdown
Contributor

@syangx38 syangx38 commented Jul 12, 2025

Closes #302

advisor.py
remove the duplicate schema_web_search

gemini.py uses the google-genai SDK, highlights:

  1. gemini api migrate doc states the new SDK promotes a stateless approach as its primary interaction model, which is a significant change from the stateful sessions that were central to the old SDK. Specially, instead of model.start_chat, now it's client.models.generate_content()
  2. As pointed out in the issue, the new Gemini API function calling requires an explicit process to first get a tool call suggestion and then a separate API call to get the final structured output (AdviseAgent etc all requires structured JSON to work).
  3. Instead of using recursive (as openai.py does), used iterative loop to manage the multi-turn conversation for tool calling.
  4. def query() method logic:
def query(chat_history, **kwargs):
    // Initial Setup
    system_instruction, prompt
    tools
    json_output_required = False // Flag to force JSON output

    // Tool-Calling Loop
    FOR LOOP:
        config = json_only_config IF json_output_required ELSE base_config
        json_output_required = False

        response = call_gemini_api(prompt, config)
        function_call = get_function_call_from(response)

        IF function_call:
            execute the local function            
            append model's response part to prompt
            append tool_result
            json_output_required = True // Set the flag 
        ELSE:
            // The model returned text
            final_response_content = get_text_from(response)
            BREAK 

    // Return the raw string content for the calling agent to handle
    RETURN final_response_content

What has been done to verify that this works as intended?

sample logs below

Applying feature engineering...
Starting training with 5-Fold Stratified Cross-Validation...
--- Fold 1/5 ---
--- Fold 2/5 ---
--- Fold 3/5 ---
--- Fold 4/5 ---
--- Fold 5/5 ---
Overall OOF AUC: 0.88561
Submission file created successfully at ./tabular-playground-may-2022/public/submission.csv
First 5 rows of submission file:
       id    target
0  800000  0.609488
1  800001  0.440478
2  800002  0.892964
3  800003  0.574231
4  800004  0.362269

Why is this the best possible solution? Were any other approaches considered?

Checked as the official doc required https://ai.google.dev/gemini-api/docs/migrate extensively.

How does this change affect users? Describe intentional changes to behavior and behavior that could have accidentally been affected by code changes. In other words, what are the regression risks?

Will implement gemini users, but users can switch to other models.

Do we need any specific form for testing your changes? If so, please attach one.

No.

Does this change require updates to documentation? If so, please file an issue here and include the link below.

No.

Before submitting this PR, please make sure you have:

  • confirmed all checks still pass OR confirm CI build passes.
  • verified that any code or assets from external sources are properly credited in comments and/or in
    the credit file.

@dosubot dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. enhancement New feature or request labels Jul 12, 2025
@syangx38 syangx38 changed the title [DRAFT] feat(model): Update gemini.py to use new Gemini API [Ready] feat(model): Update gemini.py to use new Gemini API Jul 13, 2025
@syangx38 syangx38 changed the title [Ready] feat(model): Update gemini.py to use new Gemini API [MRG] feat(model): Update gemini.py to use new Gemini API Jul 13, 2025
Copy link
Copy Markdown
Collaborator

@leeeizhang leeeizhang left a comment

Choose a reason for hiding this comment

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

LGTM! many thanks

@dosubot dosubot Bot added the lgtm This PR has been approved by a maintainer label Jul 13, 2025
Comment thread mle/agents/advisor.py
@YuanmingLeee
Copy link
Copy Markdown
Contributor

Also fix lint:
./mle/model/gemini.py:116:37: W291 trailing whitespace
./mle/model/gemini.py:129:54: W291 trailing whitespace
./mle/model/gemini.py:156:90: W291 trailing whitespace

@syangx38 syangx38 requested a review from YuanmingLeee July 13, 2025 19:23
@huangyz0918 huangyz0918 merged commit cca1ccb into MLSysOps:main Jul 13, 2025
3 checks passed
@huangyz0918
Copy link
Copy Markdown
Member

Thank you @syangx38 Summer! Nice one 👍

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request lgtm This PR has been approved by a maintainer size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Gemini] Migrate Gemini Integration from Deprecated google-generativeai SDK

5 participants