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

feat(gemini): trace google gemini Integration #10503

Merged
merged 25 commits into from
Sep 11, 2024
Merged

Conversation

Yun-Kim
Copy link
Contributor

@Yun-Kim Yun-Kim commented Sep 3, 2024

Adds support for tracing Google's Gemini Python SDK generate_content and generate_content_async methods (also indirectly traces send_message()).

This PR also moves _get_attr() convenience helper from the anthropic integration to the shared ddtrace.llmobs._utils.py module.

Features:

  • Traces generate_content/generate_content_async() and indirectly send_message/send_message_async() which uses generate_content under the hood
  • Handles streamed responses
  • Captures input/outputs (truncation included, function call/response messages included), token usage, input configuration options, and model name.

Testing:

The Gemini Python SDK does not use requests/httpx to submit requests but instead via GRPC. The vcrpy testing framework does not capture this so I wrote a mock client to store/return mock responses instead of sending actual requests to Google. This testing framework was inspired by the Gemini Python SDK's testing strategy.

Next steps (future PRs)

  • Wrap send_message() not to trace, but to attach a hash/session_id for all subsequent send_message calls
  • Integrate with LLMObs

Checklist

  • PR author has checked that all the criteria below are met
  • The PR description includes an overview of the change
  • The PR description articulates the motivation for the change
  • The change includes tests OR the PR description describes a testing strategy
  • The PR description notes risks associated with the change, if any
  • Newly-added code is easy to change
  • The change follows the library release note guidelines
  • The change includes or references documentation updates if necessary
  • Backport labels are set (if applicable)

Reviewer Checklist

  • Reviewer has checked that all the criteria below are met
  • Title is accurate
  • All changes are related to the pull request's stated goal
  • Avoids breaking API changes
  • Testing strategy adequately addresses listed risks
  • Newly-added code is easy to change
  • Release note makes sense to a user of the library
  • If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment
  • Backport labels are set in a manner that is consistent with the release branch maintenance policy

Copy link
Contributor

github-actions bot commented Sep 3, 2024

CODEOWNERS have been resolved as:

.riot/requirements/1e15a25.txt                                          @DataDog/apm-python
.riot/requirements/1f54e6b.txt                                          @DataDog/apm-python
.riot/requirements/e8247d6.txt                                          @DataDog/apm-python
.riot/requirements/ebe4ea5.txt                                          @DataDog/apm-python
ddtrace/contrib/google_generativeai/__init__.py                         @DataDog/ml-observability
ddtrace/contrib/internal/google_generativeai/_utils.py                  @DataDog/ml-observability
ddtrace/contrib/internal/google_generativeai/patch.py                   @DataDog/ml-observability
ddtrace/llmobs/_integrations/gemini.py                                  @DataDog/ml-observability
releasenotes/notes/feat-google-gemini-d5ee30b1d711bc08.yaml             @DataDog/apm-python
tests/contrib/google_generativeai/__init__.py                           @DataDog/ml-observability
tests/contrib/google_generativeai/conftest.py                           @DataDog/ml-observability
tests/contrib/google_generativeai/test_data/apple.jpg                   @DataDog/ml-observability
tests/contrib/google_generativeai/test_google_generativeai.py           @DataDog/ml-observability
tests/contrib/google_generativeai/test_google_generativeai_patch.py     @DataDog/ml-observability
tests/contrib/google_generativeai/utils.py                              @DataDog/ml-observability
tests/snapshots/tests.contrib.google_generativeai.test_google_generativeai.test_gemini_completion.json  @DataDog/apm-python
tests/snapshots/tests.contrib.google_generativeai.test_google_generativeai.test_gemini_completion_error.json  @DataDog/apm-python
tests/snapshots/tests.contrib.google_generativeai.test_google_generativeai.test_gemini_completion_image.json  @DataDog/apm-python
tests/snapshots/tests.contrib.google_generativeai.test_google_generativeai.test_gemini_completion_multiple_messages.json  @DataDog/apm-python
tests/snapshots/tests.contrib.google_generativeai.test_google_generativeai.test_gemini_completion_stream.json  @DataDog/apm-python
tests/snapshots/tests.contrib.google_generativeai.test_google_generativeai.test_gemini_completion_system_prompt.json  @DataDog/apm-python
tests/snapshots/tests.contrib.google_generativeai.test_google_generativeai.test_gemini_completion_tool_stream.json  @DataDog/apm-python
tests/snapshots/tests.contrib.google_generativeai.test_google_generativeai.test_gemini_tool_chat_completion.json  @DataDog/apm-python
tests/snapshots/tests.contrib.google_generativeai.test_google_generativeai.test_gemini_tool_completion.json  @DataDog/apm-python
.github/CODEOWNERS                                                      @DataDog/python-guild @DataDog/apm-core-python
.gitlab/tests/llmobs.yml                                                @DataDog/ml-observability
ddtrace/_monkey.py                                                      @DataDog/apm-core-python
ddtrace/contrib/internal/anthropic/_streaming.py                        @DataDog/ml-observability
ddtrace/contrib/internal/anthropic/patch.py                             @DataDog/ml-observability
ddtrace/contrib/internal/anthropic/utils.py                             @DataDog/ml-observability
ddtrace/llmobs/_integrations/__init__.py                                @DataDog/ml-observability
ddtrace/llmobs/_integrations/anthropic.py                               @DataDog/ml-observability
ddtrace/llmobs/_utils.py                                                @DataDog/ml-observability
docs/index.rst                                                          @DataDog/python-guild
docs/integrations.rst                                                   @DataDog/python-guild
docs/spelling_wordlist.txt                                              @DataDog/python-guild
riotfile.py                                                             @DataDog/apm-python
tests/.suitespec.json                                                   @DataDog/python-guild @DataDog/apm-core-python

