Merge pull request #103 from kuafuai/feat/bill

Feat/bill

backend/app/controllers/__init__.py

@@ -9,6 +9,7 @@
 from .step_requirement import bp as step_requirement_bp
 from .setting import bp as setting_bp
 from .tenant_pro import bp as tenant_bp
+from .pay_pro import bp as pay_bp
 
 def register_controllers(app):
     app.register_blueprint(user_bp)
@@ -21,4 +22,5 @@ def register_controllers(app):
     app.register_blueprint(step_code_bp)
     app.register_blueprint(step_devops_bp)
     app.register_blueprint(setting_bp)
-    app.register_blueprint(tenant_bp)
+    app.register_blueprint(tenant_bp)
+    app.register_blueprint(pay_bp)

backend/app/controllers/pay_pro.py

@@ -0,0 +1,6 @@
+from flask import Blueprint
+
+bp = Blueprint('pay', __name__, url_prefix='/pay')
+
+def test():
+    pass

backend/app/controllers/requirement.py

@@ -38,8 +38,10 @@ def setup_app():
     username = session['username']
     tenantID = session['tenant_id']
 
-    if GRADE != "base" and not Tenant.check_quota(tenantID):
-        raise Exception(_("You have exceeded your quota limit, please check your business bill."))
+    if GRADE != "base":
+        passed, msg = Tenant.check_quota(tenantID)
+        if not passed:
+            raise Exception(msg)
 
     requirement = Requirement.create_requirement(tenantID, "New requirement", "New", appID, username, sourceBranch, featureBranch,  REQUIREMENT_STATUS_NotStarted, 0, 0)
 

backend/app/controllers/step_requirement.py

@@ -30,8 +30,7 @@ def clarify():
     if len(globalContext) < 4 :
         Requirement.update_requirement(requirement_id=requirementID, requirement_name=userPrompt, status=REQUIREMENT_STATUS_InProgress)
 
-        if GRADE != "base" and not Tenant.check_quota(tenantID):
-            raise Exception(_("You have exceeded your quota limit, please check your business bill."))
+        # 开始创建一个需求账单
         if GRADE != "base":
             TenantBill.record_requirement(tenantID, userName, requirementID, userPrompt)
 

backend/app/controllers/user.py

@@ -67,7 +67,7 @@ def login():
         raise Exception(_("Invalid username or password"))
 
 
-@bp.route('/logout', methods=['POST'])
+@bp.route('/logout', methods=['GET'])
 @json_response
 def logout():
     _ = getI18n("controllers")

backend/app/controllers/workspace.py

@@ -1,7 +1,7 @@
 from flask import request, session
 from app.controllers.common import json_response
 from app.pkgs.tools.i18b import getI18n
-from app.pkgs.devops.git_tools import pullCode, pushCode
+from app.pkgs.devops.git_tools import pullCode, pushCode, gitResetWorkspace
 from app.pkgs.knowledge.app_info import getServiceGitPath
 from app.models.setting import getGitConfigList
 from app.models.requirement import Requirement
@@ -78,3 +78,32 @@ def gitpush():
         return _("Push code successfully.") + f" from {wsPath} to {gitPath} {fatureBranch}"
     else:
         raise Exception(_("Failed to push code.")+f"In the {wsPath}/{gitPath} directory, {msg}")
+
+
+@bp.route('/resetWorkspace', methods=['POST'])
+@json_response
+def resetWorkspace():
+    _ = getI18n("controllers")
+    username = session['username']
+    requirementID =  request.json.get('task_id')
+    serviceName = request.json.get('service_name')
+    wsPath = get_ws_path(requirementID)
+    req = Requirement.get_requirement_by_id(requirementID) 
+    commitMsg = req["requirement_name"]
+    fatureBranch = req["default_target_branch"]
+    gitPath, success = getServiceGitPath(req["app_id"], serviceName)
+    tenantID = session['tenant_id']
+    username = session['username']
+    gitConfigList, success = getGitConfigList(tenantID, req["app_id"])
+
+    Requirement.update_requirement(requirement_id=requirementID, status=REQUIREMENT_STATUS_Completed)
+
+    if not GIT_ENABLED:
+        raise Exception(_("Failed to reset code.")+f" You did not set Git parameters in the configuration file.")
+    else:
+        success, msg = gitResetWorkspace(wsPath, gitPath, fatureBranch, commitMsg, gitConfigList)
+
+    if success:
+        return True
+    else:
+        raise Exception(_("Failed to reset code.")+f"In the {wsPath}/{gitPath} directory, {msg}")

