fix: always set return_full_text to false for better UX (#78)

huggingface · Feb 13, 2024 · fe1f6aa · fe1f6aa
1 parent 4437c0c
commit fe1f6aa
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/llm-ls/Cargo.toml b/crates/llm-ls/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "llm-ls"
-version = "0.4.0"
+version = "0.5.2"
 edition = "2021"
 
 [[bin]]

diff --git a/crates/llm-ls/src/backend.rs b/crates/llm-ls/src/backend.rs
@@ -2,7 +2,7 @@ use super::{APIError, APIResponse, Generation, NAME, VERSION};
 use custom_types::llm_ls::{Backend, Ide};
 use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, USER_AGENT};
 use serde::{Deserialize, Serialize};
-use serde_json::{Map, Value};
+use serde_json::{json, Map, Value};
 use std::fmt::Display;
 
 use crate::error::{Error, Result};
@@ -160,6 +160,12 @@ pub(crate) fn build_body(
     match backend {
         Backend::HuggingFace { .. } | Backend::Tgi { .. } => {
             request_body.insert("inputs".to_owned(), Value::String(prompt));
+            if let Some(Value::Object(params)) = request_body.get_mut("parameters") {
+                params.insert("return_full_text".to_owned(), Value::Bool(false));
+            } else {
+                let params = json!({ "parameters": { "return_full_text": false } });
+                request_body.insert("parameters".to_owned(), params);
+            }
         }
         Backend::Ollama { .. } | Backend::OpenAi { .. } => {
             request_body.insert("prompt".to_owned(), Value::String(prompt));