From 5da3db383d178d9ee22e58e47f250f98b7966eda Mon Sep 17 00:00:00 2001
From: Kayvan Sylvan <kayvan@sylvan.com>
Date: Tue, 17 Jun 2025 20:45:03 -0700
Subject: [PATCH] feat: add citation support to perplexity AI responses

## CHANGES

- Add citation extraction from API responses
- Append citations section to response content
- Format citations as numbered markdown list
- Handle citations in streaming responses
- Store last response for citation access
- Add citations after stream completion
- Maintain backward compatibility with responses
---
 plugins/ai/perplexity/perplexity.go | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/plugins/ai/perplexity/perplexity.go b/plugins/ai/perplexity/perplexity.go
index acf3d85e..f0330f2d 100644
--- a/plugins/ai/perplexity/perplexity.go
+++ b/plugins/ai/perplexity/perplexity.go
@@ -106,7 +106,18 @@ func (c *Client) Send(ctx context.Context, msgs []*goopenai.ChatCompletionMessag
 		return "", fmt.Errorf("perplexity API request failed: %w", err) // Corrected capitalization
 	}
 
-	return resp.GetLastContent(), nil
+	content := resp.GetLastContent()
+
+	// Append citations if available
+	citations := resp.GetCitations()
+	if len(citations) > 0 {
+		content += "\n\n# CITATIONS\n\n"
+		for i, citation := range citations {
+			content += fmt.Sprintf("- [%d] %s\n", i+1, citation)
+		}
+	}
+
+	return content, nil
 }
 
 func (c *Client) SendStream(msgs []*goopenai.ChatCompletionMessage, opts *common.ChatOptions, channel chan string) error {
@@ -169,7 +180,9 @@ func (c *Client) SendStream(msgs []*goopenai.ChatCompletionMessage, opts *common
 
 	go func() {
 		defer close(channel) // Ensure the output channel is closed when this goroutine finishes
+		var lastResponse *perplexity.CompletionResponse
 		for resp := range responseChan {
+			lastResponse = &resp
 			if len(resp.Choices) > 0 {
 				content := ""
 				// Corrected: Check Delta.Content and Message.Content directly for non-emptiness
@@ -184,6 +197,17 @@ func (c *Client) SendStream(msgs []*goopenai.ChatCompletionMessage, opts *common
 				}
 			}
 		}
+
+		// Send citations at the end if available
+		if lastResponse != nil {
+			citations := lastResponse.GetCitations()
+			if len(citations) > 0 {
+				channel <- "\n\n# CITATIONS\n\n"
+				for i, citation := range citations {
+					channel <- fmt.Sprintf("- [%d] %s\n", i+1, citation)
+				}
+			}
+		}
 	}()
 
 	return nil