backend/app/pkgs/devops/git_tools.py

@@ -10,16 +10,29 @@ def pullCode(ws_path, repo_path, base_branch, feature_branch, gitConfigList):
 
     gitUrl = genCloneUrl(repo_path, gitConfig["git_url"], gitConfig["git_username"], gitConfig["git_token"])
     print(f"pullCode start {gitUrl} {base_branch} {repo_path} {ws_path}")
-    result = subprocess.run(['git', 'clone', '-b', base_branch, gitUrl, repo_path], capture_output=True, text=True, cwd=ws_path)
-    if result.returncode != 0:
-        print(result.stderr)
-        return False, "git clone failed: "+result.stderr
+    # 先从feature_branch拉代码，如果失败再从base_branch拉
+    try:
+        result = subprocess.run(['git', 'clone', '-b', feature_branch, gitUrl, repo_path], capture_output=True, text=True, cwd=ws_path)
+        if result.returncode != 0:
+            print(result.stderr)
+            return False, "git clone feature_branch failed: "+result.stderr
+    except Exception as e:
+        result = subprocess.run(['git', 'clone', '-b', base_branch, gitUrl, repo_path], capture_output=True, text=True, cwd=ws_path)
+        if result.returncode != 0:
+            print(result.stderr)
+            return False, "git clone base_branch failed: "+result.stderr
 
+        result = subprocess.run(
+            ['git', 'checkout', '-b', feature_branch], capture_output=True, text=True, cwd=ws_path+'/'+repo_path)
+        if result.returncode != 0:
+            print(result.stderr)
+            return False, "git checkout branch failed: "+result.stderr
+
     result = subprocess.run(
-        ['git', 'checkout', '-b', feature_branch], capture_output=True, text=True, cwd=ws_path+'/'+repo_path)
+        ['git', 'config', 'pull.rebase', 'false'], capture_output=True, text=True, cwd=ws_path+'/'+repo_path)
     if result.returncode != 0:
         print(result.stderr)
-        return False, "git checkout branch failed: "+result.stderr
+        return False, "git config pull.rebase false failed: "+result.stderr
 
     print(f"Code clone success. in {ws_path}")
     return True, f"Code clone success. in {ws_path}"
@@ -68,3 +81,27 @@ def genCloneUrl(gitPath, gitUrl, username, token):
     finalUrl = f"{prot}//{username}:{token}@{domain}/{gitPath}.git"
 
     return finalUrl
+
+def gitResetWorkspace(wsPath, gitPath, fatureBranch, commitMsg, gitConfigList):
+    gitConfig = gitConfigList[0]
+    gitCwd = wsPath+'/'+gitPath
+
+    try:
+        os.makedirs(gitCwd, exist_ok=True)
+    except Exception as e:
+        return False, "mkdir failed: "+str(e)
+
+    result = subprocess.run(
+        ['git', 'fetch', 'origin', fatureBranch], capture_output=True, text=True, cwd=gitCwd)
+    if result.returncode != 0:
+        print(result.stderr)
+        return False, "git fetch origin fatureBranch false failed: "+result.stderr
+
+    result = subprocess.run(
+        ['git', 'reset', '--hard', 'origin/'+fatureBranch], capture_output=True, text=True, cwd=gitCwd)
+    if result.returncode != 0:
+        print(result.stderr)
+        return False, "git reset --hard origin fatureBranch false failed: "+result.stderr
+
+    print(f"reset code success. in {wsPath}")
+    return True, f"reset code success. in {wsPath}"

backend/app/pkgs/prompt/api_basic.py

