feat: query prompt template support in chatflow (#3791)

Co-authored-by: Joel <iamjoel007@gmail.com>
This commit is contained in:
takatost
2024-04-25 18:01:53 +08:00
committed by GitHub
parent 80b9507e7a
commit 12435774ca
12 changed files with 113 additions and 18 deletions

View File

@@ -31,7 +31,8 @@ class AdvancedPromptTransform(PromptTransform):
context: Optional[str],
memory_config: Optional[MemoryConfig],
memory: Optional[TokenBufferMemory],
model_config: ModelConfigWithCredentialsEntity) -> list[PromptMessage]:
model_config: ModelConfigWithCredentialsEntity,
query_prompt_template: Optional[str] = None) -> list[PromptMessage]:
inputs = {key: str(value) for key, value in inputs.items()}
prompt_messages = []
@@ -53,6 +54,7 @@ class AdvancedPromptTransform(PromptTransform):
prompt_template=prompt_template,
inputs=inputs,
query=query,
query_prompt_template=query_prompt_template,
files=files,
context=context,
memory_config=memory_config,
@@ -121,7 +123,8 @@ class AdvancedPromptTransform(PromptTransform):
context: Optional[str],
memory_config: Optional[MemoryConfig],
memory: Optional[TokenBufferMemory],
model_config: ModelConfigWithCredentialsEntity) -> list[PromptMessage]:
model_config: ModelConfigWithCredentialsEntity,
query_prompt_template: Optional[str] = None) -> list[PromptMessage]:
"""
Get chat model prompt messages.
"""
@@ -148,6 +151,20 @@ class AdvancedPromptTransform(PromptTransform):
elif prompt_item.role == PromptMessageRole.ASSISTANT:
prompt_messages.append(AssistantPromptMessage(content=prompt))
if query and query_prompt_template:
prompt_template = PromptTemplateParser(
template=query_prompt_template,
with_variable_tmpl=self.with_variable_tmpl
)
prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs}
prompt_inputs['#sys.query#'] = query
prompt_inputs = self._set_context_variable(context, prompt_template, prompt_inputs)
query = prompt_template.format(
prompt_inputs
)
if memory and memory_config:
prompt_messages = self._append_chat_histories(memory, memory_config, prompt_messages, model_config)

View File

@@ -40,3 +40,4 @@ class MemoryConfig(BaseModel):
role_prefix: Optional[RolePrefix] = None
window: WindowConfig
query_prompt_template: Optional[str] = None