From 5e6e5826b92bf153f0f8243f2d66b98d31f853e9 Mon Sep 17 00:00:00 2001
From: Garance <garance.mary@gmail.com>
Date: Mon, 16 Sep 2024 10:28:04 -0400
Subject: [PATCH] fix(perplexity): fix preamble and context in completion
 request

---
 rig-core/src/providers/perplexity.rs | 34 ++++++++++++++++++++++------
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/rig-core/src/providers/perplexity.rs b/rig-core/src/providers/perplexity.rs
index f940da56..01703143 100644
--- a/rig-core/src/providers/perplexity.rs
+++ b/rig-core/src/providers/perplexity.rs
@@ -213,13 +213,33 @@ impl completion::CompletionModel for CompletionModel {
         &self,
         completion_request: completion::CompletionRequest,
     ) -> Result<completion::CompletionResponse<CompletionResponse>, CompletionError> {
-        let mut messages = completion_request.chat_history.clone();
-        if let Some(preamble) = completion_request.preamble {
-            messages.push(completion::Message {
-                role: "system".to_string(),
-                content: preamble,
-            });
-        }
+        // Add preamble to messages (if available)
+        let mut messages = if let Some(preamble) = &completion_request.preamble {
+            vec![completion::Message {
+                role: "system".into(),
+                content: preamble.clone(),
+            }]
+        } else {
+            vec![]
+        };
+
+        // Add context documents to chat history
+        messages.append(
+            completion_request
+                .documents
+                .into_iter()
+                .map(|doc| completion::Message {
+                    role: "system".into(),
+                    content: serde_json::to_string(&doc).expect("Document should serialize"),
+                })
+                .collect::<Vec<_>>()
+                .as_mut(),
+        );
+
+        // Add chat history to messages
+        messages.extend(completion_request.chat_history);
+
+        // Add user prompt to messages
         messages.push(completion::Message {
             role: "user".to_string(),
             content: completion_request.prompt,