diff --git a/pieces/TextTaggingPiece/models.py b/pieces/TextTaggingPiece/models.py index 43c6c17..49bf9d1 100644 --- a/pieces/TextTaggingPiece/models.py +++ b/pieces/TextTaggingPiece/models.py @@ -31,7 +31,6 @@ class InputModel(BaseModel): """ input_text: str = Field( description='Source text to be tagged.', - json_schema_extra={"from_upstream": "always"} ) openai_model: LLMModelType = Field( default=LLMModelType.gpt_3_5_turbo, diff --git a/pieces/TextTaggingPiece/piece.py b/pieces/TextTaggingPiece/piece.py index 90b637e..c1b73da 100644 --- a/pieces/TextTaggingPiece/piece.py +++ b/pieces/TextTaggingPiece/piece.py @@ -48,7 +48,23 @@ def piece_function(self, input_data: InputModel, secrets_data: SecretsModel): schema=schema, llm=llm ) - res = chain.run(input_text) + # Format display result as markdown + self.format_display_result(input_data, res) + return OutputModel(**res) + + def format_display_result(self, input_data: InputModel, result: dict): + md_text = """## Tags\n""" + for k, v in result.items(): + md_text += f"""**{k}**: {v}\n""" + md_text += "\n" + md_text += f"""**Input text**: {input_data.input_text}\n""" + file_path = f"{self.results_path}/display_result.md" + with open(file_path, "w") as f: + f.write(md_text) + self.display_result = { + "file_type": "md", + "file_path": file_path + }