Skip to content

Commit

Permalink
openapi support
Browse files Browse the repository at this point in the history
  • Loading branch information
@imalsogreg
imalsogreg committed Oct 17, 2024
1 parent b36660d commit 523bb33
Show file tree
Hide file tree
Showing 7 changed files with 280 additions and 81 deletions.
80 changes: 78 additions & 2 deletions engine/language_client_codegen/src/openapi.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use internal_baml_core::ir::{
use serde::Serialize;
use serde_json::json;

use crate::dir_writer::{FileCollector, LanguageFeatures, RemoveDirBehavior};
use crate::{dir_writer::{FileCollector, LanguageFeatures, RemoveDirBehavior}, field_type_attributes, TypeCheckAttributes};

#[derive(Default)]
pub(super) struct OpenApiLanguageFeatures {}
Expand Down Expand Up @@ -231,6 +231,17 @@ impl Serialize for OpenApiSchema<'_> {
},
"required": ["name", "provider", "options"]
})
),
( "Check",
json!({
"type": "object",
"properties": {
"name": { "type": "string" },
"expr": { "type": "string" },
"status": { "type": "string" }
}

})
)
]
.into_iter()
Expand Down Expand Up @@ -333,6 +344,40 @@ impl<'ir> TryFrom<(&'ir IntermediateRepr, &'_ crate::GeneratorArgs)> for OpenApi
}
}

pub fn type_name_for_checks(checks: &TypeCheckAttributes) -> String {
let mut name = "Checks".to_string();
let mut names: Vec<&String> = checks.0.iter().collect();
names.sort();
for check_name in names.iter() {
name.push_str("__");
name.push_str(check_name);
}
name
}

fn check() -> TypeSpecWithMeta {
TypeSpecWithMeta {
meta: TypeMetadata::default(),
type_spec: TypeSpec::Ref{ r#ref: "#components/schemas/Check".to_string() },
}
}

/// The type definition for a single "Checked_*" type. Note that we don't
/// produce a named type for each of these the way we do for SDK
/// codegeneration.
fn type_def_for_checks(checks: TypeCheckAttributes) -> TypeSpecWithMeta {
TypeSpecWithMeta {
meta: TypeMetadata::default(),
type_spec: TypeSpec::Inline(
TypeDef::Class {
properties: checks.0.iter().map(|check_name| (check_name.clone(), check())).collect(),
required: checks.0.into_iter().collect(),
additional_properties: false,
}
)
}
}

impl<'ir> TryFrom<Walker<'ir, &'ir Node<Function>>> for OpenApiMethodDef<'ir> {
type Error = anyhow::Error;

Expand Down Expand Up @@ -597,7 +642,27 @@ impl<'ir> ToTypeReferenceInTypeDefinition<'ir> for FieldType {
// something i saw suggested doing this
type_spec
}
FieldType::Constrained{base,..} => base.to_type_spec(ir)?,
FieldType::Constrained{base,..} => {
match field_type_attributes(self) {
Some(checks) => {
let base_type_ref = base.to_type_spec(ir)?;
let checks_type_spec = type_def_for_checks(checks);
TypeSpecWithMeta {
meta: TypeMetadata::default(),
type_spec: TypeSpec::Inline(
TypeDef::Class {
properties: vec![("value".to_string(), base_type_ref),("checks".to_string(), checks_type_spec)].into_iter().collect(),
required: vec!["value".to_string(), "checks".to_string()],
additional_properties: false,
}
)
}
}
None => {
base.to_type_spec(ir)?
}
}
},
})
}
}
Expand Down Expand Up @@ -632,6 +697,17 @@ struct TypeMetadata {
nullable: bool,
}

impl Default for TypeMetadata {
fn default() -> Self {
TypeMetadata {
title: None,
r#enum: None,
r#const: None,
nullable: false,
}
}
}

