diff --git a/Server/application/chukebao/controller/AiChatController.php b/Server/application/chukebao/controller/AiChatController.php
index c488a1ad..75d0416e 100644
--- a/Server/application/chukebao/controller/AiChatController.php
+++ b/Server/application/chukebao/controller/AiChatController.php
@@ -4,6 +4,7 @@ namespace app\chukebao\controller;
 
 use app\ai\controller\DouBaoAI;
 use app\chukebao\controller\TokensRecordController as tokensRecord;
+use app\chukebao\model\TokensCompany;
 use library\ResponseHelper;
 use think\Db;
 
@@ -20,11 +21,8 @@ class AiChatController extends BaseController
             return ResponseHelper::error('参数缺失');
         }
 
-        $tokens = Db::name('users')
-            ->where('id', $userId)
-            ->where('companyId', $companyId)
-            ->value('tokens');
-        if ($tokens <= 0){
+        $tokens = TokensCompany::where(['companyId' => $companyId])->value('tokens');
+        if (empty($tokens) || $tokens <= 0){
             return ResponseHelper::error('用户Tokens余额不足');
         }
 
diff --git a/Server/application/chukebao/controller/WechatChatroomController.php b/Server/application/chukebao/controller/WechatChatroomController.php
index 01d07756..76edb7e0 100644
--- a/Server/application/chukebao/controller/WechatChatroomController.php
+++ b/Server/application/chukebao/controller/WechatChatroomController.php
@@ -4,6 +4,7 @@ namespace app\chukebao\controller;
 
 use app\ai\controller\DouBaoAI;
 use app\chukebao\controller\TokensRecordController as tokensRecord;
+use app\chukebao\model\TokensCompany;
 use library\ResponseHelper;
 use think\Db;
 
@@ -96,11 +97,8 @@ class WechatChatroomController extends BaseController
             return ResponseHelper::error('参数缺失');
         }
 
-        $tokens = Db::name('users')
-            ->where('id', $userId)
-            ->where('companyId', $companyId)
-            ->value('tokens');
-        if ($tokens <= 0){
+        $tokens = TokensCompany::where(['companyId' => $companyId])->value('tokens');
+        if (empty($tokens) || $tokens <= 0){
             return ResponseHelper::error('用户Tokens余额不足');
         }
 
@@ -138,8 +136,6 @@ class WechatChatroomController extends BaseController
         }
 
 
-        exit_data($res);
-
     }
 
 }
\ No newline at end of file