Skip to content

Commit

Permalink
Add codegen for TS
Browse files Browse the repository at this point in the history
  • Loading branch information
@antoniosarosi
antoniosarosi committed Dec 18, 2024
1 parent cac1a16 commit 39141cb
Show file tree
Hide file tree
Showing 6 changed files with 51 additions and 5 deletions.
29 changes: 27 additions & 2 deletions engine/language_client_codegen/src/typescript/generate_types.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@ use anyhow::Result;
use itertools::Itertools;

use internal_baml_core::ir::{
repr::{Docstring, IntermediateRepr},
ClassWalker, EnumWalker,
repr::{Docstring, IntermediateRepr, Walker},
ClassWalker, EnumWalker, FieldType,
};

use crate::{type_check_attributes, GeneratorArgs, TypeCheckAttributes};
Expand All @@ -24,6 +24,7 @@ pub(crate) struct TypeBuilder<'ir> {
pub(crate) struct TypescriptTypes<'ir> {
enums: Vec<TypescriptEnum<'ir>>,
classes: Vec<TypescriptClass<'ir>>,
structural_recursive_alias_cycles: Vec<TypescriptTypeAlias<'ir>>,
}

struct TypescriptEnum<'ir> {
Expand All @@ -40,6 +41,11 @@ pub struct TypescriptClass<'ir> {
pub docstring: Option<String>,
}

struct TypescriptTypeAlias<'ir> {
name: Cow<'ir, str>,
target: String,
}

impl<'ir> TryFrom<(&'ir IntermediateRepr, &'ir GeneratorArgs)> for TypescriptTypes<'ir> {
type Error = anyhow::Error;

Expand All @@ -55,6 +61,10 @@ impl<'ir> TryFrom<(&'ir IntermediateRepr, &'ir GeneratorArgs)> for TypescriptTyp
.walk_classes()
.map(|e| Into::<TypescriptClass>::into(&e))
.collect::<Vec<_>>(),
structural_recursive_alias_cycles: ir
.walk_alias_cycles()
.map(TypescriptTypeAlias::from)
.collect::<Vec<_>>(),
})
}
}
Expand Down Expand Up @@ -132,6 +142,21 @@ impl<'ir> From<&ClassWalker<'ir>> for TypescriptClass<'ir> {
}
}

// TODO: Define AliasWalker to simplify type.
impl<'ir> From<Walker<'ir, (&'ir String, &'ir FieldType)>> for TypescriptTypeAlias<'ir> {
fn from(
Walker {
db,
item: (name, target),
}: Walker<(&'ir String, &'ir FieldType)>,
) -> Self {
Self {
name: Cow::Borrowed(name),
target: target.to_type_ref(db),
}
}
}

