feat(新功能): 项目信息提取/生成接口
fix(修复bug): docs(文档变更): refactor(重构): test(增加测试):
This commit is contained in:
@@ -0,0 +1,47 @@
|
||||
from langchain.output_parsers import ResponseSchema, StructuredOutputParser
|
||||
from langchain_community.chat_models import ChatTongyi
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
|
||||
from app.schemas.project_info_extraction import ProjectInfoExtractionModel
|
||||
|
||||
|
||||
class ProjectInfoExtraction:
|
||||
def __init__(self, request_data):
|
||||
# llm generate brand info init
|
||||
self.model = ChatTongyi(model="qwen2.5-14b-instruct", api_key="sk-7658298c6b99443c98184a5e634fe6ab")
|
||||
|
||||
self.response_schemas = [
|
||||
ResponseSchema(name="project_name", description="project name."),
|
||||
ResponseSchema(name="role", description="The target role of the project."),
|
||||
ResponseSchema(name="gender", description="The gender targeted by the project."),
|
||||
ResponseSchema(name="style", description="Project style.")
|
||||
]
|
||||
self.output_parser = StructuredOutputParser.from_response_schemas(self.response_schemas)
|
||||
self.format_instructions = self.output_parser.get_format_instructions()
|
||||
self.prompt = PromptTemplate(
|
||||
template="你是一个时装品牌的设计师助理。根据用户输入提取出project_name,role ,gender ,style ."
|
||||
"gender部分请用以下:menswear,womenswear,childrenwear,如果全部都适用即all."
|
||||
"如果没有以上内容,需要你根据用户输入随意发挥.\n{format_instructions}\n{question}",
|
||||
input_variables=["question"],
|
||||
partial_variables={"format_instructions": self.format_instructions}
|
||||
)
|
||||
self._input = self.prompt.format_prompt(question=request_data.prompt)
|
||||
|
||||
self.result_data = {}
|
||||
|
||||
def get_result(self):
|
||||
self.llm_extraction_project_info()
|
||||
return self.result_data
|
||||
|
||||
def llm_extraction_project_info(self):
|
||||
output = self.model(self._input.to_messages())
|
||||
project_info = self.output_parser.parse(output.content)
|
||||
self.result_data = project_info
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
request_data = ProjectInfoExtractionModel(
|
||||
prompt="海边派对主题的系列设计"
|
||||
)
|
||||
service = ProjectInfoExtraction(request_data)
|
||||
print(service.get_result())
|
||||
Reference in New Issue
Block a user