Alternative Output Structuring With PredictionGuard

This example demonstrates how to integrate PredictionGuard with custom data structures for output parsing, utilizing the Pydantic library. The focus is on structuring outputs in a specific format, in this case, a joke, ensuring the setup ends with a question mark.

Requirements

• PredictionGuard API
• Pydantic library
• LangChain library

Implementation Steps

1. Set Environment Variable for PredictionGuard: Ensure that your PredictionGuard API token is correctly set up in your environment variables.

1
from predictionguard import PredictionGuard
2
3
4
# Set your Prediction Guard API key and URL as environmental variables.
5
os.environ["PREDICTIONGUARD_API_KEY"] = "<api key>"
6
os.environ["PREDICTIONGUARD_URL"] = "<pg url>"
7
8
# You can also set them when initializing the client.
9
client = PredictionGuard(
10
    api_key="<Your PG API Key>",
11
    url="<Your PG API URL>"
12
)

2. Import Necessary Libraries: Import PredictionGuard, Pydantic for data validation, and LangChain for output parsing and prompt templating.

1
from langchain.output_parsers import PydanticOutputParser
2
from langchain.prompts import PromptTemplate
3
from pydantic import BaseModel, Field, validator

3. Define Your Data Structure: Create a Pydantic model to define the structure of a joke, including setup and punchline, with a validator to ensure the setup is properly formed.

1
class Joke(BaseModel):
2
    setup: str = Field(description="question to set up a joke")
3
    punchline: str = Field(description="answer to resolve the joke")
4
5
    @validator("setup")
6
    def question_ends_with_question_mark(cls, field):
7
        if not field.endswith("?"):
8
            raise ValueError("Badly formed question!")
9
        return field

4. Set Up an Output Parser: Utilize LangChain’s PydanticOutputParser to enforce the data structure defined by the Pydantic model on the output.

1
parser = PydanticOutputParser(pydantic_object=Joke)

5. Create and Format a Prompt: Use LangChain’s PromptTemplate to structure your query, incorporating instructions for output formatting derived from the parser.

1
prompt = PromptTemplate(
2
    template="Answer the user query.\n{format_instructions}\n{query}\n",
3
    input_variables=["query"],
4
    partial_variables={"format_instructions": parser.get_format_instructions()},
5
)

6. Generate and Parse Output: Call PredictionGuard’s text completion model “gpt-oss-120b” to generate an output based on the formatted prompt, then parse the output into the Pydantic model. Handle exceptions for parsing errors.

1
result = client.completions.create(
2
    model="gpt-oss-120b",
3
    prompt=prompt.format(query="Tell me a joke."),
4
    max_completion_tokens=200,
5
    temperature=0.1
6
)
7
8
try:
9
    joke = Joke.parse_raw(result['choices'][0]['text'])
10
    print(joke.setup)
11
    print(joke.punchline)
12
except Exception as e:
13
    print(f"Error parsing joke: {e}")

Let’s Test This Out

After running the implementation with the query “Tell me a joke.”, the structured output generated by PredictionGuard, parsed and validated by our Pydantic model, looks like this.

1
{
2
  "setup": "Why did the tomato turn red?",
3
  "punchline": "Because it saw the salad dressing."
4
}

Output Explanation

• Setup: “Why did the tomato turn red?” - This question successfully passes the validator check, ending with a question mark as required.

• Punchline: “Because it saw the salad dressing.” - Provides a humorous answer to the setup question.

This example demonstrates the effective structuring and validation of output data, ensuring that the generated joke adheres to our defined format.

Conclusion

This approach allows for the flexible and structured generation of outputs, leveraging PredictionGuard’s capabilities alongside Pydantic’s validation features.