diff --git a/src/AnthropicClient/Json/JsonSerializationOptions.cs b/src/AnthropicClient/Json/JsonSerializationOptions.cs
index 5e426b3..355aae8 100644
--- a/src/AnthropicClient/Json/JsonSerializationOptions.cs
+++ b/src/AnthropicClient/Json/JsonSerializationOptions.cs
@@ -19,6 +19,7 @@ static class JsonSerializationOptions
       new JsonStringEnumConverter(),
       new MessageBatchResultConverter(),
       new CitationConverter(),
+      new SourceConverter(),
     },
     DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
   };
diff --git a/src/AnthropicClient/Json/SourceConverter.cs b/src/AnthropicClient/Json/SourceConverter.cs
new file mode 100644
index 0000000..3847e8d
--- /dev/null
+++ b/src/AnthropicClient/Json/SourceConverter.cs
@@ -0,0 +1,46 @@
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+using AnthropicClient.Models;
+
+namespace AnthropicClient.Json;
+
+class SourceConverter : JsonConverter<Source>
+{
+  public override Source Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
+  {
+    using var jsonDocument = JsonDocument.ParseValue(ref reader);
+    var root = jsonDocument.RootElement;
+    var type = root.GetProperty("type").GetString();
+    return type switch
+    {
+      SourceType.Text => JsonSerializer.Deserialize<TextSource>(root.GetRawText(), options)!,
+      SourceType.Content => JsonSerializer.Deserialize<CustomSource>(root.GetRawText(), options)!,
+      SourceType.Base64 => JsonSerializer.Deserialize<Base64Source>(root.GetRawText(), options)!,
+      _ => throw new JsonException($"Unknown content type: {type}")
+    };
+  }
+
+  public override void Write(Utf8JsonWriter writer, Source value, JsonSerializerOptions options)
+  {
+    if (value is TextSource textSource)
+    {
+      JsonSerializer.Serialize(writer, textSource, options);
+      return;
+    }
+
+    if (value is CustomSource customSource)
+    {
+      JsonSerializer.Serialize(writer, customSource, options);
+      return;
+    }
+
+    if (value is Base64Source base64Source)
+    {
+      JsonSerializer.Serialize(writer, base64Source, options);
+      return;
+    }
+
+    JsonSerializer.Serialize(writer, value, value.GetType(), options);
+  }
+}
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/Base64Source.cs b/src/AnthropicClient/Models/Base64Source.cs
new file mode 100644
index 0000000..b175b7b
--- /dev/null
+++ b/src/AnthropicClient/Models/Base64Source.cs
@@ -0,0 +1,39 @@
+using System.Text.Json.Serialization;
+
+using AnthropicClient.Utils;
+
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents an base64 source.
+/// </summary>
+public class Base64Source : Source
+{
+  /// <summary>
+  /// Gets the media type of the source.
+  /// </summary>
+  [JsonPropertyName("media_type")]
+  public string MediaType { get; init; } = string.Empty;
+
+  /// <summary>
+  /// Gets the data of the source.
+  /// </summary>
+  public string Data { get; init; } = string.Empty;
+
+  /// <summary>
+  /// Initializes a new instance of the <see cref="Base64Source"/> class.
+  /// </summary>
+  /// <param name="mediaType">The media type of the source.</param>
+  /// <param name="data">The data of the source.</param>
+  /// <exception cref="ArgumentException">Thrown when the media type is invalid.</exception>
+  /// <exception cref="ArgumentNullException">Thrown when the media type or data is null.</exception>
+  /// <returns>A new instance of the <see cref="Base64Source"/> class.</returns>
+  public Base64Source(string mediaType, string data) : base(SourceType.Base64)
+  {
+    ArgumentValidator.ThrowIfNull(mediaType, nameof(mediaType));
+    ArgumentValidator.ThrowIfNull(data, nameof(data));
+
+    MediaType = mediaType;
+    Data = data;
+  }
+}
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/CharacterLocationCitation.cs b/src/AnthropicClient/Models/CharacterLocationCitation.cs
new file mode 100644
index 0000000..0b461cb
--- /dev/null
+++ b/src/AnthropicClient/Models/CharacterLocationCitation.cs
@@ -0,0 +1,21 @@
+using System.Text.Json.Serialization;
+
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents a citation for specific locations within text content.
+/// </summary>
+public class CharacterLocationCitation : Citation
+{
+  /// <summary>
+  /// Gets the start character index of the citation.
+  /// </summary>
+  [JsonPropertyName("start_char_index")]
+  public int StartCharIndex { get; init; }
+
+  /// <summary>
+  /// Gets the end character index of the citation.
+  /// </summary>
+  [JsonPropertyName("end_char_index")]
+  public int EndCharIndex { get; init; }
+}
diff --git a/src/AnthropicClient/Models/Citation.cs b/src/AnthropicClient/Models/Citation.cs
new file mode 100644
index 0000000..ec7ae73
--- /dev/null
+++ b/src/AnthropicClient/Models/Citation.cs
@@ -0,0 +1,32 @@
+using System.Text.Json.Serialization;
+
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents a citation
+/// </summary>
+public abstract class Citation
+{
+  /// <summary>
+  /// Gets the type of the citation.
+  /// </summary>
+  public string Type { get; init; } = string.Empty;
+
+  /// <summary>
+  /// Gets the text that is cited.
+  /// </summary>
+  [JsonPropertyName("cited_text")]
+  public string CitedText { get; init; } = string.Empty;
+
+  /// <summary>
+  /// Gets the document index of the citation.
+  /// </summary>
+  [JsonPropertyName("document_index")]
+  public int DocumentIndex { get; init; }
+
+  /// <summary>
+  /// Gets the title of the document from which the citation is made.
+  /// </summary>
+  [JsonPropertyName("document_title")]
+  public string DocumentTitle { get; init; } = string.Empty;
+}
diff --git a/src/AnthropicClient/Models/CitationOption.cs b/src/AnthropicClient/Models/CitationOption.cs
new file mode 100644
index 0000000..8917b5a
--- /dev/null
+++ b/src/AnthropicClient/Models/CitationOption.cs
@@ -0,0 +1,12 @@
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents whether citations are enabled for a document.
+/// </summary>
+public class CitationOption
+{
+  /// <summary>
+  /// Gets a value indicating whether citations are enabled for the document.
+  /// </summary>
+  public bool Enabled { get; init; }
+}
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/CitationType.cs b/src/AnthropicClient/Models/CitationType.cs
index daeda7e..d0b5536 100644
--- a/src/AnthropicClient/Models/CitationType.cs
+++ b/src/AnthropicClient/Models/CitationType.cs
@@ -1,8 +1,22 @@
 namespace AnthropicClient.Models;
 
