fix: fixed bug with base64 encoding on AWS Bedrock (#432)

* fix: fixed bug with base64 encoding on AWS Bedrock

* fix: use uuid v4 for document names

* fix: prefix documents names

Author: mkranjac

Cargo.lock

@@ -18,6 +18,7 @@ aws-sdk-bedrockruntime = "1.77.0"
 aws-smithy-types = "1.3.0"
 base64 = "0.22.1"
 async-stream = "0.3.6"
+uuid = { version = "1.16.0", features = ["v4"]}
 
 [dev-dependencies]
 anyhow = "1.0.75"

rig-bedrock/Cargo.toml

@@ -1,12 +1,12 @@
 use aws_sdk_bedrockruntime::types as aws_bedrock;
-
 use rig::{
     completion::CompletionError,
     message::{ContentFormat, Document},
 };
 
 pub(crate) use crate::types::media_types::RigDocumentMediaType;
 use base64::{prelude::BASE64_STANDARD, Engine};
+use uuid::Uuid;
 
 #[derive(Clone)]
 pub struct RigDocument(pub Document);
@@ -15,27 +15,33 @@ impl TryFrom<RigDocument> for aws_bedrock::DocumentBlock {
     type Error = CompletionError;
 
     fn try_from(value: RigDocument) -> Result<Self, Self::Error> {
-        let maybe_format = value
+        let document_media_type = value
             .0
             .media_type
             .map(|doc| RigDocumentMediaType(doc).try_into());
 
-        let format = match maybe_format {
+        let document_media_type = match document_media_type {
             Some(Ok(document_format)) => Ok(Some(document_format)),
             Some(Err(err)) => Err(err),
             None => Ok(None),
         }?;
 
-        let document_data = BASE64_STANDARD
-            .decode(value.0.data)
-            .map_err(|e| CompletionError::ProviderError(e.to_string()))?;
+        let document_data = match value.0.format {
+            Some(ContentFormat::Base64) => BASE64_STANDARD
+                .decode(value.0.data)
+                .map_err(|e| CompletionError::ProviderError(e.to_string()))?,
+            _ => value.0.data.as_bytes().to_vec(),
+        };
+
         let data = aws_smithy_types::Blob::new(document_data);
         let document_source = aws_bedrock::DocumentSource::Bytes(data);
 
+        let random_string = Uuid::new_v4().simple().to_string();
+        let document_name = format!("document-{}", random_string);
         let result = aws_bedrock::DocumentBlock::builder()
             .source(document_source)
-            .name("Document")
-            .set_format(format)
+            .name(document_name)
+            .set_format(document_media_type)
             .build()
             .map_err(|e| CompletionError::ProviderError(e.to_string()))?;
         Ok(result)
@@ -82,13 +88,35 @@ mod tests {
     fn test_document_to_aws_document() {
         let rig_document = RigDocument(Document {
             data: "data".into(),
-            format: Some(ContentFormat::Base64),
+            format: Some(ContentFormat::String),
             media_type: Some(DocumentMediaType::PDF),
         });
         let aws_document: Result<aws_bedrock::DocumentBlock, _> = rig_document.clone().try_into();
         assert_eq!(aws_document.is_ok(), true);
         let aws_document = aws_document.unwrap();
         assert_eq!(aws_document.format, aws_bedrock::DocumentFormat::Pdf);
+        let document_data = rig_document.0.data.as_bytes().to_vec();
+        let aws_document_bytes = aws_document
+            .source()
+            .unwrap()
+            .as_bytes()
+            .unwrap()
+            .as_ref()
+            .to_owned();
+
+        let doc_name = aws_document.name;
+        assert!(doc_name.starts_with("document-"));
+        assert_eq!(aws_document_bytes, document_data)
+    }
+
+    #[test]
+    fn test_base64_document_to_aws_document() {
+        let rig_document = RigDocument(Document {
+            data: "data".into(),
+            format: Some(ContentFormat::Base64),
+            media_type: Some(DocumentMediaType::PDF),
+        });
+        let aws_document: aws_bedrock::DocumentBlock = rig_document.clone().try_into().unwrap();
         let document_data = BASE64_STANDARD.decode(rig_document.0.data).unwrap();
         let aws_document_bytes = aws_document
             .source()
@@ -104,7 +132,7 @@ mod tests {
     fn test_unsupported_document_to_aws_document() {
         let rig_document = RigDocument(Document {
             data: "data".into(),
-            format: Some(ContentFormat::Base64),
+            format: Some(ContentFormat::String),
             media_type: Some(DocumentMediaType::Javascript),
         });
         let aws_document: Result<aws_bedrock::DocumentBlock, _> = rig_document.clone().try_into();