#[derive(Clone, Debug, Serialize)]
#[serde(untagged)]
enum TypeSpec {
Expand Down
137 changes: 130 additions & 7 deletions integ-tests/openapi/baml_client/openapi.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -593,7 +593,22 @@ paths:
application/json:
schema:
title: PredictAgeBareResponse
type: integer
type: object
properties:
value:
type: integer
checks:
type: object
properties:
too_big:
$ref: '#components/schemas/Check'
required:
- too_big
additionalProperties: false
required:
- value
- checks
additionalProperties: false
operationId: PredictAgeBare
/call/PredictAgeComplex:
post:
Expand Down Expand Up @@ -2648,6 +2663,15 @@ components:
- name
- provider
- options
Check:
type: object
properties:
name:
type: string
expr:
type: string
status:
type: string
Category:
enum:
- Refund
Expand Down Expand Up @@ -2910,7 +2934,25 @@ components:
type: object
properties:
age:
type: integer
type: object
properties:
value:
type: integer
checks:
type: object
properties:
no_infants:
$ref: '#components/schemas/Check'
earth_aged:
$ref: '#components/schemas/Check'
required:
- no_infants
- earth_aged
additionalProperties: false
required:
- value
- checks
additionalProperties: false
required:
- age
additionalProperties: false
Expand Down Expand Up @@ -2953,7 +2995,22 @@ components:
type: object
properties:
value:
type: string
type: object
properties:
value:
type: string
checks:
type: object
properties:
valid_email:
$ref: '#components/schemas/Check'
required:
- valid_email
additionalProperties: false
required:
- value
- checks
additionalProperties: false
required:
- value
additionalProperties: false
Expand Down Expand Up @@ -3010,9 +3067,45 @@ components:
- $ref: '#/components/schemas/Martian'
- $ref: '#/components/schemas/Earthling'
certainty:
type: integer
type: object
properties:
value:
type: integer
checks:
type: object
properties:
unreasonably_certain:
$ref: '#components/schemas/Check'
required:
- unreasonably_certain
additionalProperties: false
required:
- value
- checks
additionalProperties: false
species:
type: string
type: object
properties:
value:
type: string
checks:
type: object
properties:
trivial:
$ref: '#components/schemas/Check'
regex_bad:
$ref: '#components/schemas/Check'
regex_good:
$ref: '#components/schemas/Check'
required:
- trivial
- regex_bad
- regex_good
additionalProperties: false
required:
- value
- checks
additionalProperties: false
required:
- planetary_age
- certainty
Expand Down Expand Up @@ -3080,7 +3173,22 @@ components:
type: object
properties:
age:
type: integer
type: object
properties:
value:
type: integer
checks:
type: object
properties:
young_enough:
$ref: '#components/schemas/Check'
required:
- young_enough
additionalProperties: false
required:
- value
- checks
additionalProperties: false
required:
- age
additionalProperties: false
Expand Down Expand Up @@ -3173,7 +3281,22 @@ components:
type: object
properties:
value:
type: string
type: object
properties:
value:
type: string
checks:
type: object
properties:
valid_phone_number:
$ref: '#components/schemas/Check'
required:
- valid_phone_number
additionalProperties: false
required:
- value
- checks
additionalProperties: false
required:
- value
additionalProperties: false
Expand Down
2 changes: 1 addition & 1 deletion integ-tests/python/baml_client/partial_types.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from . import types


from .types import Checks__unreasonably_certain, Checks__regex_bad__regex_good__trivial, Checks__valid_email, Checks__valid_phone_number, Checks__too_big, Checks__earth_aged__no_infants, Checks__young_enough
from .types import Checks__unreasonably_certain, Checks__young_enough, Checks__valid_email, Checks__regex_bad__regex_good__trivial, Checks__earth_aged__no_infants, Checks__valid_phone_number, Checks__too_big


###############################################################################
Expand Down
28 changes: 14 additions & 14 deletions integ-tests/python/baml_client/types.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,30 +119,30 @@ class TestEnum(str, Enum):
F = "F"
G = "G"

class Checks__valid_email(BaseModel):
valid_email: baml_py.Check

class Checks__valid_phone_number(BaseModel):
valid_phone_number: baml_py.Check

class Checks__young_enough(BaseModel):
young_enough: baml_py.Check

class Checks__too_big(BaseModel):
too_big: baml_py.Check

class Checks__earth_aged__no_infants(BaseModel):
earth_aged: baml_py.Check
no_infants: baml_py.Check
earth_aged: baml_py.Check

class Checks__unreasonably_certain(BaseModel):
unreasonably_certain: baml_py.Check

class Checks__too_big(BaseModel):
too_big: baml_py.Check

class Checks__regex_bad__regex_good__trivial(BaseModel):
regex_bad: baml_py.Check
trivial: baml_py.Check
regex_bad: baml_py.Check
regex_good: baml_py.Check

class Checks__valid_email(BaseModel):
valid_email: baml_py.Check

class Checks__young_enough(BaseModel):
young_enough: baml_py.Check

class Checks__valid_phone_number(BaseModel):
valid_phone_number: baml_py.Check

class BigNumbers(BaseModel):


Expand Down
Loading

0 comments on commit 523bb33

Please sign in to comment.