Skip to content

feat(extractor): add with_agent_builder method for custom agent configuration#1357

Closed
atellou wants to merge 4 commits into0xPlaygrounds:mainfrom
atellou:feat/conf-agent-on-extract
Closed

feat(extractor): add with_agent_builder method for custom agent configuration#1357
atellou wants to merge 4 commits into0xPlaygrounds:mainfrom
atellou:feat/conf-agent-on-extract

Conversation

@atellou
Copy link
Copy Markdown
Contributor

@atellou atellou commented Feb 4, 2026

PR Description

Summary

This PR improves the ExtractorBuilder API by adding a closure-based customization hook that lets callers configure the underlying AgentBuilderSimple (e.g. dynamic_context, temperature, max_tokens, etc.) while keeping the extractor’s opinionated defaults (preamble + required submit tool).

Problem

ExtractorBuilder is implemented on top of AgentBuilder, but it only exposed a limited set of configuration methods. This made it awkward/impossible to reuse common agent configuration patterns (especially RAG via dynamic_context) when using extractors, and forced users toward workarounds.

Solution

  • Added ExtractorBuilder::with_agent_builder:
    • Signature: FnOnce(AgentBuilderSimple<M>) -> AgentBuilderSimple<M>
    • Enables configuring agent builder options inside the extractor builder chain.
  • Updated the gemini_rag example to use .with_agent_builder(|agent| agent.dynamic_context(...)) to attach vector search context to the extractor.

Tests

  • Added unit test extractor_builder_with_agent_builder_applies_config validating that with_agent_builder successfully applies configuration (e.g. temperature, max_tokens) to the built extractor’s inner Agent.

Examples

  • Added / updated example: rig-core/examples/gemini_rag.rs
    • Demonstrates building embeddings for questionnaire docs, adding dynamic RAG context, and extracting structured questionnaire responses using the Gemini provider.

Notes

  • cargo test -p rig-core extractor_builder_with_agent_builder_applies_config --features derive passes.

Felipe Tello Urrea added 2 commits February 4, 2026 09:44
…guration

- Format mime_guess dependency declaration in Cargo.toml
- Add gemini_rag example demonstrating RAG with structured extraction
@joshua-mo-143
Copy link
Copy Markdown
Contributor

To be honest we should probably really just re-combine the agent builders back together and have the tools thing as an enum so we don't have to deal with this. I'll open a sep issue to solve this

@joshua-mo-143
Copy link
Copy Markdown
Contributor

If you merge main back into your branch, it should now be much easier to AgentBuilder<M, WithBuilderTools> in the config rather than using one single config.

I think at the moment it will be difficult to use tool server handles with this type currently (which means that using the WithToolServer type state variant will be impossible with extractors) because the submit tool basically needs to be guaranteed to exist, but I think this is a reasonable tradeoff for now.

@atellou
Copy link
Copy Markdown
Contributor Author

atellou commented Feb 6, 2026

Mmmmm got it, thanks for the fast response. In that case, shouldn't the toolchoice function be removed from the ExtractorBuilder? Also, for my use case I only need to enable RAG. If we enable the dynamic context function, would that be ok?

    /// Add dynamic context (RAG) to the extractor.
    ///
    /// On each prompt, `sample` documents will be retrieved from the index based on the RAG text
    /// and inserted in the request.
    pub fn dynamic_context(
        mut self,
        sample: usize,
        dynamic_context: impl VectorStoreIndexDyn + Send + Sync + 'static,
    ) -> Self {
        self.agent_builder = self.agent_builder.dynamic_context(sample, dynamic_context);
        self
    }

@joshua-mo-143
Copy link
Copy Markdown
Contributor

Mmmmm got it, thanks for the fast response. In that case, shouldn't the toolchoice function be removed from the ExtractorBuilder? Also, for my use case I only need to enable RAG. If we enable the dynamic context function, would that be ok?

    /// Add dynamic context (RAG) to the extractor.
    ///
    /// On each prompt, `sample` documents will be retrieved from the index based on the RAG text
    /// and inserted in the request.
    pub fn dynamic_context(
        mut self,
        sample: usize,
        dynamic_context: impl VectorStoreIndexDyn + Send + Sync + 'static,
    ) -> Self {
        self.agent_builder = self.agent_builder.dynamic_context(sample, dynamic_context);
        self
    }

Yeah that's fine.

@atellou atellou closed this Mar 2, 2026
@atellou atellou deleted the feat/conf-agent-on-extract branch March 2, 2026 16:49
@atellou
Copy link
Copy Markdown
Contributor Author

atellou commented Mar 3, 2026

moved to this #1456

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants