WIP

burr/core/action.py

@@ -70,17 +70,40 @@ def run(self, state: State, **run_kwargs) -> dict:
         pass
 
     @property
-    def inputs(self) -> Union[list[str], tuple[list[str], list[str]]]:
+    def inputs(
+        self,
+    ) -> Tuple[Union[list[str], Dict[str, Type[Type]]], Union[list[str], Dict[str, Type[Type]]]]:
         """Represents inputs that are used for this to run.
         These correspond to the ``**run_kwargs`` in `run` above.
 
         Note that this has two possible return values:
-        1. A list of strings -- these are the keys that are required to run the function
-        2. A tuple of two lists of strings -- the first list is the required keys, the second is the optional keys
+        1. A list of strings/dict of string -> ypes -- these are the keys that are required to run the function
+        2. A tuple of two lists of strings/dict of strings -> types -- the first list is the required keys, the second is the optional keys
 
         :return: Either a list of strings (required inputs) or a tuple of two lists of strings (required and optional inputs)
         """
-        return []
+        return [], []
+
+    @property
+    def input_schema(self) -> Tuple[dict[str, Type[Type]], dict[str, Type[Type]]]:
+        """Returns the input schema for the function.
+        The input schema is a type that can be used to validate the input to the function.
+
+        Note that this is separate from inputs() for backwards compatibility --
+        inputs() can return a schema *or* a tuple of required and optional inputs.
+
+        :return: Tuple of required inputs and optional inputs with attached types
+        """
+        inputs = self.inputs
+        if len(inputs) == 1:
+            inputs = (inputs[0], {})
+        out = []
+        for input_spec in inputs:
+            if isinstance(input_spec, list):
+                out.append({key: Any for key in input_spec})
+            else:
+                out.append(input_spec)
+        return tuple(out)
 
     @property
     def optional_and_required_inputs(self) -> tuple[set[str], set[str]]:
@@ -213,11 +236,6 @@ def get_source(self) -> str:
         to display a different source"""
         return inspect.getsource(self.__class__)
 
-    def input_schema(self) -> Any:
-        """Returns the input schema for the action.
-        The input schema is a type that can be used to validate the input to the action"""
-        return None
-
     def __repr__(self):
         read_repr = ", ".join(self.reads) if self.reads else "{}"
         write_repr = ", ".join(self.writes) if self.writes else "{}"
@@ -534,7 +552,9 @@ def is_async(self) -> bool:
 
 # the following exist to share implementation between FunctionBasedStreamingAction and FunctionBasedAction
 # TODO -- think through the class hierarchy to simplify, for now this is OK
-def derive_inputs_from_fn(bound_params: dict, fn: Callable) -> tuple[list[str], list[str]]:
+def derive_inputs_from_fn(
+    bound_params: dict, fn: Callable
+) -> tuple[dict[str, Type[Type]], dict[str, Type[Type]]]:
     """Derives inputs from the function, given the bound parameters. This assumes that the function
     has inputs named `state`, as well as any number of other kwarg-boundable parameters.
 
