diff --git a/Partner-Framework/src/main/java/work/slhaf/partner/api/agent/factory/context/AgentContext.kt b/Partner-Framework/src/main/java/work/slhaf/partner/api/agent/factory/context/AgentContext.kt
index e2c64544..7d1dd011 100644
--- a/Partner-Framework/src/main/java/work/slhaf/partner/api/agent/factory/context/AgentContext.kt
+++ b/Partner-Framework/src/main/java/work/slhaf/partner/api/agent/factory/context/AgentContext.kt
@@ -15,9 +15,9 @@ object AgentContext {
         get() = _modules
 
     private val _capabilities =
-        mutableMapOf<Class<*>, Any?>()
+        mutableMapOf<Class<*>, Any>()
 
-    val capabilities: Map<Class<*>, Any?>
+    val capabilities: Map<Class<*>, Any>
         get() = _capabilities
 
     private val _additionalComponents = mutableSetOf<Any>()
@@ -34,8 +34,8 @@ object AgentContext {
         _modules[name] = module
     }
 
-    fun <T> addCapability(type: Class<T>, value: T) {
-        _capabilities[type] = value
+    fun addCapability(value: Any) {
+        _capabilities[value::class.java] = value
     }
 
     fun addAdditionalComponent(instance: Any): Boolean {