Skip to content

Commit

Permalink
Add pydantic example
Browse files Browse the repository at this point in the history
  • Loading branch information
collindutter committed Sep 11, 2024
1 parent 674a61a commit aed5f6c
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 1 deletion.
19 changes: 19 additions & 0 deletions docs/griptape-framework/structures/rulesets.md
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,25 @@ This is particularly useful when you need the LLM to return well-formed data, su
}
```

Although Griptape leverages the `schema` library, you're free to use any JSON schema generation library to define your schema!

For example, using `pydantic`:

```python
--8<-- "docs/griptape-framework/structures/src/json_schema_rule_pydantic.py"
```

```
[09/11/24 09:45:58] INFO PromptTask eae43f52829c4289a6cca9ee7950e075
Input: What is the sentiment of this message?: 'I am so happy!'
INFO PromptTask eae43f52829c4289a6cca9ee7950e075
Output: {
"answer": "The sentiment of the message is positive.",
"relevant_emojis": ["😊", "😄"]
}
answer='The sentiment of the message is positive.' relevant_emojis=['😊', '😄']
```

## Structure

### Rulesets
Expand Down
6 changes: 5 additions & 1 deletion docs/griptape-framework/structures/src/json_schema_rule.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import json

import schema

from griptape.rules.json_schema_rule import JsonSchemaRule
Expand All @@ -11,4 +13,6 @@
]
)

agent.run("What is the sentiment of this message?: 'I am so happy!'")
output = agent.run("What is the sentiment of this message?: 'I am so happy!'").output

print(json.dumps(json.loads(output.value), indent=2))
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
from __future__ import annotations

import pydantic

from griptape.rules.json_schema_rule import JsonSchemaRule
from griptape.structures import Agent


class SentimentModel(pydantic.BaseModel):
answer: str
relevant_emojis: list[str]


agent = Agent(rules=[JsonSchemaRule(SentimentModel.model_json_schema())])

output = agent.run("What is the sentiment of this message?: 'I am so happy!'").output

sentiment_analysis = SentimentModel.model_validate_json(output.value)

# Autocomplete via dot notation 🤩
print(sentiment_analysis.answer)
print(sentiment_analysis.relevant_emojis)

0 comments on commit aed5f6c

Please sign in to comment.