feat(agents): add metadata parameter to generator completion pipeline

[Hartorn]

Main goal of this PR is to add metadata both as input and output to the BaseGenerator (needed for other subclasses)

Thread metadata: dict[str, Any] | None through the entire completion pipeline (complete, middleware, _call_model) and add a metadata field to Response. Change _call_model return type from tuple to Response for extensibility. Update cursor rule accordingly.

Description

Related Issue

Type of Change

• 📚 Examples / docs / tutorials / dependencies update
• 🔧 Bug fix (non-breaking change which fixes an issue)
• 🥂 Improvement (non-breaking change which improves an existing feature)
• 🚀 New feature (non-breaking change which adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to change)
• 🔐 Security fix

Checklist

• I've read the CODE_OF_CONDUCT.md document.
• I've read the CONTRIBUTING.md guide.
• I've written tests for all new methods and classes that I created.
• I've written the docstring in Google format for all the methods and classes that I used.
• I've updated the pdm.lock running pdm update-lock (only applicable when pyproject.toml has been
  modified)

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly refactors the generator completion pipeline by introducing a flexible metadata parameter. This parameter allows for the seamless flow of arbitrary contextual data, such as trace IDs or tags, both into and out of the model generation process. By unifying the return type of the core _call_model method to a comprehensive Response object that includes this metadata, the changes enhance the system's extensibility and provide a standardized mechanism for handling additional information throughout the AI agent's operations.

Highlights

• Metadata Parameter Introduction: Introduced a metadata parameter of type dict[str, Any] | None across the generator completion pipeline, including complete, middleware, and _call_model methods, to allow passing contextual information.
• Unified Response Type: Changed the return type of the _call_model abstract method from tuple[Message, FinishReason] to a Response object, enhancing extensibility.
• Response Model Enhancement: Added a metadata field to the Response Pydantic model, enabling the return of provider-side information to the caller.
• Pipeline Adaptation: Updated the BaseGenerator and LiteLLMGenerator implementations, as well as middleware components, to correctly handle and propagate the new metadata parameter and the unified Response type.
• Test Suite Updates: Modified existing test cases to align with the updated method signatures and return types, ensuring continued functionality and correctness.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• libs/giskard-agents/.cursor/rules/generator-adapter.mdc
  • Updated the description of the _complete method to reflect its delegation to _call_model.
  • Revised the description of the _call_model method to include the new metadata parameter and its Response return type, clarifying its role in carrying context.
• libs/giskard-agents/src/giskard/agents/generators/_types.py
  • Imported Any from the typing module.
  • Added a metadata field to the Response Pydantic model, initialized as an empty dictionary.
• libs/giskard-agents/src/giskard/agents/generators/base.py
  • Removed the import of FinishReason.
  • Modified the _call_model abstract method signature to accept an optional metadata dictionary and return a Response object.
  • Updated the _complete method to accept an optional metadata dictionary and directly return the result of _call_model.
  • Adjusted the complete method signature to include an optional metadata parameter.
  • Modified the _build_chain's internal _wrapped function to accept and pass metadata to middleware calls.
  • Updated the batch_complete method signature to include an optional metadata parameter and passed it to individual complete calls.
• libs/giskard-agents/src/giskard/agents/generators/litellm_generator.py
  • Replaced the import of FinishReason with Response.
  • Modified the _call_model method signature to accept an optional metadata dictionary and return a Response object.
  • Updated the implementation of _call_model to construct and return a Response object with the message and finish reason.
• libs/giskard-agents/src/giskard/agents/generators/middleware.py
  • Imported Any from the typing module.
  • Updated the NextFn type alias to include an optional metadata dictionary parameter.
  • Modified the call method signatures for CompletionMiddleware, RetryMiddleware, and RateLimiterMiddleware to accept an optional metadata dictionary.
  • Ensured metadata is correctly passed through the middleware chain in RetryMiddleware and RateLimiterMiddleware.
• libs/giskard-agents/tests/test_generator.py
  • Removed the import of FinishReason.
  • Updated the _call_model method signatures in SpyGenerator, TaggingGenerator, and RenamedToolGenerator to accept an optional metadata dictionary and return Response objects.
  • Modified the return statements in mock _call_model implementations to return Response objects.
• libs/giskard-agents/tests/test_generator_retry.py
  • Replaced the import of FinishReason with Response.
  • Updated the _call_model method signature in MockGenerator to accept an optional metadata dictionary and return a Response object.
  • Adjusted mock side_effect assignments to return Response objects instead of tuples.
• libs/giskard-agents/tests/test_serialization.py
  • Imported Any from the typing module.
  • Removed the import of FinishReason.
  • Updated the _call_model method signature in CustomGenerator to accept an optional metadata dictionary and return a Response object.
  • Modified the return statements in CustomGenerator to return Response objects.
• libs/giskard-agents/tests/test_workflow_error_handling.py
  • Imported Any from the typing module.
  • Replaced the import of FinishReason with Response.
  • Updated the _call_model method signature in FailingGenerator to accept an optional metadata dictionary and return a Response object.
  • Modified the return statement in FailingGenerator to return a Response object.
• libs/giskard-agents/tests/test_workflow_output_parsing.py
  • Replaced the import of FinishReason with Response.
  • Updated the _call_model method signature in MockGenerator to accept an optional metadata dictionary and return a Response object.
  • Modified the return statement in MockGenerator to return a Response object.

Activity

• No specific activity (comments, reviews, progress) has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request integrates a metadata parameter across the completion pipeline, from BaseGenerator to CompletionMiddleware and _call_model, to support contextual information like trace IDs and tags. While the overall implementation is clean and consistent, there are critical issues that need addressing. Specifically, the LiteLLMGenerator implementation is incomplete, as it ignores the metadata parameter and fails to populate the response metadata. Furthermore, the introduction of a strictly typed Pydantic Response model for the generator output poses a potential Denial of Service risk if the LLM provider returns an unexpected finish_reason string.

[Hartorn]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a metadata parameter to the generator completion pipeline, allowing for better context propagation and tracing. The change is comprehensive, threading the metadata through the BaseGenerator, middlewares, and litellm implementation. The return type of _call_model has also been refactored to return a Response object, which is a good improvement for extensibility. The changes are well-implemented and all tests have been updated accordingly. I have one suggestion to improve the metadata handling in the litellm generator to ensure that caller-provided metadata is preserved in the final response.