From 3e621e8f6eae0c7920587a8557f78bc02b1e02aa Mon Sep 17 00:00:00 2001
From: Kyle Dickey <41174949+dickeyy@users.noreply.github.com>
Date: Thu, 23 Jan 2025 12:47:12 -0700
Subject: [PATCH] feat: add citations object to CompletionResponse (#20)

* add citations object to completion response

* remove temporary empty string return in the String function

* add test cases for GetCitations
---
 cmd/main.go                 |  3 +++
 perplexity_response.go      | 21 +++++++++++++++------
 perplexity_response_test.go | 21 +++++++++++++++++++++
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index 853e0ef..ff3752c 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -31,6 +31,9 @@ func main() {
 	}
 
 	fmt.Println(res.GetLastContent())
+	for i, c := range res.GetCitations() {
+		fmt.Printf("Citation %d: %s", i+1, c)
+	}
 	// fmt.Printf("%+v\n", *req)
 	// fmt.Println("*************")
 	// fmt.Printf("%+v\n", res)
diff --git a/perplexity_response.go b/perplexity_response.go
index 0422dda..2ac5166 100644
--- a/perplexity_response.go
+++ b/perplexity_response.go
@@ -22,12 +22,13 @@ type Choice struct {
 
 // CompletionResponse is a response object for the Perplexity API.
 type CompletionResponse struct {
-	ID      string   `json:"id"`
-	Model   string   `json:"model"`
-	Created int      `json:"created"`
-	Usage   Usage    `json:"usage"`
-	Object  string   `json:"object"`
-	Choices []Choice `json:"choices"`
+	ID        string    `json:"id"`
+	Model     string    `json:"model"`
+	Created   int       `json:"created"`
+	Usage     Usage     `json:"usage"`
+	Object    string    `json:"object"`
+	Choices   []Choice  `json:"choices"`
+	Citations *[]string `json:"citations,omitempty"`
 }
 
 // String returns a string representation of the CompletionResponse.
@@ -52,3 +53,11 @@ func (r *CompletionResponse) GetLastContent() string {
 	}
 	return r.Choices[len(r.Choices)-1].Message.Content
 }
+
+// GetCitations returns the citations of the completion response.
+func (r *CompletionResponse) GetCitations() []string {
+	if r.Citations == nil {
+		return []string{}
+	}
+	return *r.Citations
+}
diff --git a/perplexity_response_test.go b/perplexity_response_test.go
index 5521da6..f2c6402 100644
--- a/perplexity_response_test.go
+++ b/perplexity_response_test.go
@@ -78,3 +78,24 @@ func TestString(t *testing.T) {
 		assert.Equal(t, "{\n  \"id\": \"id\",\n  \"model\": \"model\",\n  \"created\": 1,\n  \"usage\": {\n    \"prompt_tokens\": 1,\n    \"completion_tokens\": 1,\n    \"total_tokens\": 1\n  },\n  \"object\": \"object\",\n  \"choices\": [\n    {\n      \"index\": 0,\n      \"finish_reason\": \"\",\n      \"message\": {\n        \"role\": \"assistant\",\n        \"content\": \"hello\"\n      },\n      \"delta\": {\n        \"role\": \"\",\n        \"content\": \"\"\n      }\n    }\n  ]\n}", content.String())
 	})
 }
+
+func TestGetCitations(t *testing.T) {
+	t.Run("empty response returns empty citations", func(t *testing.T) {
+		content := perplexity.CompletionResponse{}
+		assert.Equal(t, content.GetCitations(), []string{})
+	})
+
+	t.Run("nil citations returns empty citations", func(t *testing.T) {
+		content := perplexity.CompletionResponse{
+			Citations: nil,
+		}
+		assert.Equal(t, content.GetCitations(), []string{})
+	})
+
+	t.Run("case with a real citations", func(t *testing.T) {
+		content := perplexity.CompletionResponse{
+			Citations: &[]string{"citation1", "citation2"},
+		}
+		assert.Equal(t, content.GetCitations(), []string{"citation1", "citation2"})
+	})
+}