feature: add prompt to image search

devflowinc · Jan 8, 2025 · 5531e2c · 5531e2c
1 parent 00a64ac
commit 5531e2c
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git a/server/src/data/models.rs b/server/src/data/models.rs
@@ -7656,7 +7656,10 @@ impl From<(ParsedQuery, f32)> for MultiQuery {
 #[derive(Debug, Serialize, Deserialize, ToSchema, Clone, PartialEq)]
 #[serde(untagged)]
 pub enum SearchModalities {
-    Image { image_url: String },
+    Image {
+        image_url: String,
+        llm_prompt: Option<String>,
+    },
     Text(String),
 }
 

diff --git a/server/src/operators/message_operator.rs b/server/src/operators/message_operator.rs
@@ -1245,6 +1245,7 @@ pub async fn get_topic_string(
 
 pub async fn get_text_from_image(
     image_url: String,
+    prompt: Option<String>,
     dataset: &Dataset,
 ) -> Result<String, ServiceError> {
     let dataset_config = DatasetConfiguration::from_json(dataset.server_configuration.clone());
@@ -1286,12 +1287,11 @@ pub async fn get_text_from_image(
         },
     };
 
+    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();
+
     let messages = vec![
         ChatMessage::System {
-            content: ChatMessageContent::Text(
-                "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(),
-            ),
+            content: ChatMessageContent::Text(prompt.unwrap_or(default_system_prompt)),
             name: None,
         },
         ChatMessage::User {

diff --git a/server/src/operators/search_operator.rs b/server/src/operators/search_operator.rs
@@ -1618,7 +1618,10 @@ pub async fn parse_query(
     let stop_words = get_stop_words();
     let query = match query {
         SearchModalities::Text(query) => query,
-        SearchModalities::Image { image_url } => get_text_from_image(image_url, dataset).await?,
+        SearchModalities::Image {
+            image_url,
+            llm_prompt,
+        } => get_text_from_image(image_url, llm_prompt, dataset).await?,
     };
 
     let query = match remove_stop_words {