+/// <summary>
+/// The types of citations that can be returned by the Anthropic API.
+/// </summary>
 public static class CitationType
 {
+  /// <summary>
+  /// A citation that refers to a specific character in the text.
+  /// </summary>
   public const string CharacterLocation = "char_location";
+
+  /// <summary>
+  /// A citation that refers to a specific page in the text.
+  /// </summary>
   public const string PageLocation = "page_location";
+
+  /// <summary>
+  /// A citation that refers to a specific section in the text.
+  /// </summary>
   public const string ContentBlockLocation = "content_block_location";
 }
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/ContentBlockLocationCitation.cs b/src/AnthropicClient/Models/ContentBlockLocationCitation.cs
new file mode 100644
index 0000000..0cbb73a
--- /dev/null
+++ b/src/AnthropicClient/Models/ContentBlockLocationCitation.cs
@@ -0,0 +1,21 @@
+using System.Text.Json.Serialization;
+
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents a citation for content blocks within custom content.
+/// </summary>
+public class ContentBlockLocationCitation : Citation
+{
+  /// <summary>
+  /// Gets the start block index of the citation.
+  /// /// </summary>
+  [JsonPropertyName("start_block_index")]
+  public int StartBlockIndex { get; init; }
+
+  /// <summary>
+  /// Gets the end block index of the citation.
+  /// </summary>
+  [JsonPropertyName("end_block_index")]
+  public int EndBlockIndex { get; init; }
+}
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/CustomSource.cs b/src/AnthropicClient/Models/CustomSource.cs
new file mode 100644
index 0000000..98390ed
--- /dev/null
+++ b/src/AnthropicClient/Models/CustomSource.cs
@@ -0,0 +1,33 @@
+using System.Text.Json.Serialization;
+
+using AnthropicClient.Utils;
+
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents a custom source that contains a list of text content.
+/// </summary>
+public class CustomSource : Source
+{
+  /// <summary>
+  ///  Gets the list of text content that makes up the custom source.
+  /// </summary>
+  public List<TextContent> Content { get; init; } = [];
+
+  [JsonConstructor]
+  internal CustomSource() : base(SourceType.Content)
+  {
+  }
+
+  /// <summary>
+  /// Initializes a new instance of the <see cref="CustomSource"/> class.
+  /// </summary>
+  /// <returns>A new instance of the <see cref="CustomSource"/> class.</returns>
+  /// <exception cref="ArgumentNullException">Thrown when the content is null.</exception>
+  public CustomSource(List<TextContent> content) : base(SourceType.Content)
+  {
+    ArgumentValidator.ThrowIfNull(content, nameof(content));
+
+    Content = content;
+  }
+}
diff --git a/src/AnthropicClient/Models/DocumentContent.cs b/src/AnthropicClient/Models/DocumentContent.cs
index cd86c29..2e93b31 100644
--- a/src/AnthropicClient/Models/DocumentContent.cs
+++ b/src/AnthropicClient/Models/DocumentContent.cs
@@ -12,13 +12,22 @@ public class DocumentContent : Content
   /// <summary>
   /// Gets the source of the document.
   /// </summary>