@@ -23,7 +23,8 @@ def step2GenApiDoc(message, context):
         
 Without any dialogue or explanation, just output the final API only."""})
 
-    return chatCompletion(context, FAKE_API)
+    message, total_tokens, success = chatCompletion(context, FAKE_API)
+    return message, success
 
 def step1ApiDocTasks(user_prompt, apiDoc):
     context = []
@@ -48,7 +49,7 @@ def step1ApiDocTasks(user_prompt, apiDoc):
 """
 
     context.append({"role": "system", "content": content})
-    message, success = chatCompletion(context, FAKE_API)
+    message, total_tokens, success = chatCompletion(context, FAKE_API)
 
     return message, context, success
 

backend/app/pkgs/prompt/code_basic.py

@@ -31,7 +31,7 @@ def aiReferenceRepair(self, requirementID, newCode, referenceCode, fileTask, fil
     """
 
         context = [{"role": "user", "content": prompt}]
-        data, success = chatCompletion(context)
+        data, total_tokens, success = chatCompletion(context)
         return json.loads(fix_llm_json_str(data)), success
 
 
@@ -51,7 +51,7 @@ def aiAnalyzeError(self, requirementID, message, filePath):
     """
 
         context = [{"role": "user", "content": prompt}]
-        data, success = chatCompletion(context)
+        data, total_tokens, success = chatCompletion(context)
         return json.loads(fix_llm_json_str(data)), success
 
 
@@ -76,7 +76,7 @@ def aiFixError(self, requirementID, solution, code, filePath, type):
     """
 
         context = [{"role": "user", "content": prompt}]
-        data, success = chatCompletion(context)
+        data, total_tokens, success = chatCompletion(context)
         return json.loads(fix_llm_json_str(data)), success
 
 
@@ -107,7 +107,7 @@ def aiCheckCode(self, requirementID, fileTask, code, filePath):
         """
 
         context = [{"role": "user", "content": prompt}]
-        data, success = chatCompletion(context)
+        data, total_tokens, success = chatCompletion(context)
         newCode = get_code_from_str(data)
         if len(newCode) < len(code)/3*2:
             jsonData["code"] = code
@@ -141,7 +141,7 @@ def aiReviewCode(self, requirementID, fileTask, code, filePath):
     """
 
         context = [{"role": "user", "content": prompt}]
-        data, success = chatCompletion(context)
+        data, total_tokens, success = chatCompletion(context)
 
         return data, success
 
@@ -172,7 +172,7 @@ def aiMergeCode(self, requirementID, task, baseCode, newCode, filePath):
     """
 
         context = [{"role": "user", "content": prompt}]
-        data, success = chatCompletion(context)
+        data, total_tokens, success = chatCompletion(context)
         return json.loads(fix_llm_json_str(data)), success
 
 
@@ -199,5 +199,5 @@ def aiGenCode(self, requirementID, fileTask, newTask, newCode, filePath):
     """
 
         context = [{"role": "user", "content": prompt}]
-        data, success = chatCompletion(context)
+        data, total_tokens, success = chatCompletion(context)
         return json.loads(fix_llm_json_str(data)), success

backend/app/pkgs/prompt/prompt.py

@@ -1,3 +1,4 @@
+from flask import session
 from app.pkgs.prompt.subtask_pro import SubtaskPro
 from app.pkgs.prompt.requirement_basic import RequirementBasic
 from app.pkgs.prompt.requirement_pro import RequirementPro
@@ -8,8 +9,28 @@
 from app.pkgs.prompt.subtask_vue_pro import SubtaskVuePro
 from app.pkgs.prompt.code_basic import CodeBasic
 from app.pkgs.prompt.code_pro import CodePro
+from app.models.tenant_pro import Tenant
 from config import GRADE
 
+def pre_check_quota(func):
+    def wrapper(*args, **kwargs):
+        # 在方法调用前执行的代码
+        if GRADE != "base":
+            tenantID = session['tenant_id']
+            passed, msg = Tenant.check_quota(tenantID)
+            if not passed:
+                raise Exception(msg)
+            print("pre_check_quota===========")
+            print(passed)
+            print(tenantID)
+
+        # 调用原始方法
+        result = func(*args, **kwargs)
+        # 返回原始方法的结果
+        return result
+    return wrapper
+
+@pre_check_quota
 def clarifyRequirement(requirementID, userPrompt, globalContext, appArchitecture):
     if GRADE == "base":
         obj = RequirementBasic()
@@ -18,6 +39,7 @@ def clarifyRequirement(requirementID, userPrompt, globalContext, appArchitecture
 
     return obj.clarifyRequirement(requirementID, userPrompt, globalContext, appArchitecture)
 
+@pre_check_quota
 def clarifyAPI(requirementID, userPrompt, apiDoc):
     if GRADE == "base":
         obj = ApiBasic()
@@ -26,6 +48,7 @@ def clarifyAPI(requirementID, userPrompt, apiDoc):
 
     return obj.clarifyAPI(requirementID, userPrompt, apiDoc)
 
+@pre_check_quota
 def splitTask(requirementID, newfeature, serviceName, appBasePrompt, projectInfo, projectLib, serviceStruct, appID):
     if GRADE == "base":
         obj = SubtaskBasic()
@@ -39,6 +62,7 @@ def splitTask(requirementID, newfeature, serviceName, appBasePrompt, projectInfo
 
     return obj.splitTask(requirementID, newfeature, serviceName, appBasePrompt, projectInfo, projectLib, serviceStruct, appID)
 
+@pre_check_quota
 def aiReferenceRepair(requirementID, newCode, referenceCode, fileTask, filePath):
     if GRADE == "base":
         obj = CodeBasic()
@@ -47,6 +71,7 @@ def aiReferenceRepair(requirementID, newCode, referenceCode, fileTask, filePath)
 
     return obj.aiReferenceRepair(requirementID, newCode, referenceCode, fileTask, filePath)
 
+@pre_check_quota
 def aiAnalyzeError(requirementID, message, filePath):
     if GRADE == "base":
         obj = CodeBasic()
@@ -55,6 +80,7 @@ def aiAnalyzeError(requirementID, message, filePath):
 
     return obj.aiAnalyzeError(requirementID, message, filePath)
 
+@pre_check_quota
 def aiFixError(requirementID, solution, code, filePath, type):
     if GRADE == "base":
         obj = CodeBasic()
@@ -63,6 +89,7 @@ def aiFixError(requirementID, solution, code, filePath, type):
 
     return obj.aiFixError(requirementID, solution, code, filePath, type)
 
+@pre_check_quota
 def aiCheckCode(requirementID, fileTask, code, filePath):
     if GRADE == "base":
         obj = CodeBasic()
@@ -71,6 +98,7 @@ def aiCheckCode(requirementID, fileTask, code, filePath):
 
     return obj.aiCheckCode(requirementID, fileTask, code, filePath)
 
+@pre_check_quota
 def aiMergeCode(requirementID, fileTask, appName, baseCode, newCode, filePath):
     if GRADE == "base":
         obj = CodeBasic()
@@ -79,6 +107,7 @@ def aiMergeCode(requirementID, fileTask, appName, baseCode, newCode, filePath):
 
     return obj.aiMergeCode(requirementID, fileTask, appName, baseCode, newCode, filePath)
 
+@pre_check_quota
 def aiGenCode(requirementID, fileTask, newTask, newCode, filePath):
     if GRADE == "base":
         obj = CodeBasic()

backend/app/pkgs/prompt/requirement_basic.py

@@ -68,7 +68,7 @@ def clarifyRequirement(self, requirementID, userPrompt, globalContext, appArchit
             ]
             finalContext.extend(preContext)
 
-            message, success = chatCompletion(finalContext, FAKE_CLARIFY_1)
+            message, total_tokens, success = chatCompletion(finalContext, FAKE_CLARIFY_1)
 
             if message.find("Nothing more to clarify") != -1 or message.find('"question":""') != -1 :
                 return organize(firstPrompt, requirementsDetail)
@@ -120,7 +120,7 @@ def organize(firstPrompt, requirementsDetail, appArchitecture):
 """
     })
 
-    message, success = chatCompletion(Organize, FAKE_CLARIFY_2)
+    message, total_tokens, success = chatCompletion(Organize, FAKE_CLARIFY_2)
     print("------")
     print(message)
     message = fix_llm_json_str(message)