JSON Object Generator Tool (#392)

This tool became super useful while building the podcast generator. It’s
a demo of how to perform the task “Make me an object with this schema."

Author: eob

src/steamship/agents/tools/text_generation/__init__.py

@@ -1,4 +1,5 @@
 from .image_prompt_generator_tool import ImagePromptGeneratorTool
+from .json_object_generator import JsonObjectGeneratorTool
 from .personality_tool import PersonalityTool
 from .summarize_text_with_prompt_tool import SummarizeTextWithPromptTool
 from .text_rewrite_tool import TextRewritingTool
@@ -10,4 +11,5 @@
     "SummarizeTextWithPromptTool",
     "TextRewritingTool",
     "TextTranslationTool",
+    "JsonObjectGeneratorTool",
 ]

src/steamship/agents/tools/text_generation/json_object_generator.py

@@ -0,0 +1,170 @@
+import json
+import random
+from typing import Any, List, Union
+
+from steamship import Block, Steamship, SteamshipError, Task
+from steamship.agents.llms import OpenAI
+from steamship.agents.schema import AgentContext, Tool
+from steamship.agents.utils import get_llm, with_llm
+from steamship.utils.repl import ToolREPL
+
+DEFAULT_PROMPT = """INSTRUCTIONS:
+Generate a JSON object describing {table_description}.
+Always return a non-empty value for every field in the object.
+
+FIELDS DESIRED:
+{fields_desired}
+
+EXAMPLE OBJECTS:
+{example_objects}
+
+NEW OBJECT:
+{new_object_prefix}
+"""
+
+DEFAULT_PLURAL_OBJECT_DESCRIPTION = "employees of a company"
+DEFAULT_OBJECT_KEYS = ["Name", "Age", "Gender"]
+DEFAULT_OBJECT_EXAMPLES = [
+    ["Bob", 30, "Male"],
+    ["Susan", 32, "Female"],
+    ["Zhenzhong", 40, "Male"],
+    ["Luis", 32, "Male"],
+    ["Roberta", 35, "Female"],
+    ["Sofia", 30, "Female"],
+]
+DEFAULT_NEW_ROW_PREFIX_FIELDS = []
+
+
+class JsonObjectGeneratorTool(Tool):
+    """
+    Example tool to illustrate generating a new JSON object provided a set of examples.
+
+    This is useful as an example of how to generate a new structured object:
+
+    - A Person (e.g. name, gender, age)
+    - A Proposed Podcast Episode (e.g. title, description, tags)
+
+    The tool takes no input at runtime: it's a true generator parameterized only at initializtion time.
+
+    The tool's parameterization is somewhat CSV-like in construction.
+
+    """
+
+    rewrite_prompt: str = DEFAULT_PROMPT
+    """The prompt used to generate a new JSON object."""
+
+    plural_object_description: str = DEFAULT_PLURAL_OBJECT_DESCRIPTION
+    """Plural description of the object. E.g. 'employees of a company' or 'people' or 'podcast episodes'"""
+
+    object_keys: List[str] = DEFAULT_OBJECT_KEYS
+    """The keys the JSON should have."""
+
+    example_rows: List[List[str]] = DEFAULT_OBJECT_EXAMPLES
+    """List of example object values, aligned to the `object_keys` parameter."""
+
+    new_row_prefix_fields: List[str] = DEFAULT_NEW_ROW_PREFIX_FIELDS
+    """Any fields that should be hard-coded for the new row. These must be grouped as the first set of fields."""
+
+    shuffle_example_rows: bool = True
+    """Whether randomly shuffle example rows to induce a bit of variety even with low LLM temperature."""
+
+    validate_output_as_json: bool = True
+    """Whether to validate that the output is actually JSON."""
+
+    name: str = "JsonObjectTool"
+    human_description: str = "Generates a new JSON object."
+    agent_description: str = "(set at initialization time)"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.agent_description = (
+            f"Used to generate an instance of the {self.plural_object_description} table. "
+            "Input: Anything. "
+            f"Output A tab-separated row describing a new instance of the {self.plural_object_description} table."
+        )
+
+    def kv_clause(self, key: str, value: str) -> str:
+        """Return an escaped, JSON style key-value clause `"key": "value"`"""
+        value = str(value).replace('"', '\\"')
+        clause = f'"{key}": "{value}"'
+        return clause
+
+    def object_json(self, schema: List[str], values: List[str]):
+        """Render a CSV-style header row and value list into a JSON object."""
+        clauses = []
+        for field, value in zip(schema, values):
+            clauses.append(self.kv_clause(field, value))
+
+        return "{" + ", ".join(clauses) + "}"
+
+    def run(self, tool_input: List[Block], context: AgentContext) -> Union[List[Block], Task[Any]]:
+        """Ignore tool input and generate a JSON object described by the tool's configuration.
+
+        Inputs
+        ------
+        input: List[Block]
+            A list of blocks that will be ignored.
+        memory: AgentContext
+            The active AgentContext.
+
+        Output
+        ------
+        output: List[Blocks]
+            A single block containing a new row of the table described by the tool's configuration.
+        """
+
+        if self.shuffle_example_rows:
+            # Shuffle the example rows to get a bit of variety even with low temperature.
+            random.shuffle(self.example_rows)
+
+        # Generate example JSON objects with a fixed key ordering.
+        example_objects = [
+            self.object_json(self.object_keys, example_row) for example_row in self.example_rows
+        ]
+        example_objects_str = "\n".join(example_objects)
+
+        # Generate the new row line. At a minimum it's the `{` character, but it may also hard-code a number of
+        # fields that should be affixed rather than generated.
+        new_object_prefix = "{"
+        for i in range(len(self.new_row_prefix_fields)):
+            clause = self.kv_clause(self.object_keys[i], self.new_row_prefix_fields[i])
+            new_object_prefix += f"{clause}, "
+
+        # Compile the final generation prompt.
+        prompt = self.rewrite_prompt.format(
+            table_description=self.plural_object_description,
+            fields_desired=", ".join(self.object_keys),
+            example_objects=example_objects_str,
+            new_object_prefix=new_object_prefix,
+        )
+
+        # Perform the generation
+        llm = get_llm(context)
+        res = llm.complete(prompt, stop="}")
+
+        # Make sure we only generated one block; anything else violates the assumptions of this code.
+        blocks_emitted = len(res)
+        if blocks_emitted != 1:
+            raise SteamshipError(message=f"{blocks_emitted} blocks emitted; expecting 1.")
+
+        # The output JSON is generation prefix row, plus the generated content, plus a final } character
+        # The reason we have to add the final "}" character is because we used it for the stop character
+        full_json = new_object_prefix + res[0].text + "}"
+
+        if self.validate_output_as_json:
+            try:
+                json.loads(full_json)
+            except BaseException:
+                raise SteamshipError(
+                    message=f"Attempted to generate a JSON object, but did not generate valid JSON. Result: {full_json}"
+                )
+
+        res[0].text = full_json
+        return res
+
+
+if __name__ == "__main__":
+    with Steamship.temporary_workspace() as client:
+        ToolREPL(JsonObjectGeneratorTool()).run_with_client(
+            client=client, context=with_llm(llm=OpenAI(client=client))
+        )