-  public DocumentSource Source { get; init; } = new();
+  public Source Source { get; init; } = new DocumentSource();
 
-  public string Title { get; init; } = string.Empty;
+  /// <summary>
+  /// Gets the title of the document.
+  /// </summary>
+  public string? Title { get; init; }
 
-  public string Context { get; init; } = string.Empty;
+  /// <summary>
+  /// Gets the context of the document.
+  /// </summary>
+  public string? Context { get; init; }
 
-  public CitationOption Citations { get; init; } = new CitationOption();
+  /// <summary>
+  /// Gets whether citations are enabled for the document.
+  /// </summary>
+  public CitationOption? Citations { get; init; }
 
   [JsonConstructor]
   internal DocumentContent()
@@ -42,7 +51,7 @@ public class DocumentContent : Content
   {
     Validate(mediaType, data);
 
-    Source = new(mediaType, data);
+    Source = new DocumentSource(mediaType, data);
   }
 
   /// <summary>
@@ -57,26 +66,33 @@ public class DocumentContent : Content
   {
     Validate(mediaType, data);
 
-    Source = new(mediaType, data);
+    Source = new DocumentSource(mediaType, data);
   }
 
-  public DocumentContent(DocumentSource source) : base(ContentType.Document)
+  /// <summary>
+  /// Initializes a new instance of the <see cref="DocumentContent"/> class with a document source.
+  /// </summary>
+  /// <param name="source">The document source.</param>
+  /// <returns>A new instance of the <see cref="DocumentContent"/> class.</returns>
+  /// <exception cref="ArgumentNullException">Thrown when the source is null.</exception>
+  public DocumentContent(Source source) : base(ContentType.Document)
   {
     ArgumentValidator.ThrowIfNull(source, nameof(source));
 
     Source = source;
   }
 
-  public DocumentContent(DocumentSource source, CacheControl cacheControl) : base(ContentType.Document, cacheControl)
+  /// <summary>
+  /// Initializes a new instance of the <see cref="DocumentContent"/> class with a document source and cache control.
+  /// </summary>
+  /// <param name="source">The document source.</param>
+  /// <param name="cacheControl">The cache control to be used for the content.</param>
+  /// <returns>A new instance of the <see cref="DocumentContent"/> class.</returns>
+  /// <exception cref="ArgumentNullException">Thrown when the source is null.</exception>
+  public DocumentContent(Source source, CacheControl cacheControl) : base(ContentType.Document, cacheControl)
   {
     ArgumentValidator.ThrowIfNull(source, nameof(source));
 
     Source = source;
   }
 }
