add send sync for CF worker

0xPlaygrounds · Dec 31, 2024 · 204b2d5 · 204b2d5
1 parent 8e6f197
commit 204b2d5
Show file tree

Hide file tree

Showing 11 changed files with 115 additions and 5 deletions.
diff --git a/.gitignore b/.gitignore
@@ -3,3 +3,4 @@ target/
 .DS_Store
 .idea/
 .vscode/
+.devcontainer/
diff --git a/Cargo.lock b/Cargo.lock
diff --git 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,27 +41,28 @@ all = ["derive", "pdf", "rayon"]
 derive = ["dep:rig-derive"]
 pdf = ["dep:lopdf"]
 rayon = ["dep:rayon"]
+worker = ["dep:worker"]
 
 [[test]]
 name = "embed_macro"
 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
@@ -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
@@ -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<Item = String>,
@@ -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
@@ -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
@@ -41,6 +41,7 @@ impl embeddings::EmbeddingModel for EmbeddingModel {
         }
     }
 
+    #[cfg_attr(feature = "worker", worker::send)]
     async fn embed_texts(
         &self,
         documents: impl IntoIterator<Item = String> + Send,

diff --git 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<Item = String>,
@@ -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
@@ -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
@@ -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
@@ -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<Item = String>,
-Original file line number
+Diff line change
@@ Expand Up / @@ -41,6 +41,7 @@ impl embeddings::EmbeddingModel for EmbeddingModel { @@
             }
         }
+        #[cfg_attr(feature = "worker", worker::send)]
         async fn embed_texts(
             &self,
             documents: impl IntoIterator<Item = String> + Send,
@@ Expand Down @@