From b0c3b1721a904d0a958fd0b83fc3f74a155cb340 Mon Sep 17 00:00:00 2001
From: Miao Zhuang <1060950782@163.com>
Date: Wed, 6 Dec 2023 11:09:04 +0800
Subject: [PATCH] fix: prompt & bug (#698)

---
 app/pages/Import/AIImport/Create.tsx          |  2 +-
 .../Import/TaskList/TaskItem/AIImportItem.tsx | 18 ++++----
 app/pages/Import/TaskList/index.tsx           |  2 +-
 app/pages/Import/index.tsx                    | 21 +++------
 app/pages/MainPage/index.tsx                  |  2 +-
 app/pages/Schema/SchemaConfig/index.tsx       |  2 +-
 app/stores/llm.ts                             | 19 +++++---
 server/api/studio/etc/studio-api.yaml         |  4 +-
 server/api/studio/internal/model/llm-job.go   | 29 +-----------
 .../api/studio/internal/service/llm/import.go |  9 ++--
 server/api/studio/pkg/base/types.go           | 26 +++++++++++
 server/api/studio/pkg/llm/importjob.go        | 44 ++++++++++++-------
 server/api/studio/pkg/llm/job.go              |  9 ++--
 .../api/studio/pkg/llm/transformer/openai.go  |  2 +-
 14 files changed, 101 insertions(+), 88 deletions(-)

diff --git a/app/pages/Import/AIImport/Create.tsx b/app/pages/Import/AIImport/Create.tsx
index 8e8aab00..4a43e0dd 100644
--- a/app/pages/Import/AIImport/Create.tsx
+++ b/app/pages/Import/AIImport/Create.tsx
@@ -157,7 +157,7 @@ const Create = observer((props: { visible: boolean; onCancel: () => void }) => {
 
       <Form layout="vertical" style={{ display: step == 1 ? 'block' : 'none' }}>
         <Form.Item label={intl.get('llm.file')}>
-          <span>{type === 'file' ? valuse.file : valuse.filePath}</span>
+          <span>{valuse.file}</span>
         </Form.Item>
         <Form.Item label={intl.get('llm.url')}>
           <Input disabled value={llm.config.url} />
diff --git a/app/pages/Import/TaskList/TaskItem/AIImportItem.tsx b/app/pages/Import/TaskList/TaskItem/AIImportItem.tsx
index 936b2a1a..de6d2626 100644
--- a/app/pages/Import/TaskList/TaskItem/AIImportItem.tsx
+++ b/app/pages/Import/TaskList/TaskItem/AIImportItem.tsx
@@ -60,7 +60,7 @@ const AIImportItem = observer((props: IProps) => {
   const onTaskDelete = () => {
     _delete('/api/llm/import/job/' + llmJob.job_id)().then((res) => {
       if (res.code === 0) {
-        antMsg.success(intl.get('import.deleteSuccess'));
+        antMsg.success(intl.get('common.success'));
         props.onRefresh();
       }
     });
@@ -69,7 +69,7 @@ const AIImportItem = observer((props: IProps) => {
   const onTaskStop = () => {
     post('/api/llm/import/job/cancel')({ jobId: llmJob.job_id }).then((res) => {
       if (res.code === 0) {
-        antMsg.success(intl.get('import.stopImportingSuccess'));
+        antMsg.success(intl.get('common.success'));
         props.onRefresh();
       }
     });
@@ -78,7 +78,7 @@ const AIImportItem = observer((props: IProps) => {
     setRerunLoading(true);
     const res = await post('/api/llm/import/job/rerun')({ jobId: llmJob.job_id });
     if (res.code === 0) {
-      antMsg.success(intl.get('common.rerunSuccess'));
+      antMsg.success(intl.get('common.success'));
       props.onRefresh();
     }
     setRerunLoading(false);
@@ -172,11 +172,13 @@ const AIImportItem = observer((props: IProps) => {
             />
           </div>
           <div className={styles.operations}>
-            <Button className="primaryBtn" onClick={() => onViewLog(props.data)}>
-              <Tooltip title={intl.get('import.viewLogs')}>
-                <Icon type="icon-studio-btn-ddl" />
-              </Tooltip>
-            </Button>
+            {llmJob.status !== ILLMStatus.Pending && (
+              <Button className="primaryBtn" onClick={() => onViewLog(props.data)}>
+                <Tooltip title={intl.get('import.viewLogs')}>
+                  <Icon type="icon-studio-btn-ddl" />
+                </Tooltip>
+              </Button>
+            )}
             {llmJob.status === ILLMStatus.Running && (
               <Popconfirm
                 placement="left"
diff --git a/app/pages/Import/TaskList/index.tsx b/app/pages/Import/TaskList/index.tsx
index 28fe10fb..7c795d2d 100644
--- a/app/pages/Import/TaskList/index.tsx
+++ b/app/pages/Import/TaskList/index.tsx
@@ -130,7 +130,7 @@ const TaskList = () => {
               {intl.get('import.uploadTemp')}
             </Button>
           )}
-          {global.appSetting.beta.open && global.appSetting.beta.functions.llmImport.open && (
+          {global.appSetting.beta.open && global.appSetting.beta.functions?.llmImport.open && (
             <Button className="studioAddBtn" onClick={onCreateAIIMport}>
               <Icon className="studioAddBtnIcon" type="icon-studio-btn-add" />
               {intl.get('llm.aiImport')}
diff --git a/app/pages/Import/index.tsx b/app/pages/Import/index.tsx
index 850a3b18..b0a1cb94 100644
--- a/app/pages/Import/index.tsx
+++ b/app/pages/Import/index.tsx
@@ -7,6 +7,7 @@ import { useI18n } from '@vesoft-inc/i18n';
 import DatasourceList from './DatasourceList';
 import styles from './index.module.less';
 import TaskList from './TaskList';
+import llm from '@app/stores/llm';
 
 const Import = () => {
   const history = useHistory();
@@ -15,6 +16,7 @@ const Import = () => {
   const { intl } = useI18n();
   useEffect(() => {
     trackPageView('/import');
+    llm.fetchConfig();
   }, []);
   useEffect(() => {
     const path = location.pathname;
@@ -27,28 +29,15 @@ const Import = () => {
   return (
     <div className={cls(styles.nebuaImportPage, 'studioCenterLayout')}>
       <div className="studioTabHeader">
-        <Radio.Group
-          className="studioTabGroup"
-          value={tab}
-          buttonStyle="solid"
-          onChange={handleTabChange}
-        >
+        <Radio.Group className="studioTabGroup" value={tab} buttonStyle="solid" onChange={handleTabChange}>
           <Radio.Button value="tasks">{intl.get('import.importData')}</Radio.Button>
           <Radio.Button value="datasources">{intl.get('import.dataSourceManagement')}</Radio.Button>
         </Radio.Group>
       </div>
       <div>
-        <Route
-          path={`/import/datasources`}
-          exact={true}
-          component={DatasourceList}
-        />
+        <Route path={`/import/datasources`} exact={true} component={DatasourceList} />
 
-        <Route
-          path={`/import/tasks`}
-          exact={true}
-          component={TaskList}
-        />
+        <Route path={`/import/tasks`} exact={true} component={TaskList} />
       </div>
     </div>
   );
diff --git a/app/pages/MainPage/index.tsx b/app/pages/MainPage/index.tsx
index 2773059e..6f1e4d76 100644
--- a/app/pages/MainPage/index.tsx
+++ b/app/pages/MainPage/index.tsx
@@ -1,4 +1,4 @@
-import { Suspense } from 'react';
+import { Suspense, useEffect } from 'react';
 import { Layout, Spin } from 'antd';
 import { Redirect, Route, Switch } from 'react-router-dom';
 import { shouldAlwaysShowWelcome } from '@app/pages/Welcome';
diff --git a/app/pages/Schema/SchemaConfig/index.tsx b/app/pages/Schema/SchemaConfig/index.tsx
index 374970c4..f468f96e 100644
--- a/app/pages/Schema/SchemaConfig/index.tsx
+++ b/app/pages/Schema/SchemaConfig/index.tsx
@@ -36,7 +36,7 @@ const SchemaConfig = () => {
   const { schema, global } = useStore();
   const { spaces, getSpaces, switchSpace, currentSpace } = schema;
   const { currentLocale } = useI18n();
-  const showViewSchemaBetaFunc = global.appSetting.beta.open && global.appSetting.beta.functions.viewSchema.open;
+  const showViewSchemaBetaFunc = global.appSetting.beta.open && global.appSetting.beta.functions?.viewSchema?.open;
   const spaceInUrl = useMemo(() => {
     const params = new URLSearchParams(location.search);
     return params.get('space');
diff --git a/app/stores/llm.ts b/app/stores/llm.ts
index df69830f..a81b0145 100644
--- a/app/stores/llm.ts
+++ b/app/stores/llm.ts
@@ -27,14 +27,21 @@ diff
 > RETURN p.person.name;
 Question:{query_str}
 `;
-export const llmImportPrompt = `You are knowledge graph Expert.
-please extract relationship data from the text below, referring to the schema of the graph, and return the results in the following JSON format without interpreting, don't explain just return the results directly.
+export const llmImportPrompt = `As a knowledge graph expert, your task is to extract relationship data from the following text:
+----
+{text}
+----
+
+Please proceed according to the schema of the knowledge graph:
+----
+{spaceSchema}
+----
+
+Return the results directly, without interpretation. The results should be in the following JSON format:
 {
-  "nodes":[{ "name":"","type":"","props":{} }],
-  "edges":[{ "src":"","dst":"","edgeType":"","props":{} }]
+  "nodes":[{ "name":string,"type":string,"props":object }],
+  "edges":[{ "src":string,"dst":string,"edgeType":string,"props":object }]
 }
-The schema of the graph is: {spaceSchema}
-The text is: {text} 
 The result is:
 `;
 export const llmImportTask = `please excute the task below,and return the result,dont' explain,just return the result directly.
diff --git a/server/api/studio/etc/studio-api.yaml b/server/api/studio/etc/studio-api.yaml
index 3f14fd60..0ffdb6f8 100644
--- a/server/api/studio/etc/studio-api.yaml
+++ b/server/api/studio/etc/studio-api.yaml
@@ -61,5 +61,5 @@ DB:
   MaxIdleConns: 10
 LLM:
   GQLPath: "./data/llm"
-  GQLBatchSize: 1
-  MaxBlockSize: 1
\ No newline at end of file
+  GQLBatchSize: 100
+  MaxBlockSize: 0
\ No newline at end of file
diff --git a/server/api/studio/internal/model/llm-job.go b/server/api/studio/internal/model/llm-job.go
index dc8a208b..6b6282fa 100644
--- a/server/api/studio/internal/model/llm-job.go
+++ b/server/api/studio/internal/model/llm-job.go
@@ -3,6 +3,7 @@ package db
 import (
 	"time"
 
+	"github.com/vesoft-inc/nebula-studio/server/api/studio/pkg/base"
 	"gorm.io/datatypes"
 )
 
@@ -15,35 +16,9 @@ type LLMJob struct {
 	SpaceSchemaString string         `json:"space_schema_string"`
 	File              string         `json:"file"`
 	JobType           string         `json:"job_type"`
-	Status            LLMStatus      `json:"status"`
+	Status            base.LLMStatus `json:"status"`
 	PromptTemplate    string         `json:"prompt_template"`
 	Process           datatypes.JSON `json:"process"`
 	CreateTime        time.Time      `json:"create_time" gorm:"column:create_time;type:datetime;autoCreateTime"`
 	UpdateTime        time.Time      `json:"update_ime" gorm:"column:update_time;type:datetime;autoUpdateTime"`
 }
-
-type Process struct {
-	TotalSize        int     `json:"total"`
-	CurrentSize      int     `json:"current"`
-	Ratio            float64 `json:"ratio"`
-	FailedReason     string  `json:"failed_reason"`
-	PromptTokens     int     `json:"prompt_tokens"`
-	CompletionTokens int     `json:"completion_tokens"`
-}
-
-type LLMStatus string
-
-const (
-	LLMStatusRunning LLMStatus = "running"
-	LLMStatusSuccess LLMStatus = "success"
-	LLMStatusFailed  LLMStatus = "failed"
-	LLMStatusCancel  LLMStatus = "cancel"
-	LLMStatusPending LLMStatus = "pending"
-)
-
-type LLMJobType string
-
-const (
-	LLMJobTypeFile     LLMJobType = "file"
-	LLMJobTypeFilePath LLMJobType = "file_path"
-)
diff --git a/server/api/studio/internal/service/llm/import.go b/server/api/studio/internal/service/llm/import.go
index 161fd74d..82b22e51 100644
--- a/server/api/studio/internal/service/llm/import.go
+++ b/server/api/studio/internal/service/llm/import.go
@@ -11,6 +11,7 @@ import (
 	db "github.com/vesoft-inc/nebula-studio/server/api/studio/internal/model"
 	"github.com/vesoft-inc/nebula-studio/server/api/studio/internal/types"
 	"github.com/vesoft-inc/nebula-studio/server/api/studio/pkg/auth"
+	"github.com/vesoft-inc/nebula-studio/server/api/studio/pkg/base"
 	"github.com/vesoft-inc/nebula-studio/server/api/studio/pkg/llm"
 	"gorm.io/datatypes"
 )
@@ -35,7 +36,7 @@ func (g *llmService) AddImportJob(req *types.LLMImportRequest) (resp *types.LLMR
 		Space:             req.Space,
 		File:              req.File,
 		JobType:           req.Type,
-		Status:            db.LLMStatusPending,
+		Status:            base.LLMStatusPending,
 		Host:              config.Host,
 		UserName:          config.UserName,
 		PromptTemplate:    req.PromptTemplate,
@@ -95,11 +96,11 @@ func (g *llmService) HandleLLMImportJob(req *types.HandleLLMImportRequest) (resp
 	}
 
 	if req.Action == "cancel" {
-		job.Status = db.LLMStatusCancel
-		llm.ChangeRunningJobStatus(job.JobID, db.LLMStatusCancel)
+		job.Status = base.LLMStatusCancel
+		llm.ChangeRunningJobStatus(job.JobID, base.LLMStatusCancel)
 	}
 	if req.Action == "rerun" {
-		job.Status = db.LLMStatusPending
+		job.Status = base.LLMStatusPending
 		// datatypes.JSON
 		job.Process = datatypes.JSON("{}")
 		//delete log & ngql
diff --git a/server/api/studio/pkg/base/types.go b/server/api/studio/pkg/base/types.go
index de47fbb6..8afe63ad 100644
--- a/server/api/studio/pkg/base/types.go
+++ b/server/api/studio/pkg/base/types.go
@@ -8,3 +8,29 @@ const (
 	Error   StatusCode = -1
 	Success StatusCode = 0
 )
+
+type Process struct {
+	TotalSize        int     `json:"total"`
+	CurrentSize      int     `json:"current"`
+	Ratio            float64 `json:"ratio"`
+	FailedReason     string  `json:"failed_reason"`
+	PromptTokens     int     `json:"prompt_tokens"`
+	CompletionTokens int     `json:"completion_tokens"`
+}
+
+type LLMStatus string
+
+const (
+	LLMStatusRunning LLMStatus = "running"
+	LLMStatusSuccess LLMStatus = "success"
+	LLMStatusFailed  LLMStatus = "failed"
+	LLMStatusCancel  LLMStatus = "cancel"
+	LLMStatusPending LLMStatus = "pending"
+)
+
+type LLMJobType string
+
+const (
+	LLMJobTypeFile     LLMJobType = "file"
+	LLMJobTypeFilePath LLMJobType = "file_path"
+)
diff --git a/server/api/studio/pkg/llm/importjob.go b/server/api/studio/pkg/llm/importjob.go
index 50dc03d2..2e8e1a6c 100644
--- a/server/api/studio/pkg/llm/importjob.go
+++ b/server/api/studio/pkg/llm/importjob.go
@@ -16,6 +16,7 @@ import (
 	"github.com/vesoft-inc/nebula-studio/server/api/studio/internal/config"
 	db "github.com/vesoft-inc/nebula-studio/server/api/studio/internal/model"
 	"github.com/vesoft-inc/nebula-studio/server/api/studio/pkg/auth"
+	"github.com/vesoft-inc/nebula-studio/server/api/studio/pkg/base"
 	"github.com/vesoft-inc/nebula-studio/server/api/studio/pkg/client"
 	"github.com/vesoft-inc/nebula-studio/server/api/studio/pkg/pdf"
 	"gorm.io/datatypes"
@@ -48,7 +49,7 @@ type ImportJob struct {
 	LLMConfig  *db.LLMConfig
 	AuthData   *auth.AuthData
 	NSID       string
-	Process    *db.Process
+	Process    *base.Process
 	logFile    *os.File
 	Schema     Schema
 	SchemaMap  map[string]map[string]Field
@@ -59,7 +60,7 @@ func RunFileJob(job *db.LLMJob) {
 		CacheNodes: make(map[string]Node),
 		CacheEdges: make(map[string]map[string]Edge),
 		LLMJob:     job,
-		Process: &db.Process{
+		Process: &base.Process{
 			TotalSize:        0,
 			CurrentSize:      0,
 			Ratio:            0,
@@ -73,7 +74,6 @@ func RunFileJob(job *db.LLMJob) {
 		return
 	}
 	defer llmJob.CloseLogFile()
-	go llmJob.SyncProcess()
 	defer func() {
 		if err := recover(); err != nil {
 			llmJob.WriteLogFile(fmt.Sprintf("panic: %v", err), "error")
@@ -92,6 +92,7 @@ func RunFileJob(job *db.LLMJob) {
 			return
 		}
 	}()
+	go llmJob.SyncProcess(job)
 
 	llmJob.Process.Ratio = 0.01
 	err = llmJob.ParseSchema(job.SpaceSchemaString)
@@ -116,7 +117,8 @@ func RunFileJob(job *db.LLMJob) {
 
 	connectInfo, ok := auth.CtxUserInfoMap[fmt.Sprintf("%s:%s", job.Host, job.UserName)]
 	if !ok {
-		llmJob.WriteLogFile(fmt.Sprintf("get connect info error: %v", err), "error")
+		err := fmt.Errorf("get connect info error: %s", job.Host+" "+job.UserName)
+		llmJob.WriteLogFile(err.Error(), "error")
 		llmJob.SetJobFailed(err)
 		return
 	}
@@ -152,7 +154,9 @@ func RunFileJob(job *db.LLMJob) {
 		return
 	}
 	llmJob.Process.Ratio = 0.8
-	gqls, err := llmJob.MakeGQLFile()
+
+	gqlPath := filepath.Join(config.GetConfig().LLM.GQLPath, fmt.Sprintf("%s/%s.ngql", llmJob.LLMJob.JobID, llmJob.LLMJob.File))
+	gqls, err := llmJob.MakeGQLFile(gqlPath)
 	if err != nil {
 		llmJob.WriteLogFile(fmt.Sprintf("make gql file error: %v", err), "error")
 		llmJob.SetJobFailed(err)
@@ -165,13 +169,13 @@ func RunFileJob(job *db.LLMJob) {
 
 	llmJob.RunGQLFile(gqls)
 	llmJob.Process.Ratio = 1
-	llmJob.LLMJob.Status = db.LLMStatusSuccess
+	llmJob.LLMJob.Status = base.LLMStatusSuccess
 }
 
-func (i *ImportJob) SyncProcess() {
+func (i *ImportJob) SyncProcess(job *db.LLMJob) {
 	for {
 		// stop
-		if i.LLMJob.Status != db.LLMStatusRunning {
+		if job.Status != base.LLMStatusRunning {
 			return
 		}
 		jsonStr, err := json.Marshal(i.Process)
@@ -179,13 +183,13 @@ func (i *ImportJob) SyncProcess() {
 			i.WriteLogFile(fmt.Sprintf("marshal process error: %v", err), "error")
 			continue
 		}
-		i.LLMJob.Process = datatypes.JSON(jsonStr)
+		job.Process = datatypes.JSON(jsonStr)
 		time.Sleep(time.Second)
 	}
 }
 
 func (i *ImportJob) SetJobFailed(failedErr any) {
-	i.LLMJob.Status = db.LLMStatusFailed
+	i.LLMJob.Status = base.LLMStatusFailed
 	i.Process.FailedReason = fmt.Sprintf("%v", failedErr)
 }
 
@@ -227,6 +231,8 @@ func (i *ImportJob) ReadFile(filePath string) (string, error) {
 	if err != nil {
 		return "", fmt.Errorf("open file error: %v", err)
 	}
+	defer file.Close()
+
 	stat, err := file.Stat()
 	if err != nil {
 		return "", fmt.Errorf("get file stat error: %v", err)
@@ -238,7 +244,6 @@ func (i *ImportJob) ReadFile(filePath string) (string, error) {
 		file.Close()
 		return pdf.ReadPDFFile(filePath)
 	}
-	defer file.Close()
 	bytes, err := io.ReadAll(file)
 	if err != nil {
 		return "", fmt.Errorf("read file error: %v", err)
@@ -311,6 +316,11 @@ type Edge struct {
 }
 
 func (i *ImportJob) ParseText(text string) {
+	// remove ```json and ``` in text
+	text = strings.ReplaceAll(text, "```json", "")
+	text = strings.ReplaceAll(text, "```", "")
+	text = strings.ReplaceAll(text, "\n", "")
+
 	jsonObj := LLMResult{}
 	err := json.Unmarshal([]byte(text), &jsonObj)
 	if err != nil {
@@ -400,7 +410,7 @@ func (i *ImportJob) Query(prompt string) (string, error) {
 	return text, nil
 }
 
-func (i *ImportJob) MakeGQLFile() ([]string, error) {
+func (i *ImportJob) MakeGQLFile(filePath string) ([]string, error) {
 	i.WriteLogFile(fmt.Sprintf("start make gql file, nodes length: %d, edges length: %d", len(i.CacheNodes), len(i.CacheEdges)), "info")
 	gqls := make([]string, 0)
 	spaceVIDType := i.Schema.VidType
@@ -490,7 +500,6 @@ func (i *ImportJob) MakeGQLFile() ([]string, error) {
 	}
 
 	gqlStr := strings.Join(gqls, "\n")
-	filePath := filepath.Join(config.GetConfig().LLM.GQLPath, fmt.Sprintf("%s/%s.gql", i.LLMJob.JobID, i.LLMJob.File))
 	file, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
 	if err != nil {
 		return nil, err
@@ -500,7 +509,7 @@ func (i *ImportJob) MakeGQLFile() ([]string, error) {
 	if err != nil {
 		return nil, err
 	}
-	i.WriteLogFile(fmt.Sprintf("make gql file success, gqls length: %d", len(gqls)), "info")
+	i.WriteLogFile(fmt.Sprintf("make gql file success, path:%s, gqls length: %d", filePath, len(gqls)), "info")
 	return gqls, nil
 }
 
@@ -508,6 +517,9 @@ func (i *ImportJob) RunGQLFile(gqls []string) error {
 	i.WriteLogFile(fmt.Sprintf("start run gql, gqls length: %d", len(gqls)), "info")
 	batchSize := config.GetConfig().LLM.GQLBatchSize
 	for index := 0; index < len(gqls); index += batchSize {
+		if IsRunningJobStopped(i.LLMJob.JobID) {
+			return fmt.Errorf("job stopped")
+		}
 		maxEnd := index + batchSize
 		if maxEnd > len(gqls) {
 			maxEnd = len(gqls)
@@ -527,10 +539,10 @@ func (i *ImportJob) RunGQLFile(gqls []string) error {
 				}
 			}
 			if len(success) > 0 {
-				i.WriteLogFile(fmt.Sprintf("run gql success:\n %v ", success), "info")
+				i.WriteLogFile(fmt.Sprintf("run gql success:\n %s ", strings.Join(success, "\n")), "info")
 			}
 			if len(errors) > 0 {
-				i.WriteLogFile(fmt.Sprintf("run gql error:\n %v ", errors), "error")
+				i.WriteLogFile(fmt.Sprintf("run gql error:\n %s ", strings.Join(errors, "\n")), "error")
 			}
 		}
 	}
diff --git a/server/api/studio/pkg/llm/job.go b/server/api/studio/pkg/llm/job.go
index a6d3779e..7eeda10e 100644
--- a/server/api/studio/pkg/llm/job.go
+++ b/server/api/studio/pkg/llm/job.go
@@ -4,6 +4,7 @@ import (
 	"sync"
 
 	db "github.com/vesoft-inc/nebula-studio/server/api/studio/internal/model"
+	"github.com/vesoft-inc/nebula-studio/server/api/studio/pkg/base"
 	"github.com/zeromicro/go-zero/core/logx"
 	"gorm.io/gorm"
 )
@@ -15,7 +16,7 @@ var (
 
 func RunJobs(jobs []*db.LLMJob, JobRunnerMap map[string]func(job *db.LLMJob)) {
 	for _, job := range jobs {
-		job.Status = db.LLMStatusRunning
+		job.Status = base.LLMStatusRunning
 		err := db.CtxDB.Save(job).Error
 		if err != nil {
 			logx.Errorf("failed to update job status: %v", err)
@@ -42,12 +43,12 @@ func IsRunningJobStopped(jobID string) bool {
 	defer mu.Unlock()
 	job, ok := RunningJobMap[jobID]
 	if ok {
-		return job.Status == db.LLMStatusCancel
+		return job.Status == base.LLMStatusCancel
 	}
 	return false
 }
 
-func ChangeRunningJobStatus(jobID string, status db.LLMStatus) {
+func ChangeRunningJobStatus(jobID string, status base.LLMStatus) {
 	mu.Lock()
 	defer mu.Unlock()
 	job, ok := RunningJobMap[jobID]
@@ -68,7 +69,7 @@ func GetRunningJob(jobID string) *db.LLMJob {
 
 func GetPendingJobs() []*db.LLMJob {
 	var jobs []*db.LLMJob
-	err := db.CtxDB.Where("status = ?", db.LLMStatusPending).Find(&jobs).Error
+	err := db.CtxDB.Where("status = ?", base.LLMStatusPending).Find(&jobs).Error
 	if err != nil && err != gorm.ErrRecordNotFound {
 		logx.Errorf("failed to get pending jobs: %v", err)
 	}
diff --git a/server/api/studio/pkg/llm/transformer/openai.go b/server/api/studio/pkg/llm/transformer/openai.go
index af03866b..8c211f8c 100644
--- a/server/api/studio/pkg/llm/transformer/openai.go
+++ b/server/api/studio/pkg/llm/transformer/openai.go
@@ -57,7 +57,7 @@ func (o *OpenAI) HandleResponse(resp *http.Response, callback func(str string))
 		if err := scanner.Err(); err != nil {
 			fmt.Println("reading standard input:", err)
 		}
-		callback("[DONE]")
+		callback("[DONE]") // ensure has done
 		return nil, nil
 	}
 	// Read the response data