Add support for arbritrary metadata (e.g. cache_policy for anthropic) (…

…#893)

Add ability for people to support caching or any future content metadata
that someone may wish to add.

Most LLM apis have converged on:
```
{
  role: string
  content: {
   ... // multiple variants
  }[]
}
```

We now allow for people to inject arbitrary metadata into the properties
of content (or appropriately named field for any llm api all.

What is not yet supported:
* indexing and adding metadata to subinterfaces of content (i.e. setting
content.image_type.detail)
* adding metadata as a sibling to role, content

The playground and raw_curl correctly display this so if a provider
doesn't support a speicfic metadat, we do support it.

Complete e2e example:
```rust
client<llm> Foo {
  provider openai
  options {
    allowed_role_metadata: ["foo", "bar"]
  }
}

client<llm> FooWithout {
  provider openai
  options {
  }
}
template_string Foo() #"
  {{ _.role('user', foo={"type": "ephemeral"}, bar="1", cat=True) }}
  This will be have foo and bar, but not cat metadata. But only for Foo, not FooWithout.
  {{ _.role('user') }}
  This will have none of the role metadata for Foo or FooWithout.
"#
```

Metadata included:
<img width="446" alt="Screenshot 2024-08-23 at 7 41 40 AM"
src="https://github.com/user-attachments/assets/e26b1e2d-0377-42af-8629-8fcdff0db3cf">

Metadata excluded:
<img width="459" alt="Screenshot 2024-08-23 at 7 42 15 AM"
src="https://github.com/user-attachments/assets/1dc8fcf3-9a17-4e8d-8a00-29d9bfee31e0">


Potentially breaking changes:
* We now merge sequentially matching roles together into the same chunk.
```
[{
  "role": "user",
  "chunk": [A, B]
},
{
  "role": "user",
  "chunk": [C, D]
}]
```

is now converted to
```
[{
  "role": "user",
  "chunk": [A, B, C, D]
}]
```

docs/docs/snippets/clients/providers/anthropic.mdx

@@ -70,6 +70,8 @@ client<llm> MyClient {
 ```
 </ParamField>
 
+<Markdown src="../../../../snippets/allowed-role-metadata.mdx" />
+
 ## Forwarded options
 <ParamField
    path="system"

docs/docs/snippets/clients/providers/aws-bedrock.mdx

@@ -28,6 +28,19 @@ We use the AWS SDK under the hood, which will respect [all authentication mechan
   - loading the specified `AWS_PROFILE` from `~/.aws/config`
   - built-in authn for services running in EC2, ECS, Lambda, etc.
 
+## Non-forwarded options
+
+<ParamField
+  path="default_role"
+  type="string"
+>
+  The default role for any prompts that don't specify a role. **Default: `system`**
+
+  We don't have any checks for this field, you can pass any string you wish.
+</ParamField>
+
+<Markdown src="../../../../snippets/allowed-role-metadata-basic.mdx" />
+
 ## Forwarded options
 
 <ParamField

docs/docs/snippets/clients/providers/azure.mdx

@@ -96,6 +96,8 @@ client<llm> MyClient {
 ```
 </ParamField>
 
+<Markdown src="../../../../snippets/allowed-role-metadata-basic.mdx" />
+
 ## Forwarded options
 <ParamField
    path="messages"

docs/docs/snippets/clients/providers/gemini.mdx

@@ -78,6 +78,7 @@ client<llm> MyClient {
 ```
 </ParamField>
 
+<Markdown src="../../../../snippets/allowed-role-metadata-basic.mdx" />
 
 ## Forwarded options
 <ParamField

docs/docs/snippets/clients/providers/ollama.mdx

@@ -55,6 +55,8 @@ client<llm> MyClient {
 ```
 </ParamField>
 
+<Markdown src="../../../../snippets/allowed-role-metadata-basic.mdx" />
+
 ## Forwarded options
 <ParamField
    path="messages"

docs/docs/snippets/clients/providers/openai.mdx

@@ -63,6 +63,8 @@ client<llm> MyClient {
 
 </ParamField>
 
+<Markdown src="../../../../snippets/allowed-role-metadata-basic.mdx" />
+
 ## Forwarded options
 
 <ParamField

docs/docs/snippets/clients/providers/vertex.mdx

@@ -108,7 +108,7 @@ The options are passed through directly to the API, barring a few. Here's a shor
 
 
   <Accordion title='Example Credentials Content'>
-  ```json Credentials
+```json Credentials
     {
       "type": "service_account",
       "project_id": "my-project-id",
@@ -122,7 +122,7 @@ The options are passed through directly to the API, barring a few. Here's a shor
       "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/...",
       "universe_domain": "googleapis.com"
     }
-    ```
+```
 
 </Accordion>
 
@@ -187,6 +187,8 @@ client<llm> MyClient {
 ```
 </ParamField>
 
+<Markdown src="../../../../snippets/allowed-role-metadata-basic.mdx" />
+
 ## Forwarded options
 <ParamField
   path="safetySettings"

docs/snippets/allowed-role-metadata-basic.mdx

@@ -0,0 +1,34 @@
+<ParamField
+  path="allowed_role_metadata"
+  type="string[]"
+>
+  Which role metadata should we forward to the API? **Default: `[]`**
+
+  For example you can set this to `["foo", "bar"]` to forward the cache policy to the API.
+
+  If you do not set `allowed_role_metadata`, we will not forward any role metadata to the API even if it is set in the prompt.
+
+  Then in your prompt you can use something like:
+  ```baml
+  client<llm> Foo {
+    provider openai
+    options {
+      allowed_role_metadata: ["foo", "bar"]
+    }
+  }
+
+  client<llm> FooWithout {
+    provider openai
+    options {
+    }
+  }
+  template_string Foo() #"
+    {{ _.role('user', foo={"type": "ephemeral"}, bar="1", cat=True) }}
+    This will be have foo and bar, but not cat metadata. But only for Foo, not FooWithout.
+    {{ _.role('user') }}
+    This will have none of the role metadata for Foo or FooWithout.
+  "#
+  ```
+
+  You can use the playground to see the raw curl request to see what is being sent to the API.
+</ParamField>

docs/snippets/allowed-role-metadata.mdx

@@ -0,0 +1,35 @@
+<ParamField
+  path="allowed_role_metadata"
+  type="string[]"
+>
+  Which role metadata should we forward to the API? **Default: `[]`**
+
+  For example you can set this to `["cache_control"]` to forward the cache policy to the API.
+
+  If you do not set `allowed_role_metadata`, we will not forward any role metadata to the API even if it is set in the prompt.
+
+  Then in your prompt you can use something like:
+  ```baml
+  client<llm> Foo {
+    provider anthropic
+    options {
+      allowed_role_metadata: ["cache_control"]
+    }
+  }
+
+  client<llm> FooWithout {
+    provider anthropic
+    options {
+    }
+  }
+
+  template_string Foo() #"
+    {{ _.role('user', cache_control={"type": "ephemeral"}) }}
+    This will be cached for Foo, but not for FooWithout!
+    {{ _.role('user') }}
+    This will not be cached for Foo or FooWithout!
+  "#
+  ```
+
+  You can use the playground to see the raw curl request to see what is being sent to the API.
+</ParamField>

engine/baml-lib/jinja/Cargo.toml

@@ -23,6 +23,7 @@ minijinja = { version = "1.0.16", default-features = false, features = [
   "unstable_machinery_serde",
   "custom_syntax",
   "internal_debug",
+  "deserialization",
   # We don't want to use these features:
   # multi_template
   # loader

engine/baml-lib/jinja/src/chat_message_part.rs

@@ -0,0 +1,80 @@
+use std::collections::{HashMap, HashSet};
+
+use baml_types::{BamlMedia, BamlMediaContent};
+use serde::Serialize;
+
+use crate::RenderedChatMessage;
+
+#[derive(Debug, PartialEq, Serialize, Clone)]
+pub enum ChatMessagePart {
+    // raw user-provided text
+    Text(String),
+    Media(BamlMedia),
+    WithMeta(Box<ChatMessagePart>, HashMap<String, serde_json::Value>),
+}
+
+impl ChatMessagePart {
+    pub fn with_meta(self, meta: HashMap<String, serde_json::Value>) -> ChatMessagePart {
+        match self {
+            ChatMessagePart::WithMeta(part, mut existing_meta) => {
+                existing_meta.extend(meta);
+                ChatMessagePart::WithMeta(part, existing_meta)
+            }
+            _ => ChatMessagePart::WithMeta(Box::new(self), meta),
+        }
+    }
+
+    pub fn as_text(&self) -> Option<&String> {
+        match self {
+            ChatMessagePart::Text(t) => Some(t),
+            ChatMessagePart::WithMeta(t, _) => t.as_text(),
+            ChatMessagePart::Media(_) => None,
+        }
+    }
+
+    pub fn as_media(&self) -> Option<&BamlMedia> {
+        match self {
+            ChatMessagePart::Media(m) => Some(m),
+            ChatMessagePart::WithMeta(t, _) => t.as_media(),
+            ChatMessagePart::Text(_) => None,
+        }
+    }
+
+    pub fn meta(&self) -> Option<&HashMap<String, serde_json::Value>> {
+        match self {
+            ChatMessagePart::WithMeta(_, meta) => Some(meta),
+            _ => None,
+        }
+    }
+
+    pub fn as_completion(self) -> String {
+        match self {
+            ChatMessagePart::Text(t) => t,
+            ChatMessagePart::Media(_) => "".to_string(), // we are choosing to ignore the image for now
+            ChatMessagePart::WithMeta(p, _) => p.as_completion(),
+        }
+    }
+}
+
+impl std::fmt::Display for ChatMessagePart {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            ChatMessagePart::Text(t) => write!(f, "{}", t),
+            ChatMessagePart::Media(media) => match &media.content {
+                BamlMediaContent::Url(url) => {
+                    write!(f, "<{}_placeholder: {}>", media.media_type, url.url)
+                }
+                BamlMediaContent::Base64(_) => {
+                    write!(f, "<{}_placeholder base64>", media.media_type)
+                }
+                BamlMediaContent::File(file) => write!(
+                    f,
+                    "<{}_placeholder: {}>",
+                    media.media_type,
+                    file.relpath.to_string_lossy()
+                ),
+            },
+            ChatMessagePart::WithMeta(part, meta) => write!(f, "{:?}::{}", meta, part),
+        }
+    }
+}