Improve AI Product detection

v3xlabs · Jan 20, 2025 · f5f5834 · f5f5834
1 parent 52544cc
commit f5f5834
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 2 deletions.
diff --git a/engine/src/modules/intelligence/gemini/actor.rs b/engine/src/modules/intelligence/gemini/actor.rs
@@ -21,6 +21,8 @@ impl Actor for GeminiActor {
     ) -> Result<CalculatedResponse, anyhow::Error> {
         let body = GeminiStructuredContentRequest::from_conversation(conversation, strategy);
 
+        tracing::info!("body: {}", serde_json::to_string(&body).unwrap());
+
         let client = reqwest::Client::new();
         let api_key = intelligence.gemini.as_ref().unwrap().api_key.as_str();
 

diff --git a/engine/src/modules/intelligence/gemini/structured.rs b/engine/src/modules/intelligence/gemini/structured.rs
@@ -51,7 +51,7 @@ impl GeminiStructuredContentRequest {
                 true => Some(GeminiStructuredContentRequestToolConfig {
                     function_calling_config: GeminiStructuredContentRequestToolConfigFunctionCallingConfig {
                         allowed_function_names: strategy.allowed_functions.clone(),
-                        mode: serde_json::to_string(strategy.function_mode.as_ref().unwrap_or(&FunctionMode::Auto)).unwrap()
+                        mode: strategy.function_mode.as_ref().unwrap_or(&FunctionMode::Auto).to_string()
                     }
                 }),
                 false => None,

diff --git a/engine/src/modules/intelligence/structured/strategy.rs b/engine/src/modules/intelligence/structured/strategy.rs
@@ -3,6 +3,7 @@ use crate::modules::intelligence::actions::{
     SmartAction, SmartActionType,
 };
 use serde::{Deserialize, Serialize};
+use serde_with::SerializeDisplay;
 
 use super::Conversation;
 
@@ -13,6 +14,8 @@ pub enum Strategy {
     UPCForcedSearch,
     /// 8 rounds max
     Basic,
+    /// 8 rounds max
+    ProductOptimized,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -25,6 +28,16 @@ pub enum FunctionMode {
     None,
 }
 
+impl ToString for FunctionMode {
+    fn to_string(&self) -> String {
+        match self {
+            FunctionMode::Any => "ANY".to_string(),
+            FunctionMode::Auto => "AUTO".to_string(),
+            FunctionMode::None => "NONE".to_string(),
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct StrategyConfig {
     pub max_rounds: u8,
@@ -61,6 +74,30 @@ impl Strategy {
                     _ => FunctionMode::None,
                 }),
             }),
+            Strategy::ProductOptimized => Ok(StrategyConfig {
+                max_rounds: 8,
+                tasks: vec![
+                    SmartActionType::SearchKagi,
+                    SmartActionType::SearchUPCEAN,
+                    SmartActionType::ExtractLDJSON,
+                ],
+                allowed_functions: None,
+                // allowed_functions: match conversation.index {
+                //     // less then 5
+                //     0..8 => Some(vec![
+                //         SearchKagiTask::name(),
+                //         ExtractLDJsonTask::name(),
+                //         SearchUPCEANDatabaseTask::name(),
+                //     ]),
+                //     _ => None,
+                // },
+                function_mode: Some(match conversation.index {
+                    0..3 => FunctionMode::Any,
+                    3..7 => FunctionMode::Auto,
+                    _ => FunctionMode::None,
+                }),
+                // function_mode: Some(FunctionMode::Auto),
+            }),
             Strategy::Basic => Ok(StrategyConfig {
                 max_rounds: 8,
                 tasks: vec![

diff --git a/engine/src/modules/intelligence/tasks/ingress_product.rs b/engine/src/modules/intelligence/tasks/ingress_product.rs
@@ -28,7 +28,7 @@ impl IngressProductTask {
         let actor = GeminiActor;
 
         let conversation = Conversation {
-            strategy: Strategy::Basic,
+            strategy: Strategy::ProductOptimized,
             index: 0,
             messages: vec![ConversationMessage {
                 role: "user".to_string(),