feat: return usage in ChatCompletionStreamResponse (#506)

router/src/lib.rs

@@ -621,6 +621,8 @@ struct ChatCompletionStreamResponse {
  created: i64,
  model: String,
  choices: Vec<ChatCompletionStreamResponseChoice>,
+ #[serde(skip_serializing_if = "Option::is_none")]
+ usage: Option<UsageInfo>,
 }
 
 #[derive(Serialize, ToSchema, PartialEq)]
@@ -824,6 +826,20 @@ impl From<GenerateResponse> for ChatCompletionResponse {
 
 impl From<StreamResponse> for ChatCompletionStreamResponse {
  fn from(resp: StreamResponse) -> Self {
+ let prompt_tokens = resp.details.as_ref().map(|x| x.prompt_tokens).unwrap_or(0);
+ let completion_tokens = resp.details.as_ref().map(|x| x.generated_tokens);
+ let total_tokens = prompt_tokens + completion_tokens.unwrap_or(0);
+
+ let usage: Option<UsageInfo> = if completion_tokens.is_some() {
+ Some(UsageInfo {
+ prompt_tokens: prompt_tokens,
+ total_tokens: total_tokens,
+ completion_tokens: completion_tokens,
+ })
+ } else {
+ None
+ };
+
  let finish_reason = resp
  .details
  .map(|x| CompletionFinishReason::from(x.finish_reason));
@@ -849,6 +865,7 @@ impl From<StreamResponse> for ChatCompletionStreamResponse {
  },
  finish_reason: finish_reason,
  }],
+ usage: usage,
  }
  }
 }