@pr-commenter
Copy link

pr-commenter bot commented Sep 3, 2024

Benchmarks

Benchmark execution time: 2024-09-11 20:02:53

Comparing candidate commit eecda47 in PR branch yunkim/gemini-integration with baseline commit 150beb7 in branch main.

Found 0 performance improvements and 0 performance regressions! Performance is the same for 353 metrics, 47 unstable metrics.

@datadog-dd-trace-py-rkomorn
Copy link

datadog-dd-trace-py-rkomorn bot commented Sep 9, 2024

Datadog Report

Branch report: yunkim/gemini-integration
Commit report: eecda47
Test service: dd-trace-py

✅ 0 Failed, 9717 Passed, 355 Skipped, 2h 9m 47.19s Total duration (1m 35.08s time saved)

@Yun-Kim Yun-Kim marked this pull request as ready for review September 10, 2024 15:42
@Yun-Kim Yun-Kim requested review from a team as code owners September 10, 2024 15:42
@Yun-Kim Yun-Kim requested a review from a team as a code owner September 10, 2024 15:42
Copy link
Contributor

@erikayasuda erikayasuda left a comment

Choose a reason for hiding this comment

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

Just some nits and clarifying questions

riotfile.py Outdated Show resolved Hide resolved
docs/spelling_wordlist.txt Show resolved Hide resolved
@Yun-Kim Yun-Kim merged commit c0fd013 into main Sep 11, 2024
506 of 507 checks passed
@Yun-Kim Yun-Kim deleted the yunkim/gemini-integration branch September 11, 2024 20:45
github-actions bot pushed a commit that referenced this pull request Sep 11, 2024
Adds support for tracing Google's Gemini Python SDK `generate_content`
and `generate_content_async` methods (also indirectly traces
`send_message()`).

This PR also moves `_get_attr()` convenience helper from the anthropic
integration to the shared `ddtrace.llmobs._utils.py` module.

### Features:
- Traces `generate_content/generate_content_async()` and indirectly
`send_message/send_message_async()` which uses `generate_content` under
the hood
- Handles streamed responses
- Captures input/outputs (truncation included, function call/response
messages included), token usage, input configuration options, and model
name.

### Testing:
The Gemini Python SDK does not use requests/httpx to submit requests but
instead via GRPC. The vcrpy testing framework does not capture this so I
wrote a mock client to store/return mock responses instead of sending
actual requests to Google. This testing framework was inspired by the
[Gemini Python SDK's testing
strategy](https://github.com/google-gemini/generative-ai-python/blob/9407dcde5666ba58831227d3acf8bd0e3f3b4f81/tests/test_generative_models.py#L41).

### Next steps (future PRs)
- Wrap `send_message()` not to trace, but to attach a hash/session_id
for all subsequent send_message calls
- Integrate with LLMObs

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

(cherry picked from commit c0fd013)
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.

5 participants