增加模型训练和重试时的余额检测

imiphp · Oct 19, 2023 · ae4b011 · ae4b011
1 parent d9e602b
commit ae4b011
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 0 deletions.
diff --git a/server/Module/Embedding/Service/EmbeddingRetryParser.php b/server/Module/Embedding/Service/EmbeddingRetryParser.php
@@ -50,6 +50,12 @@ public function __construct(private int $memberId)
         $this->config = EmbeddingConfig::__getConfigAsync();
     }
 
+    public function checkBalance(): void
+    {
+        $modelConfig = $this->config->getEmbeddingModelConfig($this->model);
+        $this->memberCardService->checkBalance($this->memberId, 1, paying: $modelConfig ? $modelConfig->paying : false);
+    }
+
     public function asyncRun(): void
     {
         Coroutine::create(function () {

diff --git a/server/Module/Embedding/Service/EmbeddingService.php b/server/Module/Embedding/Service/EmbeddingService.php
@@ -39,6 +39,8 @@ public function upload(int $memberId, string $fileName, string $clientFileName,
     {
         $parser = App::newInstance(EmbeddingUploadParser::class, $memberId, $fileName, $clientFileName, $ip, $id, $override, $directory, $sectionSeparator, $sectionSplitLength, $sectionSplitByTitle);
 
+        $parser->checkBalance();
+
         return $parser->upload();
     }
 
@@ -457,6 +459,7 @@ public function retryProject(string $id, int $memberId = 0, bool $force = false)
         $project->status = EmbeddingStatus::TRAINING;
         $project->update();
         $retryParser = App::newInstance(EmbeddingRetryParser::class, $memberId);
+        $retryParser->checkBalance();
         $retryParser->asyncRun();
         foreach ($this->fileList($id, $memberId, $force ? 0 : EmbeddingStatus::FAILED) as $file)
         {
@@ -487,6 +490,7 @@ public function retryFile(string|EmbeddingFile $file, int $memberId = 0, ?Embedd
             $project->status = EmbeddingStatus::TRAINING;
             $project->update();
             $retryParser = App::newInstance(EmbeddingRetryParser::class, $memberId);
+            $retryParser->checkBalance();
             $retryParser->asyncRun();
         }
         foreach ($this->sectionList($file->projectId, $file->id, $memberId, $force ? 0 : EmbeddingStatus::FAILED) as $section)
@@ -534,6 +538,7 @@ public function retrySection(string|EmbeddingSection $_section, int $memberId =
             $project->update();
 
             $retryParser = App::newInstance(EmbeddingRetryParser::class, $memberId);
+            $retryParser->checkBalance();
             $retryParser->asyncRun();
         }
         $retryParser->push($section);

diff --git a/server/Module/Embedding/Service/EmbeddingUploadParser.php b/server/Module/Embedding/Service/EmbeddingUploadParser.php
@@ -91,6 +91,12 @@ public function __construct(private int $memberId, private string $fileName, pri
         $this->setSectionSeparator($sectionSeparator);
     }
 
+    public function checkBalance(): void
+    {
+        $modelConfig = $this->config->getEmbeddingModelConfig($this->model);
+        $this->memberCardService->checkBalance($this->memberId, 1, paying: $modelConfig ? $modelConfig->paying : false);
+    }
+
     public function setSectionSeparator(string $sectionSeparator): void
     {
         $this->sectionSeparator = stripcslashes($sectionSeparator);