Skip to content

Commit 5531e2c

Browse files
densumeshcdxker
authored andcommitted
feature: add prompt to image search
1 parent 00a64ac commit 5531e2c

File tree

3 files changed

+12
-6
lines changed

3 files changed

+12
-6
lines changed

server/src/data/models.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7656,7 +7656,10 @@ impl From<(ParsedQuery, f32)> for MultiQuery {
76567656
#[derive(Debug, Serialize, Deserialize, ToSchema, Clone, PartialEq)]
76577657
#[serde(untagged)]
76587658
pub enum SearchModalities {
7659-
Image { image_url: String },
7659+
Image {
7660+
image_url: String,
7661+
llm_prompt: Option<String>,
7662+
},
76607663
Text(String),
76617664
}
76627665

server/src/operators/message_operator.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1245,6 +1245,7 @@ pub async fn get_topic_string(
12451245

12461246
pub async fn get_text_from_image(
12471247
image_url: String,
1248+
prompt: Option<String>,
12481249
dataset: &Dataset,
12491250
) -> Result<String, ServiceError> {
12501251
let dataset_config = DatasetConfiguration::from_json(dataset.server_configuration.clone());
@@ -1286,12 +1287,11 @@ pub async fn get_text_from_image(
12861287
},
12871288
};
12881289

1290+
let default_system_prompt = "Please describe the image and turn the description into a search query. DO NOT INCLUDE ANY OTHER CONTEXT OR INFORMATION. JUST OUTPUT THE SEARCH QUERY AND NOTHING ELSE".to_string();
1291+
12891292
let messages = vec![
12901293
ChatMessage::System {
1291-
content: ChatMessageContent::Text(
1292-
"Please describe the image and turn the description into a search query. DO NOT INCLUDE ANY OTHER CONTEXT OR INFORMATION. JUST OUTPUT THE SEARCH QUERY AND NOTHING ELSE"
1293-
.to_string(),
1294-
),
1294+
content: ChatMessageContent::Text(prompt.unwrap_or(default_system_prompt)),
12951295
name: None,
12961296
},
12971297
ChatMessage::User {

server/src/operators/search_operator.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1618,7 +1618,10 @@ pub async fn parse_query(
16181618
let stop_words = get_stop_words();
16191619
let query = match query {
16201620
SearchModalities::Text(query) => query,
1621-
SearchModalities::Image { image_url } => get_text_from_image(image_url, dataset).await?,
1621+
SearchModalities::Image {
1622+
image_url,
1623+
llm_prompt,
1624+
} => get_text_from_image(image_url, llm_prompt, dataset).await?,
16221625
};
16231626

16241627
let query = match remove_stop_words {

0 commit comments

Comments
 (0)