Repository navigation
feat: add stdlib building blocks for custom LLM providers - #4164
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
⏭️ Performance benchmarks were skippedPerf benchmarks (CodSpeed) are opt-in on pull requests — they no longer run on every push. They always run automatically after merge to To run them on this PR, do any of the following, then push a commit (or re-run CI):
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds JSON-path navigation, JSON Schema lowering, SAP parsing, prompt-to-text conversion, and unknown-error builtins. It also adds feature-gated native WebSocket support with registry-backed streams, plus updated native artifact size baselines. ChangesProvider-oriented standard library surfaces
Native WebSocket support
Artifact size gate updates
Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant ProviderTest
participant baml_schema
participant IoNamespaceSchema
participant schema_json_schema
ProviderTest->>baml_schema: provider_schema()
baml_schema->>IoNamespaceSchema: json_schema(ProviderSchemaExample)
IoNamespaceSchema->>schema_json_schema: lower RuntimeTy
schema_json_schema-->>IoNamespaceSchema: JSON Schema value
IoNamespaceSchema-->>ProviderTest: reflected JSON schema
ProviderTest->>baml_schema: baml.json.path(schema, selector)
baml_schema-->>ProviderTest: typed schema value
sequenceDiagram
participant ProviderTest
participant baml_sap
participant baml_llm
participant prompt_to_text
ProviderTest->>baml_sap: parse<int>("42")
baml_sap->>baml_llm: __sap_parse_final(text, stream cache)
baml_llm-->>ProviderTest: parsed integer
ProviderTest->>prompt_to_text: prompt_to_text(PromptAst)
prompt_to_text-->>ProviderTest: flattened text
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
baml_language/crates/sys_native/src/io_impls.rs (2)
2486-2488: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winBinary frames are silently lossy-decoded to string, despite the "text-oriented" contract.
ws.bamldocuments this surface as "text-oriented" andnext()as returning "the next text frame", butMessage::Binarypayloads are silently converted viafrom_utf8_lossy, which replaces invalid UTF-8 with replacement characters — corrupting data instead of surfacing an error for callers that don't expect binary frames.♻️ Proposed fix: treat unexpected binary frames as an error
- Some(Ok(Message::Binary(bytes))) => { - return Ok(Some(String::from_utf8_lossy(&bytes).into_owned())); - } + Some(Ok(Message::Binary(bytes))) => { + return Err(VmBamlError::Io { + message: format!( + "received unexpected binary WebSocket frame ({} bytes) on a text-oriented stream", + bytes.len() + ), + } + .into()); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@baml_language/crates/sys_native/src/io_impls.rs` around lines 2486 - 2488, Update the Message::Binary branch in the websocket next() handling to return an error for unexpected binary frames instead of converting bytes with String::from_utf8_lossy. Preserve the existing successful text-frame behavior and use the module’s established error type or constructor to report the binary-frame condition.
2525-2542: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win
close()never sends a WebSocket Close frame.It only removes the registry entry and lets the sink/source drop; the underlying TCP connection is torn down without a WS-level close handshake, so the remote peer sees an abnormal closure instead of a clean one. Consider sending a
Message::Close(None)(and flushing) on the sink before/while removing from the registry.Separately, note that
close()only unregisters the resource — anysend()/next()call that already resolved its ownArcclone viaget_ws_stream()before this runs will keep operating on the still-alive connection until it naturally completes, soclose()doesn't forcibly interrupt in-flight operations.♻️ Proposed fix to send a Close frame
fn close( &self, _heap: &Arc<BexHeap>, _call_id: CallId, stream: owned::ws::WsStream, _ctx: &SysOpContext, ) -> SysOpOutput<()> { use bex_resource_types::ResourceRegistryRef; - if let Ok(handle) = stream - ._handle - .downcast::<bex_resource_types::ResourceHandle>() - { - crate::registry::REGISTRY.remove(handle.key()); - } - SysOpOutput::ok(()) + if let Ok(handle) = stream + ._handle + .downcast::<bex_resource_types::ResourceHandle>() + { + if let Some((sink, _)) = crate::registry::REGISTRY.get_ws_stream(handle.key()) { + use futures::SinkExt; + let _ = sink.lock().await.send(tokio_tungstenite::tungstenite::Message::Close(None)).await; + } + crate::registry::REGISTRY.remove(handle.key()); + } + SysOpOutput::ok(()) }(sketch only —
closewould need to become async / returnSysOpOutput::async_opfor the.await, mirroringsend/next.)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@baml_language/crates/sys_native/src/io_impls.rs` around lines 2525 - 2542, Update the WebSocket close flow in close to become asynchronous as needed, send Message::Close(None) through the stream’s sink, and flush it before removing the resource from REGISTRY. Preserve the existing successful SysOpOutput result and resource unregistration behavior; do not attempt to forcibly interrupt in-flight operations using previously resolved Arc clones.baml_language/crates/sys_ops/src/lib.rs (1)
683-960: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd in-crate Rust unit tests for the
schemalowering.The module carries non-trivial logic (cycle handling, union/null widening,
$ref/$defs, JSON-pointer escaping) but is only exercised via the.bamlintegration test. Direct unit tests would pin edge cases (self-referential/mutually-recursive classes, nullable unions, enum aliases).As per coding guidelines: "Prefer writing Rust unit tests over integration tests where possible".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@baml_language/crates/sys_ops/src/lib.rs` around lines 683 - 960, Add in-crate Rust unit tests within the schema module covering json_schema lowering for self-referential and mutually recursive classes, nullable unions, JSON-pointer escaping in class references, and enum aliases. Build the required SysOpContext definitions using existing test helpers or fixtures, and assert the generated $ref/$defs structure and representative schema values directly.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@baml_language/crates/sys_native/src/io_impls.rs`:
- Around line 2581-2586: Wrap the tokio_tungstenite::connect_async future in the
existing timeout_from_nanos helper within SysOpOutput::async_op, passing the
request timeout and preserving the established timeout error mapping. Keep the
existing VmBamlError::Io handling for handshake failures, while ensuring an
incomplete TCP/TLS/WebSocket handshake cannot exceed the configured request
timeout.
In `@baml_language/crates/sys_ops/src/lib.rs`:
- Around line 786-793: The TypeAlias branch in ty_schema recursively resolves
aliases without detecting cycles, causing the builtin json alias to overflow the
stack. Add visited alias-key tracking across recursive ty_schema resolution, and
on re-entry emit the existing $ref/Unsupported schema outcome; add a Rust unit
test covering the builtin recursive json alias.
---
Nitpick comments:
In `@baml_language/crates/sys_native/src/io_impls.rs`:
- Around line 2486-2488: Update the Message::Binary branch in the websocket
next() handling to return an error for unexpected binary frames instead of
converting bytes with String::from_utf8_lossy. Preserve the existing successful
text-frame behavior and use the module’s established error type or constructor
to report the binary-frame condition.
- Around line 2525-2542: Update the WebSocket close flow in close to become
asynchronous as needed, send Message::Close(None) through the stream’s sink, and
flush it before removing the resource from REGISTRY. Preserve the existing
successful SysOpOutput result and resource unregistration behavior; do not
attempt to forcibly interrupt in-flight operations using previously resolved Arc
clones.
In `@baml_language/crates/sys_ops/src/lib.rs`:
- Around line 683-960: Add in-crate Rust unit tests within the schema module
covering json_schema lowering for self-referential and mutually recursive
classes, nullable unions, JSON-pointer escaping in class references, and enum
aliases. Build the required SysOpContext definitions using existing test helpers
or fixtures, and assert the generated $ref/$defs structure and representative
schema values directly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 7a3c0456-138a-4baa-a3fe-a52f3aab224f
⛔ Files ignored due to path filters (1)
baml_language/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (18)
baml_language/crates/baml_builtins2/baml_std/baml/ns_errors/capability.bamlbaml_language/crates/baml_builtins2/baml_std/baml/ns_json/json.bamlbaml_language/crates/baml_builtins2/baml_std/baml/ns_llm/llm.bamlbaml_language/crates/baml_builtins2/baml_std/baml/ns_llm/llm_types.bamlbaml_language/crates/baml_builtins2/baml_std/baml/ns_sap/sap.bamlbaml_language/crates/baml_builtins2/baml_std/baml/ns_schema/schema.bamlbaml_language/crates/baml_builtins2/baml_std/baml/ns_ws/ws.bamlbaml_language/crates/baml_builtins2/src/lib.rsbaml_language/crates/baml_compiler2_emit/src/lib.rsbaml_language/crates/baml_compiler2_tir/src/interfaces/impl_rules.rsbaml_language/crates/baml_tests/baml_src/ns_interface_direct_methods/interface_direct_methods.bamlbaml_language/crates/baml_tests/baml_src/ns_provider_stdlib/provider_stdlib.bamlbaml_language/crates/bex_resource_types/src/lib.rsbaml_language/crates/sys_native/Cargo.tomlbaml_language/crates/sys_native/src/io_impls.rsbaml_language/crates/sys_native/src/registry.rsbaml_language/crates/sys_ops/Cargo.tomlbaml_language/crates/sys_ops/src/lib.rs
Binary size checks passed✅ 7 passed
Generated by |
4bd3ea0 to
236c051
Compare
236c051 to
10f6e4b
Compare
10f6e4b to
c1424c5
Compare
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit c1424c5. Configure here.
69c69a5 to
147a55f
Compare
Summary
Why
BAML-authored LLM providers need reusable stdlib primitives for parsing, request/schema adaptation, realtime transport, and provider-neutral failure handling. These APIs are the minimal built-in and runtime support needed for the custom-provider fixtures to type-check with a freshly built CLI.
Validation
Note
Medium Risk
New public I/O surfaces (WebSockets, schema lowering) and dependency growth affect shipped CLI/pack sizes and runtime networking; changes are covered by tests and size gates but expand the attack/ops surface for providers.
Overview
Introduces stdlib building blocks for BAML-authored LLM providers: jq-style
baml.json.path/path_orwithJsonPathError, publicbaml.sap.parse/parse_type(via newbaml.llm.__new_stream_cache),baml.schema.json_schema(Rust I/O), and a textbaml.wsclient (connect,send,next,close). Also addsbaml.errors.UnknownErrorfor errors without a known capability channel.Registers the new builtin modules in
baml_builtins2, wires WebSocket support throughtokio-tungstenite/ rustls inCargo.lock, and adds provider stdlib BAML tests plus compiler/CLI snapshot updates.Size-gate baselines and per-platform ceilings in
.cargo/size-gate.tomland.ci/size-gate/*are bumped (notably packed-program grows on all native targets).Reviewed by Cursor Bugbot for commit 147a55f. Bugbot is set up for automated code reviews on this repo. Configure here.
Summary by CodeRabbit