From 80116a1281887fca50bc78fba5d4d35168002fcb Mon Sep 17 00:00:00 2001 From: Ronald Chan Date: Tue, 31 Dec 2024 15:13:46 +0800 Subject: [PATCH] feat: add send sync for CF worker (#176) --- .gitignore | 1 + Cargo.lock | 97 +++++++++++++++++++ rig-core/Cargo.toml | 12 ++- .../src/providers/anthropic/completion.rs | 1 + rig-core/src/providers/cohere.rs | 2 + rig-core/src/providers/gemini/completion.rs | 1 + rig-core/src/providers/gemini/embedding.rs | 1 + rig-core/src/providers/openai.rs | 2 + rig-core/src/providers/perplexity.rs | 1 + rig-core/src/providers/xai/completion.rs | 1 + rig-core/src/providers/xai/embedding.rs | 1 + 11 files changed, 115 insertions(+), 5 deletions(-) diff --git a/.gitignore b/.gitignore index 7f04b051..ed3fad9e 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ target/ .DS_Store .idea/ .vscode/ +.devcontainer/ \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 6c86bb4c..ccd00505 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4824,6 +4824,7 @@ dependencies = [ "tokio-test", "tracing", "tracing-subscriber", + "worker", ] [[package]] @@ -5273,6 +5274,28 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-wasm-bindgen" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3b143e2833c57ab9ad3ea280d21fd34e285a42837aeb0ee301f4f41890fa00e" +dependencies = [ + "js-sys", + "serde", + "wasm-bindgen", +] + +[[package]] +name = "serde-wasm-bindgen" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8302e169f0eddcc139c70f139d19d6467353af16f9fce27e8c30158036a1e16b" +dependencies = [ + "js-sys", + "serde", + "wasm-bindgen", +] + [[package]] name = "serde_bytes" version = "0.11.15" @@ -6880,6 +6903,80 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "worker" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "727789ca7eff9733efbea9d0e97779edc1cf1926e98aee7d7d8afe32805458aa" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures-channel", + "futures-util", + "http 1.1.0", + "http-body 1.0.1", + "js-sys", + "matchit", + "pin-project", + "serde", + "serde-wasm-bindgen 0.6.5", + "serde_json", + "serde_urlencoded", + "tokio", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", + "worker-kv", + "worker-macros", + "worker-sys", +] + +[[package]] +name = "worker-kv" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f06d4d1416a9f8346ee9123b0d9a11b3cfa38e6cfb5a139698017d1597c4d41" +dependencies = [ + "js-sys", + "serde", + "serde-wasm-bindgen 0.5.0", + "serde_json", + "thiserror 1.0.69", + "wasm-bindgen", + "wasm-bindgen-futures", +] + +[[package]] +name = "worker-macros" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d625c24570ba9207a2617476013335f28a95cbe513e59bb814ffba092a18058" +dependencies = [ + "async-trait", + "proc-macro2", + "quote", + "syn 2.0.89", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-bindgen-macro-support", + "worker-sys", +] + +[[package]] +name = "worker-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34563340d41016b4381257c5a16b0d2bc590dbe00500ecfbebcaa16f5f85ce90" +dependencies = [ + "cfg-if", + "js-sys", + "wasm-bindgen", + "web-sys", +] + [[package]] name = "write16" version = "1.0.0" diff --git a/rig-core/Cargo.toml b/rig-core/Cargo.toml index afb3c1d9..68cda329 100644 --- a/rig-core/Cargo.toml +++ b/rig-core/Cargo.toml @@ -27,6 +27,7 @@ rig-derive = { version = "0.1.0", path = "./rig-core-derive", optional = true } glob = "0.3.1" lopdf = { version = "0.34.0", optional = true } rayon = { version = "1.10.0", optional = true} +worker = { version = "0.5", optional = true } [dev-dependencies] anyhow = "1.0.75" @@ -40,6 +41,7 @@ all = ["derive", "pdf", "rayon"] derive = ["dep:rig-derive"] pdf = ["dep:lopdf"] rayon = ["dep:rayon"] +worker = ["dep:worker"] [[test]] name = "embed_macro" @@ -47,20 +49,20 @@ required-features = ["derive"] [[example]] name = "rag" -required-features = ["derive"] +required-features = ["derive"] [[example]] name = "vector_search" -required-features = ["derive"] +required-features = ["derive"] [[example]] name = "vector_search_cohere" -required-features = ["derive"] +required-features = ["derive"] [[example]] name = "gemini_embeddings" -required-features = ["derive"] +required-features = ["derive"] [[example]] name = "xai_embeddings" -required-features = ["derive"] +required-features = ["derive"] diff --git a/rig-core/src/providers/anthropic/completion.rs b/rig-core/src/providers/anthropic/completion.rs index 56c00bed..4d4d1d1c 100644 --- a/rig-core/src/providers/anthropic/completion.rs +++ b/rig-core/src/providers/anthropic/completion.rs @@ -189,6 +189,7 @@ enum ToolChoice { impl completion::CompletionModel for CompletionModel { type Response = CompletionResponse; + #[cfg_attr(feature = "worker", worker::send)] async fn completion( &self, completion_request: completion::CompletionRequest, diff --git a/rig-core/src/providers/cohere.rs b/rig-core/src/providers/cohere.rs index 883204d4..d362ec44 100644 --- a/rig-core/src/providers/cohere.rs +++ b/rig-core/src/providers/cohere.rs @@ -211,6 +211,7 @@ impl embeddings::EmbeddingModel for EmbeddingModel { self.ndims } + #[cfg_attr(feature = "worker", worker::send)] async fn embed_texts( &self, documents: impl IntoIterator, @@ -512,6 +513,7 @@ impl CompletionModel { impl completion::CompletionModel for CompletionModel { type Response = CompletionResponse; + #[cfg_attr(feature = "worker", worker::send)] async fn completion( &self, completion_request: completion::CompletionRequest, diff --git a/rig-core/src/providers/gemini/completion.rs b/rig-core/src/providers/gemini/completion.rs index 73b9ec6f..c4ec70d2 100644 --- a/rig-core/src/providers/gemini/completion.rs +++ b/rig-core/src/providers/gemini/completion.rs @@ -45,6 +45,7 @@ impl CompletionModel { impl completion::CompletionModel for CompletionModel { type Response = GenerateContentResponse; + #[cfg_attr(feature = "worker", worker::send)] async fn completion( &self, mut completion_request: CompletionRequest, diff --git a/rig-core/src/providers/gemini/embedding.rs b/rig-core/src/providers/gemini/embedding.rs index c2b76e02..87cdd329 100644 --- a/rig-core/src/providers/gemini/embedding.rs +++ b/rig-core/src/providers/gemini/embedding.rs @@ -41,6 +41,7 @@ impl embeddings::EmbeddingModel for EmbeddingModel { } } + #[cfg_attr(feature = "worker", worker::send)] async fn embed_texts( &self, documents: impl IntoIterator + Send, diff --git a/rig-core/src/providers/openai.rs b/rig-core/src/providers/openai.rs index 27cb3b08..194863be 100644 --- a/rig-core/src/providers/openai.rs +++ b/rig-core/src/providers/openai.rs @@ -249,6 +249,7 @@ impl embeddings::EmbeddingModel for EmbeddingModel { self.ndims } + #[cfg_attr(feature = "worker", worker::send)] async fn embed_texts( &self, documents: impl IntoIterator, @@ -475,6 +476,7 @@ impl CompletionModel { impl completion::CompletionModel for CompletionModel { type Response = CompletionResponse; + #[cfg_attr(feature = "worker", worker::send)] async fn completion( &self, mut completion_request: CompletionRequest, diff --git a/rig-core/src/providers/perplexity.rs b/rig-core/src/providers/perplexity.rs index fa1e34fb..a60b2cf6 100644 --- a/rig-core/src/providers/perplexity.rs +++ b/rig-core/src/providers/perplexity.rs @@ -198,6 +198,7 @@ impl CompletionModel { impl completion::CompletionModel for CompletionModel { type Response = CompletionResponse; + #[cfg_attr(feature = "worker", worker::send)] async fn completion( &self, completion_request: completion::CompletionRequest, diff --git a/rig-core/src/providers/xai/completion.rs b/rig-core/src/providers/xai/completion.rs index 560e35d9..2ae81be9 100644 --- a/rig-core/src/providers/xai/completion.rs +++ b/rig-core/src/providers/xai/completion.rs @@ -38,6 +38,7 @@ impl CompletionModel { impl completion::CompletionModel for CompletionModel { type Response = CompletionResponse; + #[cfg_attr(feature = "worker", worker::send)] async fn completion( &self, mut completion_request: completion::CompletionRequest, diff --git a/rig-core/src/providers/xai/embedding.rs b/rig-core/src/providers/xai/embedding.rs index f0ad1c92..651eea45 100644 --- a/rig-core/src/providers/xai/embedding.rs +++ b/rig-core/src/providers/xai/embedding.rs @@ -69,6 +69,7 @@ impl embeddings::EmbeddingModel for EmbeddingModel { self.ndims } + #[cfg_attr(feature = "worker", worker::send)] async fn embed_texts( &self, documents: impl IntoIterator,