From 5d5c4452c20b5af8097b4e821454863d5717c0f9 Mon Sep 17 00:00:00 2001
From: Moritz Althaus <moritz.althaus@aleph-alpha.com>
Date: Wed, 23 Oct 2024 08:55:47 +0200
Subject: [PATCH] refactor: make ChatOutput public

---
 src/chat.rs | 11 ++++++-----
 src/lib.rs  |  5 ++---
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/chat.rs b/src/chat.rs
index d2e537e..4f4af8e 100644
--- a/src/chat.rs
+++ b/src/chat.rs
@@ -83,15 +83,16 @@ impl<'a> TaskChat<'a> {
 }
 
 #[derive(Deserialize, Debug, PartialEq, Eq)]
-pub struct Choice {
+pub struct ChatOutput {
     pub message: Message<'static>,
     pub finish_reason: String,
 }
 
 #[derive(Deserialize, Debug, PartialEq, Eq)]
-pub struct ChatResponse {
-    pub choices: Vec<Choice>,
+pub struct ResponseChat {
+    pub choices: Vec<ChatOutput>,
 }
+
 #[derive(Serialize)]
 struct ChatBody<'a> {
     /// Name of the model tasked with completing the prompt. E.g. `luminous-base"`.
@@ -125,9 +126,9 @@ impl<'a> ChatBody<'a> {
 }
 
 impl<'a> Task for TaskChat<'a> {
-    type Output = Choice;
+    type Output = ChatOutput;
 
-    type ResponseBody = ChatResponse;
+    type ResponseBody = ResponseChat;
 
     fn build_request(
         &self,
diff --git a/src/lib.rs b/src/lib.rs
index 2ca2cc7..57575fe 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -34,13 +34,12 @@ mod semantic_embedding;
 mod tokenization;
 use std::time::Duration;
 
-use chat::Choice;
 use http::HttpClient;
 use semantic_embedding::{BatchSemanticEmbeddingOutput, SemanticEmbeddingOutput};
 use tokenizers::Tokenizer;
 
 pub use self::{
-    chat::{Role, TaskChat},
+    chat::{ChatOutput, Message, Role, TaskChat},
     completion::{CompletionOutput, Sampling, Stopping, TaskCompletion},
     detokenization::{DetokenizationOutput, TaskDetokenization},
     explanation::{
@@ -213,7 +212,7 @@ impl Client {
     ///     Ok(())
     /// }
     /// ```
-    pub async fn chat(&self, task: &TaskChat<'_>, model: &str, how: &How) -> Result<Choice, Error> {
+    pub async fn chat(&self, task: &TaskChat<'_>, model: &str, how: &How) -> Result<ChatOutput, Error> {
         self.http_client
             .output_of(&task.with_model(model), how)
             .await