From db47bc58c54b26748d152b240a8a22032b727a0c Mon Sep 17 00:00:00 2001
From: Paul Hendricks <phendricks@nvidia.com>
Date: Tue, 24 Jun 2025 16:16:28 -0400
Subject: [PATCH] refactor: using async_openai::types::Logprobs

---
 lib/llm/src/protocols/openai/completions.rs        | 14 +-------------
 .../src/protocols/openai/completions/aggregator.rs |  4 ++--
 2 files changed, 3 insertions(+), 15 deletions(-)

diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs
index ea561d3c67..127fbe48a0 100644
--- a/lib/llm/src/protocols/openai/completions.rs
+++ b/lib/llm/src/protocols/openai/completions.rs
@@ -13,8 +13,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::HashMap;
-
 use derive_builder::Builder;
 use dynamo_runtime::protocols::annotated::AnnotationsProvider;
 use serde::{Deserialize, Serialize};
@@ -92,7 +90,7 @@ pub struct CompletionChoice {
 
     #[serde(skip_serializing_if = "Option::is_none")]
     #[builder(default, setter(strip_option))]
-    pub logprobs: Option<LogprobResult>,
+    pub logprobs: Option<async_openai::types::Logprobs>,
 }
 
 impl ContentProvider for CompletionChoice {
@@ -107,16 +105,6 @@ impl CompletionChoice {
     }
 }
 
-// TODO: validate this is the correct format
-/// Legacy OpenAI LogprobResult component
-#[derive(Clone, Debug, Deserialize, Serialize)]
-pub struct LogprobResult {
-    pub tokens: Vec<String>,
-    pub token_logprobs: Vec<f32>,
-    pub top_logprobs: Vec<HashMap<String, f32>>,
-    pub text_offset: Vec<i32>,
-}
-
 pub fn prompt_to_string(prompt: &async_openai::types::Prompt) -> String {
     match prompt {
         async_openai::types::Prompt::String(s) => s.clone(),
diff --git a/lib/llm/src/protocols/openai/completions/aggregator.rs b/lib/llm/src/protocols/openai/completions/aggregator.rs
index b2380fbf89..91990ce110 100644
--- a/lib/llm/src/protocols/openai/completions/aggregator.rs
+++ b/lib/llm/src/protocols/openai/completions/aggregator.rs
@@ -18,7 +18,7 @@ use std::{collections::HashMap, str::FromStr};
 use anyhow::Result;
 use futures::StreamExt;
 
-use super::{CompletionChoice, CompletionResponse, LogprobResult};
+use super::{CompletionChoice, CompletionResponse};
 use crate::protocols::{
     codec::{Message, SseCodecError},
     common::FinishReason,
@@ -40,7 +40,7 @@ struct DeltaChoice {
     index: u64,
     text: String,
     finish_reason: Option<FinishReason>,
-    logprobs: Option<LogprobResult>,
+    logprobs: Option<async_openai::types::Logprobs>,
 }
 
 impl Default for DeltaAggregator {