feat: auto rule generator (#273)
John Wang
3 years ago
committed by
GitHub
parent
44a1aa5e44
commit
490858a4d5
@ -0,0 +1,75 @@
|
||||
from flask_login import login_required, current_user
|
||||
from flask_restful import Resource, reqparse
|
||||
|
||||
from controllers.console import api
|
||||
from controllers.console.app.error import ProviderNotInitializeError, ProviderQuotaExceededError, \
|
||||
CompletionRequestError, ProviderModelCurrentlyNotSupportError
|
||||
from controllers.console.setup import setup_required
|
||||
from controllers.console.wraps import account_initialization_required
|
||||
from core.generator.llm_generator import LLMGenerator
|
||||
from core.llm.error import ProviderTokenNotInitError, QuotaExceededError, LLMBadRequestError, LLMAPIConnectionError, \
|
||||
LLMAPIUnavailableError, LLMRateLimitError, LLMAuthorizationError, ModelCurrentlyNotSupportError
|
||||
|
||||
|
||||
class IntroductionGenerateApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def post(self):
|
||||
parser = reqparse.RequestParser()
|
||||
parser.add_argument('prompt_template', type=str, required=True, location='json')
|
||||
args = parser.parse_args()
|
||||
|
||||
account = current_user
|
||||
|
||||
try:
|
||||
answer = LLMGenerator.generate_introduction(
|
||||
account.current_tenant_id,
|
||||
args['prompt_template']
|
||||
)
|
||||
except ProviderTokenNotInitError:
|
||||
raise ProviderNotInitializeError()
|
||||
except QuotaExceededError:
|
||||
raise ProviderQuotaExceededError()
|
||||
except ModelCurrentlyNotSupportError:
|
||||
raise ProviderModelCurrentlyNotSupportError()
|
||||
except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
|
||||
LLMRateLimitError, LLMAuthorizationError) as e:
|
||||
raise CompletionRequestError(str(e))
|
||||
|
||||
return {'introduction': answer}
|
||||
|
||||
|
||||
class RuleGenerateApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def post(self):
|
||||
parser = reqparse.RequestParser()
|
||||
parser.add_argument('audiences', type=str, required=True, nullable=False, location='json')
|
||||
parser.add_argument('hoping_to_solve', type=str, required=True, nullable=False, location='json')
|
||||
args = parser.parse_args()
|
||||
|
||||
account = current_user
|
||||
|
||||
try:
|
||||
rules = LLMGenerator.generate_rule_config(
|
||||
account.current_tenant_id,
|
||||
args['audiences'],
|
||||
args['hoping_to_solve']
|
||||
)
|
||||
except ProviderTokenNotInitError:
|
||||
raise ProviderNotInitializeError()
|
||||
except QuotaExceededError:
|
||||
raise ProviderQuotaExceededError()
|
||||
except ModelCurrentlyNotSupportError:
|
||||
raise ProviderModelCurrentlyNotSupportError()
|
||||
except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
|
||||
LLMRateLimitError, LLMAuthorizationError) as e:
|
||||
raise CompletionRequestError(str(e))
|
||||
|
||||
return rules
|
||||
|
||||
|
||||
api.add_resource(IntroductionGenerateApi, '/introduction-generate')
|
||||
api.add_resource(RuleGenerateApi, '/rule-generate')
|
||||
@ -0,0 +1,32 @@
|
||||
from typing import Any
|
||||
|
||||
from langchain.schema import BaseOutputParser, OutputParserException
|
||||
from core.prompt.prompts import RULE_CONFIG_GENERATE_TEMPLATE
|
||||
from libs.json_in_md_parser import parse_and_check_json_markdown
|
||||
|
||||
|
||||
class RuleConfigGeneratorOutputParser(BaseOutputParser):
|
||||
|
||||
def get_format_instructions(self) -> str:
|
||||
return RULE_CONFIG_GENERATE_TEMPLATE
|
||||
|
||||
def parse(self, text: str) -> Any:
|
||||
try:
|
||||
expected_keys = ["prompt", "variables", "opening_statement"]
|
||||
parsed = parse_and_check_json_markdown(text, expected_keys)
|
||||
if not isinstance(parsed["prompt"], str):
|
||||
raise ValueError("Expected 'prompt' to be a string.")
|
||||
if not isinstance(parsed["variables"], list):
|
||||
raise ValueError(
|
||||
f"Expected 'variables' to be a list."
|
||||
)
|
||||
if not isinstance(parsed["opening_statement"], str):
|
||||
raise ValueError(
|
||||
f"Expected 'opening_statement' to be a str."
|
||||
)
|
||||
return parsed
|
||||
except Exception as e:
|
||||
raise OutputParserException(
|
||||
f"Parsing text\n{text}\n of rule config generator raised following error:\n{e}"
|
||||
)
|
||||
|
||||
@ -0,0 +1,38 @@
|
||||
import json
|
||||
from typing import List
|
||||
|
||||
from langchain.schema import OutputParserException
|
||||
|
||||
|
||||
def parse_json_markdown(json_string: str) -> dict:
|
||||
# Remove the triple backticks if present
|
||||
json_string = json_string.strip()
|
||||
start_index = json_string.find("```json")
|
||||
end_index = json_string.find("```", start_index + len("```json"))
|
||||
|
||||
if start_index != -1 and end_index != -1:
|
||||
extracted_content = json_string[start_index + len("```json"):end_index].strip()
|
||||
|
||||
# Parse the JSON string into a Python dictionary
|
||||
parsed = json.loads(extracted_content)
|
||||
elif json_string.startswith("{"):
|
||||
# Parse the JSON string into a Python dictionary
|
||||
parsed = json.loads(json_string)
|
||||
else:
|
||||
raise Exception("Could not find JSON block in the output.")
|
||||
|
||||
return parsed
|
||||
|
||||
|
||||
def parse_and_check_json_markdown(text: str, expected_keys: List[str]) -> dict:
|
||||
try:
|
||||
json_obj = parse_json_markdown(text)
|
||||
except json.JSONDecodeError as e:
|
||||
raise OutputParserException(f"Got invalid JSON object. Error: {e}")
|
||||
for key in expected_keys:
|
||||
if key not in json_obj:
|
||||
raise OutputParserException(
|
||||
f"Got invalid return object. Expected key `{key}` "
|
||||
f"to be present, but got {json_obj}"
|
||||
)
|
||||
return json_obj
|
||||
Loading…
Reference in New Issue