import sutro as so
from pydantic import BaseModel

texts = [
    "The airplane is flying at an altitude of 30,000 feet.",
    "The capital of France is Paris.",
    "The best way to cook a steak is on a grill."
]

class AviationOutput(BaseModel):
    is_aviation_related: bool
    answer_justification: str

results = so.infer(texts, system_prompt="Is the following sentence aviation related? Please provide a justification for your answer.", output_schema=AviationOutput)

print(results)

import sutro as so

texts = [
    "The airplane is flying at an altitude of 30,000 feet.",
    "The capital of France is Paris.",
    "The best way to cook a steak is on a grill."
]

json_schema = {
    "type": "object",
    "properties": {
        "is_aviation_related": {
            "type": "boolean"
        },
        "answer_justification": {
            "type": "string"
        }
    },
    "required": ["is_aviation_related", "answer_justification"]
}

results = so.infer(texts, system_prompt="Is the following sentence aviation related? Please provide a justification for your answer.", output_schema=json_schema)

print(results)