Skip to content

Commit

Permalink
add structured example
Browse files Browse the repository at this point in the history
  • Loading branch information
alex-dixon committed Aug 4, 2024
1 parent e6ce85f commit 64e05bd
Showing 1 changed file with 73 additions and 0 deletions.
73 changes: 73 additions & 0 deletions examples/calculator_structured.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
import dataclasses
import ell
from typing import Any, Literal, Type, Union
import pydantic

from ell.stores.sql import SQLiteStore

ell.config.verbose = True


@dataclasses.dataclass
class Add:
op: Literal["+"]
a: float
b: float


@dataclasses.dataclass
class Sub:
op: Literal["-"]
a: float
b: float


@dataclasses.dataclass
class Mul:
op: Literal["*"]
a: float
b: float


@dataclasses.dataclass
class Div:
op: Literal["/"]
a: float
b: float


CalcOp = Union[Add, Sub, Mul, Div]


@ell.lm(model="gpt-4o", temperature=0.1)
def parse_json(task: str, type: Type[Any]):
return [
ell.system(
f"""You are a JSON parser. You respond only in JSON. Do not format using markdown."""
),
ell.user(
f"""You are given the following task: "{task}"
Parse the task into the following type:
{pydantic.TypeAdapter(type).json_schema()}
"""
)
]


def calc_structured(task: str) -> float:
output = parse_json(task, CalcOp)
structured = pydantic.TypeAdapter(CalcOp).validate_json(output)
match structured.op:
case "+":
return structured.a + structured.b
case "-":
return structured.a - structured.b
case "*":
return structured.a * structured.b
case "/":
return structured.a / structured.b


if __name__ == "__main__":
ell.set_store(SQLiteStore('sqlite_example'), autocommit=True)
print(calc_structured("What is two plus two?"))

0 comments on commit 64e05bd

Please sign in to comment.