From f762b931c34ecb458fdf6dbaa157e0509df136d9 Mon Sep 17 00:00:00 2001
From: Nico Arqueros <nico@dcspark.io>
Date: Wed, 13 Nov 2024 20:57:16 -0600
Subject: [PATCH] fix

---
 shinkai-libs/shinkai-db/src/db/db_inbox.rs    | 22 +++++++++++++++----
 .../src/schemas/smart_inbox.rs                | 17 +++++++++++++-
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/shinkai-libs/shinkai-db/src/db/db_inbox.rs b/shinkai-libs/shinkai-db/src/db/db_inbox.rs
index 7f58bc527..08647846d 100644
--- a/shinkai-libs/shinkai-db/src/db/db_inbox.rs
+++ b/shinkai-libs/shinkai-db/src/db/db_inbox.rs
@@ -543,11 +543,25 @@ impl ShinkaiDB {
                                     let job = self.get_job_with_options(&unique_id, false, false)?;
                                     let agent_id = job.parent_agent_or_llm_provider_id;
 
-                                    // TODO: add caching so we don't call this every time for the same agent_id
-                                    match self.get_llm_provider(&agent_id, &p) {
+                                    // Check if the agent_id is an LLM provider
+                                    let agent_subset = match self.get_llm_provider(&agent_id, &p) {
                                         Ok(agent) => agent.map(LLMProviderSubset::from_serialized_llm_provider),
-                                        Err(_) => None,
-                                    }
+                                        Err(_) => {
+                                            // If not found as an LLM provider, check if it exists as an agent
+                                            match self.get_agent(&agent_id) {
+                                                Ok(Some(agent)) => {
+                                                    // Fetch the serialized LLM provider
+                                                    if let Ok(Some(serialized_llm_provider)) = self.get_llm_provider(&agent.llm_provider_id, &p) {
+                                                        Some(LLMProviderSubset::from_agent(agent, serialized_llm_provider))
+                                                    } else {
+                                                        None
+                                                    }
+                                                }
+                                                _ => None,
+                                            }
+                                        }
+                                    };
+                                    agent_subset
                                 }
                                 _ => None,
                             }
diff --git a/shinkai-libs/shinkai-message-primitives/src/schemas/smart_inbox.rs b/shinkai-libs/shinkai-message-primitives/src/schemas/smart_inbox.rs
index ec09f20d2..59427f244 100644
--- a/shinkai-libs/shinkai-message-primitives/src/schemas/smart_inbox.rs
+++ b/shinkai-libs/shinkai-message-primitives/src/schemas/smart_inbox.rs
@@ -4,7 +4,14 @@ use utoipa::ToSchema;
 
 use crate::shinkai_message::{shinkai_message::ShinkaiMessage, shinkai_message_schemas::V2ChatMessage};
 
-use super::{job_config::JobConfig, llm_providers::serialized_llm_provider::{LLMProviderInterface, SerializedLLMProvider}, shinkai_name::ShinkaiName};
+use super::{
+    job_config::JobConfig,
+    llm_providers::{
+        agent::Agent,
+        serialized_llm_provider::{LLMProviderInterface, SerializedLLMProvider},
+    },
+    shinkai_name::ShinkaiName,
+};
 
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, ToSchema)]
 pub struct LLMProviderSubset {
@@ -21,6 +28,14 @@ impl LLMProviderSubset {
             model: serialized_llm_provider.model,
         }
     }
+
+    pub fn from_agent(agent: Agent, serialized_llm_provider: SerializedLLMProvider) -> Self {
+        Self {
+            id: agent.agent_id,
+            full_identity_name: agent.full_identity_name,
+            model: serialized_llm_provider.model,
+        }
+    }
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone, ToSchema)]