-
-
-public class CitationOption
-{
-  public bool Enabled { get; init; }
-}
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/DocumentSource.cs b/src/AnthropicClient/Models/DocumentSource.cs
index e18f978..e7029f2 100644
--- a/src/AnthropicClient/Models/DocumentSource.cs
+++ b/src/AnthropicClient/Models/DocumentSource.cs
@@ -7,26 +7,10 @@ namespace AnthropicClient.Models;
 /// <summary>
 /// Represents a document source.
 /// </summary>
-public class DocumentSource
+public class DocumentSource : Base64Source
 {
-  /// <summary>
-  /// Gets the media type of the document.
-  /// </summary>
-  [JsonPropertyName("media_type")]
-  public string MediaType { get; init; } = string.Empty;
-
-  /// <summary>
-  /// Gets the data of the document.
-  /// </summary>
-  public string Data { get; init; } = string.Empty;
-
-  /// <summary>
-  /// Gets the type of encoding of the document data.
-  /// </summary>
-  public string Type { get; init; } = "base64";
-
   [JsonConstructor]
-  internal DocumentSource()
+  internal DocumentSource() : base(string.Empty, string.Empty)
   {
   }
 
@@ -38,12 +22,7 @@ public class DocumentSource
   /// <exception cref="ArgumentException">Thrown when the media type is invalid.</exception>
   /// <exception cref="ArgumentNullException">Thrown when the media type or data is null.</exception>
   /// <returns>A new instance of the <see cref="DocumentSource"/> class.</returns>
-  public DocumentSource(string mediaType, string data)
+  public DocumentSource(string mediaType, string data) : base(mediaType, data)
   {
-    ArgumentValidator.ThrowIfNull(mediaType, nameof(mediaType));
-    ArgumentValidator.ThrowIfNull(data, nameof(data));
-
-    MediaType = mediaType;
-    Data = data;
   }
 }
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/ImageContent.cs b/src/AnthropicClient/Models/ImageContent.cs
index b096fab..1d4a949 100644
--- a/src/AnthropicClient/Models/ImageContent.cs
+++ b/src/AnthropicClient/Models/ImageContent.cs
@@ -12,7 +12,7 @@ public class ImageContent : Content
   /// <summary>
   /// Gets the source of the image.
   /// </summary>
-  public ImageSource Source { get; init; } = new();
+  public Source Source { get; init; } = new ImageSource();
 
   [JsonConstructor]
   internal ImageContent()
@@ -36,7 +36,7 @@ public class ImageContent : Content
   {
     Validate(mediaType, data);
 
-    Source = new(mediaType, data);
+    Source = new ImageSource(mediaType, data);
   }
 
   /// <summary>
@@ -51,6 +51,6 @@ public class ImageContent : Content
   {
     Validate(mediaType, data);
 
-    Source = new(mediaType, data);
+    Source = new ImageSource(mediaType, data);
   }
 }
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/ImageSource.cs b/src/AnthropicClient/Models/ImageSource.cs
index cc68c74..35257c5 100644
--- a/src/AnthropicClient/Models/ImageSource.cs
+++ b/src/AnthropicClient/Models/ImageSource.cs
@@ -7,26 +7,10 @@ namespace AnthropicClient.Models;
 /// <summary>
 /// Represents an image source.
 /// </summary>
