[feat] Add retrieval_info and allow full environment references into traces#4469
Draft
jp-agenta wants to merge 5 commits into
Draft
[feat] Add retrieval_info and allow full environment references into traces#4469jp-agenta wants to merge 5 commits into
jp-agenta wants to merge 5 commits into
Conversation
Previously RetrievalInfo lived under api/oss/src/apis/fastapi/git/ and the routers assembled it directly, including making duplicate environment-revision DB calls to gather environment-side references. This commit moves the DTO and helpers to api/oss/src/core/git/ for symmetry with ResolutionInfo, threads retrieval_info through the services as the third element of a (revision, resolution_info, retrieval_info) tuple, and lets the routers act as pass-throughs. Highlights: - core/git/dtos.py: RetrievalInfo Pydantic model - core/git/utils.py: build_retrieval_info / revision_references helpers - Six services updated to return the 3-tuple; environment-backed paths merge environment + target references inside the service, no router-level DB call - Queries and testsets gain thin retrieve_*_revision wrappers for symmetry - /revisions/retrieve and /revisions/resolve responses carry retrieval_info - Unit + acceptance test coverage for all six entities, both retrieve and resolve, direct and environment-backed No request-field renames in this commit; only the internal refactor and the new retrieval_info on retrieve/resolve responses.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
dosubot
Bot
added
size:XL
4 tasks
junaway
changed the title
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR adds retrieval provenance metadata to revision retrieval responses across multiple entity types (applications, environments, evaluators, workflows, queries, testsets). A new ChangesRetrieval Info Feature
Python SDK Resolver Middleware Integration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
|
@mmabrouk |
Contributor
Railway Preview Environment
Updated at 2026-05-28T07:36:42.357Z |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
api/oss/src/core/applications/service.py (1)
595-627:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDerive the environment selector key before lookup.
This branch now returns
None, None, Nonewheneverkeyis omitted, even though the API contract inapi/oss/src/apis/fastapi/applications/models.py(Lines 403-409) says the server derives it fromapplication_ref. That breaks documented environment-backed retrievals like{application_slug}.revision.Suggested fix
environment_retrieval_info: Optional[RetrievalInfo] = None is_environment_backed = bool( environment_ref or environment_variant_ref or environment_revision_ref ) if is_environment_backed: + if key is None and application_ref and application_ref.slug: + key = f"{application_ref.slug}.revision" + environments_service = self.workflows_service.environments_service if not environments_service: return None, None, Noneapi/oss/src/core/workflows/service.py (1)
1153-1233: 🛠️ Refactor suggestion | 🟠 Major | 🏗️ Heavy liftReplace the positional 3-tuple with a typed result DTO.
This expands a brittle unpack-by-position contract across layers, and callers are already discarding elements positionally in
_ensure_request_revision(). A dedicatedBaseModelresult would make the contract self-describing and safer to evolve.As per coding guidelines, Service methods must return typed DTOs (Pydantic
BaseModelsubclasses), not raw dicts, tuples, orAny.
🧹 Nitpick comments (2)
api/oss/tests/pytest/unit/applications/test_router.py (1)
195-206: ⚡ Quick winAssert the router actually delegates to
applications_service.retrieve_application_revision.This test validates the returned payload, but not the service invocation itself. Add an explicit await assertion to protect the router-to-service contract.
Proposed test hardening
assert response.retrieval_info is not None assert response.retrieval_info.key == "demo-app.revision" assert response.retrieval_info.references["environment"].id == environment_id assert response.retrieval_info.references["environment_revision"].version == "7" assert ( response.retrieval_info.references["application_revision"].id == application_revision_id ) + applications_service.retrieve_application_revision.assert_awaited_once() # The router no longer does its own env round-trip; the applications service # owns environment lookup and reference assembly. environments_service.retrieve_environment_revision.assert_not_awaited()api/oss/src/apis/fastapi/queries/models.py (1)
3-3: ⚡ Quick winDocument the new
retrieval_infofield in the public schema.Line 152 adds a new response field, but unlike the other
retrieval_infoadditions in this PR it has noField(...)metadata, so the OpenAPI schema won't explain what clients should expect here.Suggested patch
-from pydantic import BaseModel +from pydantic import BaseModel, Field @@ class QueryRevisionResponse(BaseModel): count: int = 0 query_revision: Optional[QueryRevision] = None - retrieval_info: Optional[RetrievalInfo] = None + retrieval_info: Optional[RetrievalInfo] = Field( + default=None, + description="References used to retrieve the top-level revision.", + )Also applies to: 149-152
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: c0dd2d2d-ca02-4739-a11b-d11dbc5678e5
📒 Files selected for processing (45)
api/oss/src/apis/fastapi/applications/models.pyapi/oss/src/apis/fastapi/applications/router.pyapi/oss/src/apis/fastapi/environments/models.pyapi/oss/src/apis/fastapi/environments/router.pyapi/oss/src/apis/fastapi/evaluators/models.pyapi/oss/src/apis/fastapi/evaluators/router.pyapi/oss/src/apis/fastapi/legacy_variants/router.pyapi/oss/src/apis/fastapi/queries/models.pyapi/oss/src/apis/fastapi/queries/router.pyapi/oss/src/apis/fastapi/testsets/models.pyapi/oss/src/apis/fastapi/testsets/router.pyapi/oss/src/apis/fastapi/workflows/models.pyapi/oss/src/apis/fastapi/workflows/router.pyapi/oss/src/core/applications/service.pyapi/oss/src/core/environments/service.pyapi/oss/src/core/evaluators/service.pyapi/oss/src/core/git/dtos.pyapi/oss/src/core/git/utils.pyapi/oss/src/core/queries/service.pyapi/oss/src/core/testsets/service.pyapi/oss/src/core/workflows/service.pyapi/oss/tests/pytest/acceptance/applications/test_application_retrieval_info.pyapi/oss/tests/pytest/acceptance/environments/test_environment_retrieval_info.pyapi/oss/tests/pytest/acceptance/evaluators/test_evaluator_retrieval_info.pyapi/oss/tests/pytest/acceptance/queries/test_query_retrieval_info.pyapi/oss/tests/pytest/acceptance/testsets/test_testset_retrieval_info.pyapi/oss/tests/pytest/acceptance/workflows/test_workflow_retrieval_info.pyapi/oss/tests/pytest/unit/applications/__init__.pyapi/oss/tests/pytest/unit/applications/test_retrieve_resolve.pyapi/oss/tests/pytest/unit/applications/test_router.pyapi/oss/tests/pytest/unit/environments/__init__.pyapi/oss/tests/pytest/unit/environments/test_retrieve_resolve.pyapi/oss/tests/pytest/unit/evaluators/test_retrieve_resolve.pyapi/oss/tests/pytest/unit/git/__init__.pyapi/oss/tests/pytest/unit/git/test_retrieval_info_utils.pyapi/oss/tests/pytest/unit/legacy_variants/test_configs_fetch.pyapi/oss/tests/pytest/unit/queries/__init__.pyapi/oss/tests/pytest/unit/queries/test_retrieve.pyapi/oss/tests/pytest/unit/testsets/__init__.pyapi/oss/tests/pytest/unit/testsets/test_retrieve.pyapi/oss/tests/pytest/unit/workflows/__init__.pyapi/oss/tests/pytest/unit/workflows/test_flag_ownership.pyapi/oss/tests/pytest/unit/workflows/test_retrieve_resolve.pysdks/python/agenta/sdk/middlewares/running/resolver.pysdks/python/oss/tests/pytest/utils/test_resolver_middleware.py
Resolved conflicts by keeping both import sets (RetrievalInfo/build_retrieval_info from HEAD and validate_* functions from incoming) across all service and router files. In environments/service.py, adopted the incoming simplification of retrieve_environment_revision to delegate to fetch_environment_revision (which includes the full validation set), restored build_retrieval_info call and fixed the incorrect 2-tuple return to 3-tuple. In the SDK test file, kept HEAD's test_stores_retrieval_references_on_tracing_context method and appended the incoming TestResolverMiddlewareTracingParameters class. Dropped unused VariantForkError import from all three routers (now handled by handle_git_exceptions). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Refactor of the six revisioned services (workflows, applications, evaluators, environments, queries, testsets) to thread a new RetrievalInfo (third element of the return tuple) alongside the existing (revision, resolution_info). Environment-backed lookups now merge environment + target references inside the service, removing a duplicate retrieve_environment_revision call in the routers. retrieval_info is exposed on /revisions/retrieve and /revisions/resolve responses for all six entities. The SDK resolver middleware learns a resolve_references_with_info variant that merges returned references into the tracing context.
Changes:
- Move
RetrievalInfoDTO intocore/git, addbuild_retrieval_info/revision_referencesutilities, and addretrieval_infofields to the relevant FastAPI response models. - Update all six services and their routers (plus legacy-variants pass-throughs) to produce/consume the 3-tuple, and add
retrieve_query_revision/retrieve_testset_revisionwrappers. - Extend SDK resolver middleware to merge server-provided retrieval references into
TracingContext, and add unit + acceptance tests across all six entities.
Reviewed changes
Copilot reviewed 39 out of 45 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| api/oss/src/core/git/dtos.py | Adds RetrievalInfo DTO with references + key. |
| api/oss/src/core/git/utils.py | New build_retrieval_info / revision_references helpers. |
| api/oss/src/core/workflows/service.py | Returns (revision, resolution_info, retrieval_info); merges env refs when env-backed. |
| api/oss/src/core/applications/service.py | Same triple-tuple refactor for applications. |
| api/oss/src/core/evaluators/service.py | Same triple-tuple refactor for evaluators. |
| api/oss/src/core/environments/service.py | Returns triple and removes stale commented log. |
| api/oss/src/core/queries/service.py | New retrieve_query_revision wrapper emitting retrieval_info. |
| api/oss/src/core/testsets/service.py | New retrieve_testset_revision wrapper emitting retrieval_info. |
| api/oss/src/apis/fastapi/{workflows,applications,evaluators,environments,queries,testsets}/router.py | Pass retrieval_info through; resolve endpoints synthesize it via build_retrieval_info. |
| api/oss/src/apis/fastapi/{workflows,applications,evaluators,environments,queries,testsets}/models.py | Add retrieval_info field to response models. |
| api/oss/src/apis/fastapi/legacy_variants/router.py | Updated all unpacks to the 3-tuple shape. |
| sdks/python/agenta/sdk/middlewares/running/resolver.py | New resolve_references_with_info + _merge_tracing_references; existing resolve_references becomes a thin wrapper. |
| api/oss/tests/pytest/unit/git/test_retrieval_info_utils.py | Unit tests for helper merging and key handling. |
| api/oss/tests/pytest/unit/{applications,evaluators,environments,workflows,queries,testsets}/test_retrieve_resolve.py (+ test_router.py, test_flag_ownership.py) | New/updated service-level unit tests covering direct + env-backed paths. |
| api/oss/tests/pytest/unit/legacy_variants/test_configs_fetch.py | Updated mock return values to 3-tuple. |
| api/oss/tests/pytest/acceptance/{applications,evaluators,environments,workflows,queries,testsets}/test_*_retrieval_info.py | End-to-end acceptance coverage for the new field. |
| sdks/python/oss/tests/pytest/utils/test_resolver_middleware.py | Asserts retrieval references are merged onto the tracing context. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Contributor
|
@CodeRabbit review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
api/oss/src/core/environments/service.py (1)
796-810: 🛠️ Refactor suggestion | 🟠 Major | 🏗️ Heavy liftReturn a typed result model here instead of another positional tuple.
This new 3-tuple is already forcing positional
_unpacking in downstream callers, and that coupling will keep spreading. Please wrapenvironment_revision,resolution_info, andretrieval_infoin a Pydantic DTO and return that from the service.As per coding guidelines, "Service methods must return typed DTOs (Pydantic `BaseModel` subclasses), not raw dicts, tuples, or `Any`."📦 Suggested direction
+class EnvironmentRevisionRetrieveResult(BaseModel): + environment_revision: Optional[EnvironmentRevision] = None + resolution_info: Optional[ResolutionInfo] = None + retrieval_info: Optional[RetrievalInfo] = None + async def retrieve_environment_revision( self, *, project_id: UUID, ... - ) -> tuple[ - Optional[EnvironmentRevision], - Optional[ResolutionInfo], - Optional[RetrievalInfo], - ]: + ) -> EnvironmentRevisionRetrieveResult:api/oss/src/core/applications/service.py (1)
591-610: 🛠️ Refactor suggestion | 🟠 Major | 🏗️ Heavy liftUse a typed DTO for this retrieve contract instead of a positional tuple.
This API now carries three values and has already started leaking positional unpacking into callers. Please return a Pydantic result model with named fields for the revision, resolution metadata, and retrieval metadata.
As per coding guidelines, "Service methods must return typed DTOs (Pydantic `BaseModel` subclasses), not raw dicts, tuples, or `Any`."📦 Suggested direction
+class ApplicationRevisionRetrieveResult(BaseModel): + application_revision: Optional[ApplicationRevision] = None + resolution_info: Optional[ResolutionInfo] = None + retrieval_info: Optional[RetrievalInfo] = None + async def retrieve_application_revision( self, *, project_id: UUID, ... - ) -> tuple[ - Optional[ApplicationRevision], - Optional[ResolutionInfo], - Optional[RetrievalInfo], - ]: + ) -> ApplicationRevisionRetrieveResult:
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 51a706b5-471b-4d0a-be11-d35fd2868e2b
📒 Files selected for processing (27)
api/oss/src/apis/fastapi/applications/models.pyapi/oss/src/apis/fastapi/applications/router.pyapi/oss/src/apis/fastapi/environments/models.pyapi/oss/src/apis/fastapi/environments/router.pyapi/oss/src/apis/fastapi/evaluators/models.pyapi/oss/src/apis/fastapi/evaluators/router.pyapi/oss/src/apis/fastapi/legacy_variants/router.pyapi/oss/src/apis/fastapi/queries/models.pyapi/oss/src/apis/fastapi/queries/router.pyapi/oss/src/apis/fastapi/testsets/models.pyapi/oss/src/apis/fastapi/testsets/router.pyapi/oss/src/apis/fastapi/workflows/models.pyapi/oss/src/apis/fastapi/workflows/router.pyapi/oss/src/core/applications/service.pyapi/oss/src/core/environments/service.pyapi/oss/src/core/evaluators/service.pyapi/oss/src/core/queries/service.pyapi/oss/src/core/testsets/service.pyapi/oss/src/core/workflows/service.pyapi/oss/tests/pytest/acceptance/applications/test_application_retrieval_info.pyapi/oss/tests/pytest/acceptance/environments/test_environment_retrieval_info.pyapi/oss/tests/pytest/acceptance/evaluators/test_evaluator_retrieval_info.pyapi/oss/tests/pytest/acceptance/queries/test_query_retrieval_info.pyapi/oss/tests/pytest/acceptance/testsets/test_testset_retrieval_info.pyapi/oss/tests/pytest/acceptance/workflows/test_workflow_retrieval_info.pysdks/python/agenta/sdk/middlewares/running/resolver.pysdks/python/oss/tests/pytest/utils/test_resolver_middleware.py
🚧 Files skipped from review as they are similar to previous changes (21)
- api/oss/src/apis/fastapi/workflows/models.py
- api/oss/src/apis/fastapi/applications/models.py
- api/oss/src/apis/fastapi/evaluators/models.py
- api/oss/src/apis/fastapi/environments/models.py
- api/oss/src/apis/fastapi/queries/models.py
- api/oss/src/apis/fastapi/evaluators/router.py
- sdks/python/oss/tests/pytest/utils/test_resolver_middleware.py
- api/oss/src/apis/fastapi/workflows/router.py
- api/oss/src/apis/fastapi/applications/router.py
- api/oss/src/apis/fastapi/queries/router.py
- api/oss/src/core/testsets/service.py
- api/oss/tests/pytest/acceptance/testsets/test_testset_retrieval_info.py
- api/oss/tests/pytest/acceptance/environments/test_environment_retrieval_info.py
- api/oss/src/core/queries/service.py
- api/oss/tests/pytest/acceptance/workflows/test_workflow_retrieval_info.py
- api/oss/src/core/evaluators/service.py
- api/oss/tests/pytest/acceptance/evaluators/test_evaluator_retrieval_info.py
- api/oss/tests/pytest/acceptance/queries/test_query_retrieval_info.py
- api/oss/tests/pytest/acceptance/applications/test_application_retrieval_info.py
- sdks/python/agenta/sdk/middlewares/running/resolver.py
- api/oss/src/core/workflows/service.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
RetrievalInfo(DTO + helpers) fromapi/oss/src/apis/fastapi/git/toapi/oss/src/core/git/, for symmetry withResolutionInfo.retrieval_infothrough the six revisioned services as the third element of a(revision, resolution_info, retrieval_info)tuple; routers become pass-throughs.retrieve_environment_revisionDB call the routers previously made./revisions/retrieveand/revisions/resolveresponses now carryretrieval_infofor all six entities (workflows, applications, evaluators, environments, queries, testsets).retrieve_*_revisionwrappers for symmetry with the other entities.This is the first of two stacked PRs. It contains no request-field renames — only the internal refactor and the new
retrieval_infoon retrieve/resolve responses. The field-name unification follows in the stacked PR.Test plan
api/oss/tests/pytest/unit/git/test_retrieval_info_utils.py+ per-entitytest_retrieve_resolve.pytest_*_retrieval_info.py(direct + environment-backed + resolve)🤖 Generated with Claude Code