pub fn type_name_for_checks(checks: &TypeCheckAttributes) -> String {
checks
.0
Expand Down
5 changes: 5 additions & 0 deletions engine/language_client_codegen/src/typescript/templates/types.ts.j2
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,3 +52,8 @@ export interface {{cls.name}} {
{%- endif %}
}
{% endfor %}

{#- Type Aliases -#}
{% for alias in structural_recursive_alias_cycles %}
type {{alias.name}} = {{alias.target}}
{% endfor %}
2 changes: 1 addition & 1 deletion integ-tests/python/baml_client/inlinedbaml.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@
"test-files/functions/output/recursive-type-aliases.baml": "class LinkedListAliasNode {\n value int\n next LinkedListAliasNode?\n}\n\n// Simple alias that points to recursive type.\ntype LinkedListAlias = LinkedListAliasNode\n\nfunction AliasThatPointsToRecursiveType(list: LinkedListAlias) -> LinkedListAlias {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given linked list back:\n \n {{ list }}\n \n {{ ctx.output_format }}\n \"#\n}\n\n// Class that points to an alias that points to a recursive type.\nclass ClassToRecAlias {\n list LinkedListAlias\n}\n\nfunction ClassThatPointsToRecursiveClassThroughAlias(cls: ClassToRecAlias) -> ClassToRecAlias {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given object back:\n \n {{ cls }}\n \n {{ ctx.output_format }}\n \"#\n}\n\n// This is tricky cause this class should be hoisted, but classes and aliases\n// are two different types in the AST. This test will make sure they can interop.\nclass NodeWithAliasIndirection {\n value int\n next NodeIndirection?\n}\n\ntype NodeIndirection = NodeWithAliasIndirection\n\nfunction RecursiveClassWithAliasIndirection(cls: NodeWithAliasIndirection) -> NodeWithAliasIndirection {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given object back:\n \n {{ cls }}\n \n {{ ctx.output_format }}\n \"#\n}\n",
"test-files/functions/output/serialization-error.baml": "class DummyOutput {\n nonce string\n nonce2 string\n @@dynamic\n}\n\nfunction DummyOutputFunction(input: string) -> DummyOutput {\n client GPT35\n prompt #\"\n Say \"hello there\".\n \"#\n}",
"test-files/functions/output/string-list.baml": "function FnOutputStringList(input: string) -> string[] {\n client GPT35\n prompt #\"\n Return a list of strings in json format like [\"string1\", \"string2\", \"string3\"].\n\n JSON:\n \"#\n}\n\ntest FnOutputStringList {\n functions [FnOutputStringList]\n args {\n input \"example input\"\n }\n}\n",
"test-files/functions/output/type-aliases.baml": "type Primitive = int | string | bool | float\n\ntype List = string[]\n\ntype Graph = map<string, string[]>\n\ntype Combination = Primitive | List | Graph\n\nfunction PrimitiveAlias(p: Primitive) -> Primitive {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given value back: {{ p }}\n \"#\n}\n\nfunction MapAlias(m: Graph) -> Graph {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given Graph back:\n\n {{ m }}\n\n {{ ctx.output_format }}\n \"#\n}\n\nfunction NestedAlias(c: Combination) -> Combination {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given value back:\n\n {{ c }}\n\n {{ ctx.output_format }}\n \"#\n}\n\n// Test attribute merging.\ntype Currency = int @check(gt_ten, {{ this > 10 }})\ntype Amount = Currency @assert({{ this > 0 }})\n\nclass MergeAttrs {\n amount Amount @description(\"In USD\")\n}\n\n// This should be allowed.\ntype MultipleAttrs = int @assert({{ this > 0 }}) @check(gt_ten, {{ this > 10 }})\n\nfunction MergeAliasAttributes(money: int) -> MergeAttrs {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given integer in the specified format:\n\n {{ money }}\n\n {{ ctx.output_format }}\n \"#\n}\n\nfunction ReturnAliasWithMergedAttributes(money: Amount) -> Amount {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given integer without additional context:\n\n {{ money }}\n\n {{ ctx.output_format }}\n \"#\n}\n\nfunction AliasWithMultipleAttrs(money: MultipleAttrs) -> MultipleAttrs {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given integer without additional context:\n\n {{ money }}\n\n {{ ctx.output_format }}\n \"#\n}\n\ntype RecursiveMapAlias = map<string, RecursiveMapAlias>\n\nfunction SimpleRecursiveMapAlias(input: RecursiveMapAlias) -> RecursiveMapAlias {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given value:\n\n {{ input }}\n\n {{ ctx.output_format }}\n \"#\n}\n\ntype RecursiveListAlias = RecursiveListAlias[]\n\nfunction SimpleRecursiveListAlias(input: RecursiveListAlias) -> RecursiveListAlias {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given JSON array:\n\n {{ input }}\n\n {{ ctx.output_format }}\n \"#\n}\n\ntype RecAliasOne = RecAliasTwo\ntype RecAliasTwo = RecAliasThree\ntype RecAliasThree = RecAliasOne[]\n\nfunction RecursiveAliasCycle(input: RecAliasOne) -> RecAliasOne {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given JSON array:\n\n {{ input }}\n\n {{ ctx.output_format }}\n \"#\n}\n\ntype JsonValue = int | string | bool | float | JsonObject | JsonArray\ntype JsonObject = map<string, JsonValue>\ntype JsonArray = JsonValue[]\n\nfunction JsonTypeAliasCycle(input: JsonValue) -> JsonValue {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given input back:\n\n {{ input }}\n\n {{ ctx.output_format }}\n \"#\n}",
"test-files/functions/output/type-aliases.baml": "type Primitive = int | string | bool | float\n\ntype List = string[]\n\ntype Graph = map<string, string[]>\n\ntype Combination = Primitive | List | Graph\n\nfunction PrimitiveAlias(p: Primitive) -> Primitive {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given value back: {{ p }}\n \"#\n}\n\nfunction MapAlias(m: Graph) -> Graph {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given Graph back:\n\n {{ m }}\n\n {{ ctx.output_format }}\n \"#\n}\n\nfunction NestedAlias(c: Combination) -> Combination {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given value back:\n\n {{ c }}\n\n {{ ctx.output_format }}\n \"#\n}\n\n// Test attribute merging.\ntype Currency = int @check(gt_ten, {{ this > 10 }})\ntype Amount = Currency @assert({{ this > 0 }})\n\nclass MergeAttrs {\n amount Amount @description(\"In USD\")\n}\n\n// This should be allowed.\ntype MultipleAttrs = int @assert({{ this > 0 }}) @check(gt_ten, {{ this > 10 }})\n\nfunction MergeAliasAttributes(money: int) -> MergeAttrs {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given integer in the specified format:\n\n {{ money }}\n\n {{ ctx.output_format }}\n \"#\n}\n\nfunction ReturnAliasWithMergedAttributes(money: Amount) -> Amount {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given integer without additional context:\n\n {{ money }}\n\n {{ ctx.output_format }}\n \"#\n}\n\nfunction AliasWithMultipleAttrs(money: MultipleAttrs) -> MultipleAttrs {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given integer without additional context:\n\n {{ money }}\n\n {{ ctx.output_format }}\n \"#\n}\n\ntype RecursiveMapAlias = map<string, RecursiveMapAlias>\n\nfunction SimpleRecursiveMapAlias(input: RecursiveMapAlias) -> RecursiveMapAlias {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given value:\n\n {{ input }}\n\n {{ ctx.output_format }}\n \"#\n}\n\ntype RecursiveListAlias = RecursiveListAlias[]\n\nfunction SimpleRecursiveListAlias(input: RecursiveListAlias) -> RecursiveListAlias {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given JSON array:\n\n {{ input }}\n\n {{ ctx.output_format }}\n \"#\n}\n\ntype RecAliasOne = RecAliasTwo\ntype RecAliasTwo = RecAliasThree\ntype RecAliasThree = RecAliasOne[]\n\nfunction RecursiveAliasCycle(input: RecAliasOne) -> RecAliasOne {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given JSON array:\n\n {{ input }}\n\n {{ ctx.output_format }}\n \"#\n}\n\ntype JsonValue = int | string | bool | float | JsonObject | JsonArray\ntype JsonObject = map<string, JsonValue>\ntype JsonArray = JsonValue[]\n\nfunction JsonTypeAliasCycle(input: JsonValue) -> JsonValue {\n client \"openai/gpt-4o\"\n prompt r#\"\n Return the given input back:\n\n {{ input }}\n\n {{ ctx.output_format }}\n \"#\n}\n",
"test-files/functions/output/unions.baml": "class UnionTest_ReturnType {\n prop1 string | bool\n prop2 (float | bool)[]\n prop3 (bool[] | int[])\n}\n\nfunction UnionTest_Function(input: string | bool) -> UnionTest_ReturnType {\n client GPT35\n prompt #\"\n Return a JSON blob with this schema: \n {{ctx.output_format}}\n\n JSON:\n \"#\n}\n\ntest UnionTest_Function {\n functions [UnionTest_Function]\n args {\n input \"example input\"\n }\n}\n",
"test-files/functions/prompts/no-chat-messages.baml": "\n\nfunction PromptTestClaude(input: string) -> string {\n client Sonnet\n prompt #\"\n Tell me a haiku about {{ input }}\n \"#\n}\n\n\nfunction PromptTestStreaming(input: string) -> string {\n client GPT35\n prompt #\"\n Tell me a short story about {{ input }}\n \"#\n}\n\ntest TestName {\n functions [PromptTestStreaming]\n args {\n input #\"\n hello world\n \"#\n }\n}\n",
"test-files/functions/prompts/with-chat-messages.baml": "\nfunction PromptTestOpenAIChat(input: string) -> string {\n client GPT35\n prompt #\"\n {{ _.role(\"system\") }}\n You are an assistant that always responds in a very excited way with emojis and also outputs this word 4 times after giving a response: {{ input }}\n \n {{ _.role(\"user\") }}\n Tell me a haiku about {{ input }}\n \"#\n}\n\nfunction PromptTestOpenAIChatNoSystem(input: string) -> string {\n client GPT35\n prompt #\"\n You are an assistant that always responds in a very excited way with emojis and also outputs this word 4 times after giving a response: {{ input }}\n \n {{ _.role(\"user\") }}\n Tell me a haiku about {{ input }}\n \"#\n}\n\nfunction PromptTestClaudeChat(input: string) -> string {\n client Claude\n prompt #\"\n {{ _.role(\"system\") }}\n You are an assistant that always responds in a very excited way with emojis and also outputs this word 4 times after giving a response: {{ input }}\n \n {{ _.role(\"user\") }}\n Tell me a haiku about {{ input }}\n \"#\n}\n\nfunction PromptTestClaudeChatNoSystem(input: string) -> string {\n client Claude\n prompt #\"\n You are an assistant that always responds in a very excited way with emojis and also outputs this word 4 times after giving a response: {{ input }}\n \n {{ _.role(\"user\") }}\n Tell me a haiku about {{ input }}\n \"#\n}\n\ntest TestSystemAndNonSystemChat1 {\n functions [PromptTestClaude, PromptTestOpenAI, PromptTestOpenAIChat, PromptTestOpenAIChatNoSystem, PromptTestClaudeChat, PromptTestClaudeChatNoSystem]\n args {\n input \"cats\"\n }\n}\n\ntest TestSystemAndNonSystemChat2 {\n functions [PromptTestClaude, PromptTestOpenAI, PromptTestOpenAIChat, PromptTestOpenAIChatNoSystem, PromptTestClaudeChat, PromptTestClaudeChatNoSystem]\n args {\n input \"lion\"\n }\n}",
Expand Down
Loading

0 comments on commit 39141cb

Please sign in to comment.