# ==========================================
# 1. Proposed Fix for ingestion (after_run)
# ==========================================
# ... inside after_run ...
if messages:
    await self.mem0_client.add( 
        messages=messages,
        user_id=self.user_id,
        agent_id=self.agent_id,
        app_id=self.application_id,  # Fixed: Map application_id to native app_id parameter
    )


# ==========================================
# 2. Proposed Fix for retrieval (before_run)
# ==========================================
async def before_run(
        self,
        *,
        agent: SupportsAgentRun,
        session: AgentSession,
        context: SessionContext,
        state: dict[str, Any],
    ) -> None:
        """Search Mem0 for relevant memories and add to the session context."""
        self._validate_filters()
        input_text = "\n".join(msg.text for msg in context.input_messages if msg and msg.text and msg.text.strip())
        if not input_text.strip():
            return

        # Query entity partitions independently to bypass strict logical AND limitations
        search_tasks = []
       
        if not self.application_id:
            return
        
        # Query User partition independently
        if self.user_id:
            search_tasks.append(self.mem0_client.search(query=input_text, filters={"user_id" : self.user_id, "app_id" : self.application_id}, top_k=5, threshold=0.25)) 
        
        # Query Agent partition independently
        if self.agent_id:
            search_tasks.append(self.mem0_client.search(query=input_text, filters={"agent_id" : self.agent_id, "app_id" : self.application_id}, top_k=5, threshold=0.25))

        if not search_tasks:
            return
            
        results = await asyncio.gather(*search_tasks, return_exceptions=True)
        # Merge and deduplicate results
        memories = []
        seen_memory_ids = set()
        
        for search_response in results:
            if isinstance(search_response, Exception):
                continue
                
            current_memories = []
            if isinstance(search_response, list):
                current_memories = search_response
            elif isinstance(search_response, dict) and "results" in search_response:
                current_memories = search_response["results"]
            else:
                current_memories = [search_response]
                
            for mem in current_memories:
                mem_id = mem.get("id")
                if mem_id and mem_id not in seen_memory_ids:
                    seen_memory_ids.add(mem_id)
                    memories.append(mem)

        line_separated_memories = "\n".join(memory.get("memory", "") for memory in memories)
        if line_separated_memories:
            context.extend_messages(
                self.source_id,
                [Message(role="user", contents=[f"{self.context_prompt}\n{line_separated_memories}"])],
            )