feat(新功能): 项目信息提取/生成接口

fix(修复bug):
docs(文档变更):
refactor(重构):
test(增加测试):
This commit is contained in:
zchengrong
2025-05-15 14:49:33 +08:00
parent 3095d2654e
commit 6cb32d11a8
4 changed files with 88 additions and 2 deletions

View File

@@ -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())