-public class ImageSource
+public class ImageSource : Base64Source
 {
-  /// <summary>
-  /// Gets the media type of the image.
-  /// </summary>
-  [JsonPropertyName("media_type")]
-  public string MediaType { get; init; } = string.Empty;
-
-  /// <summary>
-  /// Gets the data of the image.
-  /// </summary>
-  public string Data { get; init; } = string.Empty;
-
-  /// <summary>
-  /// Gets the type of encoding of the image data.
-  /// </summary>
-  public string Type { get; init; } = "base64";
-
   [JsonConstructor]
-  internal ImageSource()
+  internal ImageSource() : base(string.Empty, string.Empty)
   {
   }
 
@@ -38,17 +22,13 @@ public class ImageSource
   /// <exception cref="ArgumentException">Thrown when the media type is invalid.</exception>
   /// <exception cref="ArgumentNullException">Thrown when the media type or data is null.</exception>
   /// <returns>A new instance of the <see cref="ImageSource"/> class.</returns>
-  public ImageSource(string mediaType, string data)
+  public ImageSource(string mediaType, string data) : base(mediaType, data)
   {
-    ArgumentValidator.ThrowIfNull(mediaType, nameof(mediaType));
-
     if (ImageType.IsValidImageType(mediaType) is false)
     {
       throw new ArgumentException($"Invalid media type: {mediaType}");
     }
 
-    ArgumentValidator.ThrowIfNull(data, nameof(data));
-
     MediaType = mediaType;
     Data = data;
   }
diff --git a/src/AnthropicClient/Models/PageLocationCitation.cs b/src/AnthropicClient/Models/PageLocationCitation.cs
new file mode 100644
index 0000000..040c262
--- /dev/null
+++ b/src/AnthropicClient/Models/PageLocationCitation.cs
@@ -0,0 +1,21 @@
+using System.Text.Json.Serialization;
+
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents a citation for text within a page of a document.
+/// </summary>
+public class PageLocationCitation : Citation
+{
+  /// <summary>
+  /// Gets the start page number of the citation.
+  /// </summary>
+  [JsonPropertyName("start_page_number")]
+  public int StartPageNumber { get; init; }
+
+  /// <summary>
+  /// Gets the end page number of the citation.
+  /// </summary>
+  [JsonPropertyName("end_page_number")]
+  public int EndPageNumber { get; init; }
+}
diff --git a/src/AnthropicClient/Models/Source.cs b/src/AnthropicClient/Models/Source.cs
new file mode 100644
index 0000000..9f0b153
--- /dev/null
+++ b/src/AnthropicClient/Models/Source.cs
@@ -0,0 +1,22 @@
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents a base class for sources.
+/// </summary>
+public abstract class Source
+{
+  /// <summary>
+  /// Gets the type of the source.
+  /// </summary>
+  public string Type { get; init; }
+
+  /// <summary>
+  /// Initializes a new instance of the <see cref="Source"/> class.
+  /// </summary>
+  /// <param name="type">The type of the source.</param>
+  /// <returns>A new instance of the <see cref="Source"/> class.</returns>
+  protected Source(string type)
+  {
+    Type = type;
+  }
+}
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/SourceType.cs b/src/AnthropicClient/Models/SourceType.cs
new file mode 100644
index 0000000..7079bbe
--- /dev/null
+++ b/src/AnthropicClient/Models/SourceType.cs
@@ -0,0 +1,22 @@
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents the types of document sources that can be used in the Anthropic API.
+/// </summary>
+public static class SourceType
+{
+  /// <summary>
+  /// The base64 encoded document source type.
+  /// </summary>
+  public const string Base64 = "base64";
+
+  /// <summary>
+  /// The custom content document source type.
+  /// </summary>
+  public const string Content = "content";
+
+  /// <summary>
+  /// The text document source type.
+  /// </summary>
+  public const string Text = "text";
+}
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/TextContent.cs b/src/AnthropicClient/Models/TextContent.cs
index 924ec1b..4984783 100644
--- a/src/AnthropicClient/Models/TextContent.cs
+++ b/src/AnthropicClient/Models/TextContent.cs
@@ -14,7 +14,10 @@ public class TextContent : Content
   /// </summary>
   public string Text { get; init; } = string.Empty;
 
-  public Citation[] Citations { get; init; } = [];
+  /// <summary>
+  /// Gets the citations associated with the text content.
+  /// </summary>
+  public Citation[]? Citations { get; init; }
 
   [JsonConstructor]
   internal TextContent() : base(ContentType.Text)
@@ -53,44 +56,3 @@ public class TextContent : Content
     Text = text;
   }
 }
