feat: observe decorators

[abhishekg999]

📝 Summary

• 1. ...

✅ Checklist

• Docs updated (if necessary)
• Changelogs are updated (if necessary)

[gemini-code-assist]

Summary of Changes

Hello @abhishekg999, 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 extends the tracing capabilities by enabling the observation of Python generators, both synchronous and asynchronous. This feature ensures that the execution flow of generators, including individual yielded items and any exceptions, is properly captured and integrated into the tracing system. The changes also include improvements to context management within decorators and robust error handling, alongside dedicated tests to verify the new functionality.

Highlights

• Generator Observation: Introduced the ability to observe both synchronous and asynchronous Python generators using the @observe decorator, allowing for detailed tracing of each yielded item.
• Context Preservation: Implemented contextvars to ensure that the execution context (e.g., customer ID, agent ID) is correctly preserved across generator yields, maintaining trace continuity.
• Improved Error Handling and Span Management: Enhanced error handling within observed generators, ensuring exceptions are properly recorded in spans. Spans for generators are now explicitly ended, providing more accurate trace durations.
• Refactored Decorator Logic: The observe and agent decorators in base_tracer.py have been refactored for cleaner code, including a new _serialize helper for consistent attribute serialization and a build_context helper for agent.
• New Test Coverage: Added a comprehensive test suite (test_generator.py) to validate the new generator observation features, covering basic functionality, context preservation, exception handling, and partial consumption.

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 introduces a significant feature: support for observing synchronous and asynchronous generators within the tracer. This is implemented by wrapping generator objects and creating child spans for each yielded item, which is a solid approach. The PR also includes comprehensive tests for this new functionality, covering various scenarios like context preservation and exception handling. Additionally, there are several good refactorings, such as simplifying the agent decorator's context creation and improving type hints.

My review includes a few suggestions:

• A potential bug in _ObservedAsyncGenerator where context propagation might be lost.
• A simplification for creating child span names within the generator wrappers.
• A recommendation to remove an unnecessary type: ignore comment.
• Suggestions to strengthen a couple of tests in test_tracer.py that were weakened during refactoring.

Overall, this is a great contribution that enhances the tracing capabilities.

[propel-code-bot]

[CriticalError]

This FireworksTrainer creates its own SpanStore and InMemorySpanExporter. However, the spans generated by @self.tracer.observe within this class will be processed by the exporter configured in the tracer instance passed during initialization, not this new, local one.

As a result, self.span_store will remain empty, and _extract_message_history_from_spans will fail to retrieve any trace data, likely breaking the training logic that depends on it.

To fix this, the trainer should either receive the SpanStore as a constructor argument or have a mechanism to access the SpanStore from the provided tracer instance.

Context for Agents

[**CriticalError**]

This `FireworksTrainer` creates its own `SpanStore` and `InMemorySpanExporter`. However, the spans generated by `@self.tracer.observe` within this class will be processed by the exporter configured in the `tracer` instance passed during initialization, not this new, local one.

As a result, `self.span_store` will remain empty, and `_extract_message_history_from_spans` will fail to retrieve any trace data, likely breaking the training logic that depends on it.

To fix this, the trainer should either receive the `SpanStore` as a constructor argument or have a mechanism to access the `SpanStore` from the provided `tracer` instance.

File: src/judgeval/v1/trainers/fireworks_trainer.py
Line: 48

[propel-code-bot]

[BestPractice]

The type hint for the client parameter is Optional[JudgmentSyncClient] = None, but the implementation raises a ValueError if it's None. This makes the parameter mandatory in practice. The type hint should be updated to reflect this requirement to avoid confusion for developers using this class.

Suggested Change

Suggested change

	config: TrainerConfig,
	trainable_model: TrainableModel,
	tracer: Tracer,
	project_name: Optional[str] = None,
	client: Optional["JudgmentSyncClient"] = None,
	client: Optional[JudgmentSyncClient] = None,
	config: TrainerConfig,
	trainable_model: TrainableModel,
	tracer: Tracer,
	project_name: Optional[str] = None,
	client: JudgmentSyncClient,

⚡ Committable suggestion

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

Context for Agents

[**BestPractice**]

The type hint for the `client` parameter is `Optional[JudgmentSyncClient] = None`, but the implementation raises a `ValueError` if it's `None`. This makes the parameter mandatory in practice. The type hint should be updated to reflect this requirement to avoid confusion for developers using this class.

<details>
<summary>Suggested Change</summary>

```suggestion
        config: TrainerConfig,
        trainable_model: TrainableModel,
        tracer: Tracer,
        project_name: Optional[str] = None,
        client: JudgmentSyncClient,
```

⚡ **Committable suggestion**

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

</details>

File: src/judgeval/v1/trainers/fireworks_trainer.py
Line: 45

[alanzhang25]

very nice! just the two comments

[propel-code-bot]

✔️ Propel has finished reviewing this change.