@@ -543,20 +563,30 @@ def derive_inputs_from_fn(bound_params: dict, fn: Callable) -> tuple[list[str],
     :return: Required and optional inputs
     """
     sig = inspect.signature(fn)
-    required_inputs, optional_inputs = [], []
+    required_inputs, optional_inputs = {}, {}
     for param_name, param in sig.parameters.items():
         if param_name != "state" and param_name not in bound_params:
             if param.default is inspect.Parameter.empty:
                 # has no default means its required
-                required_inputs.append(param_name)
+                required_inputs[param_name] = (
+                    param.annotation if param.annotation != inspect.Parameter.empty else Any
+                )
             else:
                 # has a default means its optional
-                optional_inputs.append(param_name)
+                optional_inputs[param_name] = (
+                    param.annotation if param.annotation != inspect.Parameter.empty else Any
+                )
     return required_inputs, optional_inputs
 
 
 FunctionBasedActionType = Union["FunctionBasedAction", "FunctionBasedStreamingAction"]
 
+InputSpec = [
+    Tuple[Union[list[str], Dict[str, Type[Type]], Union[list[str], Dict[str, Type[Type]]]]]
+]
+
+OptionalInputType = Optional[InputSpec]
+
 
 class FunctionBasedAction(SingleStepAction):
     ACTION_FUNCTION = "action_function"
@@ -567,7 +597,9 @@ def __init__(
         reads: List[str],
         writes: List[str],
         bound_params: Optional[dict] = None,
-        input_spec: Optional[tuple[list[str], list[str]]] = None,
+        input_spec: Optional[
+            Tuple[Union[list[str], Dict[str, Type[Type]]], Union[list[str], Dict[str, Type[Type]]]]
+        ] = None,
         originating_fn: Optional[Callable] = None,
         schema: ActionSchema = DEFAULT_SCHEMA,
     ):
@@ -609,7 +641,9 @@ def writes(self) -> list[str]:
         return self._writes
 
     @property
-    def inputs(self) -> tuple[list[str], list[str]]:
+    def inputs(
+        self,
+    ) -> Tuple[Union[list[str], Dict[str, Type[Type]]], Union[list[str], Dict[str, Type[Type]]]]:
         return self._inputs
 
     @property
@@ -1044,7 +1078,9 @@ def __init__(
         reads: List[str],
         writes: List[str],
         bound_params: Optional[dict] = None,
-        input_spec: Optional[tuple[list[str], list[str]]] = None,
+        input_spec: Optional[
+            Tuple[Union[list[str], Dict[str, Type[Type]]], Union[list[str], Dict[str, Type[Type]]]]
+        ] = None,
         originating_fn: Optional[Callable] = None,
         schema: ActionSchema = DEFAULT_SCHEMA,
     ):
@@ -1117,7 +1153,9 @@ def with_params(self, **kwargs: Any) -> "FunctionBasedStreamingAction":
         )
 
     @property
-    def inputs(self) -> tuple[list[str], list[str]]:
+    def inputs(
+        self,
+    ) -> Tuple[Union[list[str], Dict[str, Type[Type]]], Union[list[str], Dict[str, Type[Type]]]]:
         return self._inputs
 
     @property

burr/tracking/server/run.py

@@ -32,6 +32,7 @@
 
     # dynamic importing due to the dashes (which make reading the examples on github easier)
     email_assistant = importlib.import_module("burr.examples.email-assistant.server")
+    email_assistant_typed = importlib.import_module("burr.examples.fastapi.server")
     chatbot = importlib.import_module("burr.examples.multi-modal-chatbot.server")
     streaming_chatbot = importlib.import_module("burr.examples.streaming-fastapi.server")
 
@@ -249,6 +250,7 @@ async def version() -> dict:
 # Examples -- todo -- put them behind `if` statements
 app.include_router(chatbot.router, prefix="/api/v0/chatbot")
 app.include_router(email_assistant.router, prefix="/api/v0/email_assistant")
+app.include_router(email_assistant_typed.router, prefix="/api/v0/email_assistant_typed")
 app.include_router(streaming_chatbot.router, prefix="/api/v0/streaming_chatbot")
 
 if SERVE_STATIC:

telemetry/ui/src/api/index.ts

@@ -12,17 +12,22 @@ export type { ApplicationLogs } from './models/ApplicationLogs';
 export type { ApplicationModel } from './models/ApplicationModel';
 export type { ApplicationPage } from './models/ApplicationPage';
 export type { ApplicationSummary } from './models/ApplicationSummary';
+export type { AppResponse } from './models/AppResponse';
 export type { AttributeModel } from './models/AttributeModel';
 export type { BackendSpec } from './models/BackendSpec';
 export type { BeginEntryModel } from './models/BeginEntryModel';
 export type { BeginSpanModel } from './models/BeginSpanModel';
+export { burr__examples__email_assistant__server__EmailAssistantState } from './models/burr__examples__email_assistant__server__EmailAssistantState';
 export { ChatItem } from './models/ChatItem';
 export { ChildApplicationModel } from './models/ChildApplicationModel';
+export type { ClarificationAnswers } from './models/ClarificationAnswers';
+export type { ClarificationQuestions } from './models/ClarificationQuestions';
 export type { DraftInit } from './models/DraftInit';
-export { EmailAssistantState } from './models/EmailAssistantState';
+export type { Email } from './models/Email';
 export type { EndEntryModel } from './models/EndEntryModel';
 export type { EndSpanModel } from './models/EndSpanModel';
 export type { EndStreamModel } from './models/EndStreamModel';
+export type { examples__fastapi__application__EmailAssistantState } from './models/examples__fastapi__application__EmailAssistantState';
 export type { Feedback } from './models/Feedback';
 export type { FirstItemStreamModel } from './models/FirstItemStreamModel';
 export type { HTTPValidationError } from './models/HTTPValidationError';

telemetry/ui/src/api/models/AppResponse.ts

@@ -0,0 +1,10 @@
+/* generated using openapi-typescript-codegen -- do no edit */
+/* istanbul ignore file */
+/* tslint:disable */
+/* eslint-disable */
+import type { examples__fastapi__application__EmailAssistantState } from './examples__fastapi__application__EmailAssistantState';
+export type AppResponse = {
+  app_id: string;
+  next_step: 'process_input' | 'clarify_instructions' | 'process_feedback' | null;
+  state: examples__fastapi__application__EmailAssistantState;
+};

telemetry/ui/src/api/models/ClarificationAnswers.ts

@@ -0,0 +1,7 @@
+/* generated using openapi-typescript-codegen -- do no edit */
+/* istanbul ignore file */
+/* tslint:disable */
+/* eslint-disable */
+export type ClarificationAnswers = {
+  answers: Array<string>;
+};

telemetry/ui/src/api/models/ClarificationQuestions.ts

@@ -0,0 +1,7 @@
+/* generated using openapi-typescript-codegen -- do no edit */
+/* istanbul ignore file */
+/* tslint:disable */
+/* eslint-disable */
+export type ClarificationQuestions = {
+  question: Array<string>;
+};

telemetry/ui/src/api/models/Email.ts

@@ -0,0 +1,8 @@
+/* generated using openapi-typescript-codegen -- do no edit */
+/* istanbul ignore file */
+/* tslint:disable */
+/* eslint-disable */
+export type Email = {
+  subject: string;
+  contents: string;
+};

telemetry/ui/src/api/models/EmailAssistantState.ts → telemetry/ui/src/api/models/burr__examples__email_assistant__server__EmailAssistantState.ts

@@ -2,7 +2,7 @@
 /* istanbul ignore file */
 /* tslint:disable */
 /* eslint-disable */
-export type EmailAssistantState = {
+export type burr__examples__email_assistant__server__EmailAssistantState = {
   app_id: string;
   email_to_respond: string | null;
   response_instructions: string | null;
@@ -11,9 +11,9 @@ export type EmailAssistantState = {
   drafts: Array<string>;
   feedback_history: Array<string>;
   final_draft: string | null;
-  next_step: EmailAssistantState.next_step;
+  next_step: burr__examples__email_assistant__server__EmailAssistantState.next_step;
 };
-export namespace EmailAssistantState {
+export namespace burr__examples__email_assistant__server__EmailAssistantState {
   export enum next_step {
     PROCESS_INPUT = 'process_input',
     CLARIFY_INSTRUCTIONS = 'clarify_instructions',

telemetry/ui/src/api/models/examples__fastapi__application__EmailAssistantState.ts

@@ -0,0 +1,18 @@
+/* generated using openapi-typescript-codegen -- do no edit */
+/* istanbul ignore file */
+/* tslint:disable */
+/* eslint-disable */
+import type { ClarificationAnswers } from './ClarificationAnswers';
+import type { ClarificationQuestions } from './ClarificationQuestions';
+import type { Email } from './Email';
+export type examples__fastapi__application__EmailAssistantState = {
+  email_to_respond?: string | null;
+  response_instructions?: string | null;
+  questions?: ClarificationQuestions | null;
+  answers?: ClarificationAnswers | null;
+  draft_history?: Array<Email>;
+  current_draft?: Email | null;
+  feedback_history?: Array<string>;
+  feedback?: string | null;
+  final_draft?: string | null;
+};