Merge pull request #73 from 0xPlaygrounds/misc/cleanup-small-improvem…

…ents Small cleanup/changes before 0.3.0
0xPlaygrounds · Oct 24, 2024 · fb6bc54 · fb6bc54
2 parents 0d0083b + a0bee83
commit fb6bc54
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 6 deletions.
diff --git a/rig-core/src/completion.rs b/rig-core/src/completion.rs
@@ -266,7 +266,7 @@ pub struct CompletionRequest {
 }
 
 impl CompletionRequest {
-    pub fn prompt_with_context(&self) -> String {
+    pub(crate) fn prompt_with_context(&self) -> String {
         if !self.documents.is_empty() {
             format!(
                 "<attachments>\n{}</attachments>\n\n{}",

diff --git a/rig-core/src/embeddings.rs b/rig-core/src/embeddings.rs
@@ -97,7 +97,7 @@ pub trait EmbeddingModel: Clone + Sync + Send {
     /// Embed multiple documents in a single request
     fn embed_documents(
         &self,
-        documents: Vec<String>,
+        documents: impl IntoIterator<Item = String> + Send,
     ) -> impl std::future::Future<Output = Result<Vec<Embedding>, EmbeddingError>> + Send;
 }
 

diff --git a/rig-core/src/lib.rs b/rig-core/src/lib.rs
@@ -61,17 +61,19 @@
 //! Rig natively supports the following completion and embedding model providers:
 //! - OpenAI
 //! - Cohere
+//! - Anthropic
+//! - Perplexity
 //!
 //! Rig currently has the following integration companion crates:
 //! - `rig-mongodb`: Vector store implementation for MongoDB
-//!
+//! - `rig-lancedb`: Vector store implementation for LanceDB
 
 pub mod agent;
 pub mod cli_chatbot;
 pub mod completion;
 pub mod embeddings;
 pub mod extractor;
-pub mod json_utils;
+pub(crate) mod json_utils;
 pub mod providers;
 pub mod tool;
 pub mod vector_store;
diff --git a/rig-core/src/providers/cohere.rs b/rig-core/src/providers/cohere.rs
@@ -199,8 +199,10 @@ impl embeddings::EmbeddingModel for EmbeddingModel {
 
     async fn embed_documents(
         &self,
-        documents: Vec<String>,
+        documents: impl IntoIterator<Item = String>,
     ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
+        let documents = documents.into_iter().collect::<Vec<_>>();
+
         let response = self
             .client
             .post("/v1/embed")

diff --git a/rig-core/src/providers/openai.rs b/rig-core/src/providers/openai.rs
@@ -241,8 +241,10 @@ impl embeddings::EmbeddingModel for EmbeddingModel {
 
     async fn embed_documents(
         &self,
-        documents: Vec<String>,
+        documents: impl IntoIterator<Item = String>,
     ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
+        let documents = documents.into_iter().collect::<Vec<_>>();
+
         let response = self
             .client
             .post("/v1/embeddings")