-
-public abstract class Citation
-{
-  public string Type { get; init; } = string.Empty;
-
-  [JsonPropertyName("cited_text")]
-  public string CitedText { get; init; } = string.Empty;
-
-  [JsonPropertyName("document_index")]
-  public int DocumentIndex { get; init; }
-
-  [JsonPropertyName("document_title")]
-  public string DocumentTitle { get; init; } = string.Empty;
-}
-
-public class CharacterLocationCitation : Citation
-{
-  [JsonPropertyName("start_char_index")]
-  public int StartCharIndex { get; init; }
-
-  [JsonPropertyName("end_char_index")]
-  public int EndCharIndex { get; init; }
-}
-
-public class PageLocationCitation : Citation
-{
-  [JsonPropertyName("start_page_number")]
-  public int StartPageNumber { get; init; }
-
-  [JsonPropertyName("end_page_number")]
-  public int EndPageNumber { get; init; }
-}
-
-public class ContentBlockLocationCitation : Citation
-{
-  [JsonPropertyName("start_block_index")]
-  public int StartBlockIndex { get; init; }
-
-  [JsonPropertyName("end_block_index")]
-  public int EndBlockIndex { get; init; }
-}
\ No newline at end of file
diff --git a/src/AnthropicClient/Models/TextDocumentSource.cs b/src/AnthropicClient/Models/TextDocumentSource.cs
deleted file mode 100644
index d4dbf71..0000000
--- a/src/AnthropicClient/Models/TextDocumentSource.cs
+++ /dev/null
@@ -1,11 +0,0 @@
-namespace AnthropicClient.Models;
-
-public class TextDocumentSource : DocumentSource
-{
-  public TextDocumentSource(string data) : base("text/plain", data)
-  {
-    Type = "text";
-  }
-}
-
-
diff --git a/src/AnthropicClient/Models/TextSource.cs b/src/AnthropicClient/Models/TextSource.cs
new file mode 100644
index 0000000..9549b0c
--- /dev/null
+++ b/src/AnthropicClient/Models/TextSource.cs
@@ -0,0 +1,33 @@
+using System.Text.Json.Serialization;
+
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Represents a text document source.
+/// </summary>
+public class TextSource : Source
+{
+  /// <summary>
+  /// Gets the media type of the source.
+  /// </summary>
+  [JsonPropertyName("media_type")]
+  public string MediaType { get; } = "text/plain";
+
+  /// <summary>
+  /// Gets the data of the source.
+  /// </summary>
+  public string Data { get; init; } = string.Empty;
+
+  /// <summary>
+  /// Initializes a new instance of the <see cref="TextSource"/> class.
+  /// </summary>
+  /// <param name="data">The data of the document.</param>
+  /// <exception cref="ArgumentNullException">Thrown when the data is null.</exception>
+  /// <returns>A new instance of the <see cref="TextSource"/> class.</returns>
+  public TextSource(string data) : base(SourceType.Text)
+  {
+    Data = data;
+  }
+}
+
+
diff --git a/tests/AnthropicClient.Tests/EndToEnd/AnthropicApiClientTests.cs b/tests/AnthropicClient.Tests/EndToEnd/AnthropicApiClientTests.cs
index 1e0bafb..7c4ca20 100644
--- a/tests/AnthropicClient.Tests/EndToEnd/AnthropicApiClientTests.cs
+++ b/tests/AnthropicClient.Tests/EndToEnd/AnthropicApiClientTests.cs
@@ -287,16 +287,16 @@ public class AnthropicApiClientTests(ConfigurationFixture configFixture) : EndTo
   }
 
   [Fact]
