From 45a8a19eee45c211ec491e954e7f04ccc1bd3e12 Mon Sep 17 00:00:00 2001
From: Phil Scott <phil.scott@gmail.com>
Date: Thu, 2 Jan 2025 17:11:55 -0500
Subject: [PATCH] Adds an exception for missing templates.

---
 LLama/Exceptions/RuntimeError.cs | 19 +++++++++++++++++++
 LLama/LLamaTemplate.cs           | 14 ++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/LLama/Exceptions/RuntimeError.cs b/LLama/Exceptions/RuntimeError.cs
index 4db77911e..bf20543ca 100644
--- a/LLama/Exceptions/RuntimeError.cs
+++ b/LLama/Exceptions/RuntimeError.cs
@@ -58,6 +58,25 @@ public LLamaDecodeError(DecodeResult returnCode)
     }
 }
 
+/// <summary>
+/// `llama_decode` return a non-zero status code
+/// </summary>
+public class MissingTemplateException
+    : RuntimeError
+{
+    /// <inheritdoc />
+    public MissingTemplateException()
+        : base("llama_chat_apply_template failed: template not found")
+    {
+    }
+    
+    /// <inheritdoc />
+    public MissingTemplateException(string message)
+        : base($"llama_chat_apply_template failed: template not found for '{message}'")
+    {
+    } 
+}
+
 /// <summary>
 /// `llama_get_logits_ith` returned null, indicating that the index was invalid
 /// </summary>
diff --git a/LLama/LLamaTemplate.cs b/LLama/LLamaTemplate.cs
index 8f9c9daa2..e82cbccb4 100644
--- a/LLama/LLamaTemplate.cs
+++ b/LLama/LLamaTemplate.cs
@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Text;
+using LLama.Exceptions;
 using LLama.Native;
 
 namespace LLama;
@@ -250,6 +251,19 @@ public ReadOnlySpan<byte> Apply()
                 {
                     // Run templater and discover true length
                     var outputLength = ApplyInternal(_nativeChatMessages.AsSpan(0, Count), output);
+                    
+                    // if we have a return code of -1, the template was not found.
+                    if (outputLength == -1)
+                    {
+                        if (_customTemplate != null)
+                        {
+                            throw new MissingTemplateException(Encoding.GetString(_customTemplate));
+                        }
+                        else
+                        {
+                            throw new MissingTemplateException();    
+                        }
+                    }    
 
                     // If length was too big for output buffer run it again
                     if (outputLength > output.Length)