-  public async Task CreateMessageAsync_WhenCitationsAreEnabled_ItShouldReturnCitationsInResponse()
+  public async Task CreateMessageAsync_WhenCitationsAreEnabledForTextDocumentSource_ItShouldReturnCitationsInResponse()
   {
     var request = new MessageRequest(
-      model: AnthropicModels.Claude3Haiku,
+      model: AnthropicModels.Claude35HaikuLatest,
       messages: [
         new(
           MessageRole.User,
           [
             new DocumentContent(
-              new TextDocumentSource("The grass is green. The sky is blue.")
+              new TextSource("The grass is green. The sky is blue.")
             )
             {
               Title = "My Document",
@@ -312,7 +312,90 @@ public class AnthropicApiClientTests(ConfigurationFixture configFixture) : EndTo
     var result = await _client.CreateMessageAsync(request);
 
     result.IsSuccess.Should().BeTrue();
-    result.Value.Content.OfType<TextContent>().SelectMany(static c => c.Citations).Should().NotBeEmpty();
+
+    var citations = result.Value
+      .Content
+      .OfType<TextContent>()
+      .SelectMany(static c =>
+      {
+        return c.Citations is null ? [] : c.Citations;
+      });
+
+    citations.OfType<CharacterLocationCitation>().Should().NotBeEmpty();
+  }
+
+  [Fact]
+  public async Task CreateMessageAsync_WhenCitationsAreEnabledForPDFDocumentSource_ItShouldReturnCitationsInResponse()
+  {
+    var pdfPath = TestFileHelper.GetTestFilePath("addendum.pdf");
+    var bytes = await File.ReadAllBytesAsync(pdfPath);
+    var base64Data = Convert.ToBase64String(bytes);
+
+    var request = new MessageRequest(
+      model: AnthropicModels.Claude35HaikuLatest,
+      messages: [
+        new(
+          MessageRole.User,
+          [
+            new DocumentContent("application/pdf", base64Data)
+            {
+              Title = "My PDF Document",
+              Context = "This is a trustworthy document.",
+              Citations = new() { Enabled = true }
+            },
+            new TextContent("What is the title of this paper?"),
+          ]
+        )
+      ]
+    );
+
+    var result = await _client.CreateMessageAsync(request);
+
+    result.IsSuccess.Should().BeTrue();
+
+    var citations = result.Value
+      .Content
+      .OfType<TextContent>()
+      .SelectMany(static c => c.Citations is null ? [] : c.Citations);
+
+    citations.OfType<PageLocationCitation>().Should().NotBeEmpty();
+  }
+
+  [Fact]
+  public async Task CreateMessageAsync_WhenCitationsAreEnabledForCustomDocumentSource_ItShouldReturnCitationsInResponse()
+  {
+    var request = new MessageRequest(
+      model: AnthropicModels.Claude35HaikuLatest,
+      messages: [
+        new(
+          MessageRole.User,
+          [
+            new DocumentContent(
+              new CustomSource([
+                new TextContent("The grass is green. The sky is blue.")
+              ])
+            )
+            {
+              Title = "My Custom Document",
+              Context = "This is a trustworthy document.",
+              Citations = new() { Enabled = true }
+            },
+            new TextContent("What color is the grass and sky?"),
+          ]
+        )
+      ]
+    );
+
+    var result = await _client.CreateMessageAsync(request);
+
+    result.IsSuccess.Should().BeTrue();
+
+    var citations = result.Value
+      .Content
+      .OfType<TextContent>()
+      .SelectMany(static c => c.Citations is null ? [] : c.Citations);
+
+    citations.OfType<ContentBlockLocationCitation>().Should().NotBeEmpty();
   }
 
   [Fact]
diff --git a/tests/AnthropicClient.Tests/Unit/Models/MessageRequestTests.cs b/tests/AnthropicClient.Tests/Unit/Models/MessageRequestTests.cs
index 2e7e87c..da3f848 100644
--- a/tests/AnthropicClient.Tests/Unit/Models/MessageRequestTests.cs
+++ b/tests/AnthropicClient.Tests/Unit/Models/MessageRequestTests.cs
@@ -558,8 +558,9 @@ public class MessageRequestTests : SerializationTest
 
     var imageContent = messageRequest.Messages[0].Content[0] as ImageContent;
     imageContent!.Type.Should().Be("image");
-    imageContent.Source.MediaType.Should().Be("image/jpeg");
-    imageContent.Source.Data.Should().Be("data");
+
+    imageContent.Source.As<ImageSource>().MediaType.Should().Be("image/jpeg");
+    imageContent.Source.As<ImageSource>().Data.Should().Be("data");
   }
 
   [Fact]