From 46ed8facb35c1b17978bbb21bdf20246c19d3bf0 Mon Sep 17 00:00:00 2001 From: Hamada Salhab Date: Thu, 3 Oct 2024 00:47:50 +0300 Subject: [PATCH 1/3] feat(agents-api): Add static checking for Jinja templates & Python expressions in task creation | Add validation for subworkflows (#570) Closes #535 > [!IMPORTANT] > Add static validation for Python expressions and Jinja templates in task creation and validate subworkflows in `openapi_model.py` and `steps.tsp`. > > - **Validation**: > - Add `validate_python_expression()` and `validate_jinja_template()` in `openapi_model.py`. > - Integrate validation into `EvaluateStep`, `ToolCallStep`, `PromptStep`, `SetStep`, `LogStep`, `ReturnStep`, `YieldStep`, `IfElseWorkflowStep`, and `MapReduceStep` in `openapi_model.py`. > - **Models**: > - Update `CreateTaskRequest` in `openapi_model.py` to validate subworkflows using `WorkflowType`. > - Add `YieldStep` to `MappableWorkflowStep` and `NonConditionalWorkflowStep` in `steps.tsp`. > - **Misc**: > - Reorder `YieldStep` in `Tasks.py` to maintain consistency. > > This description was created by [Ellipsis](https://www.ellipsis.dev?ref=julep-ai%2Fjulep&utm_source=github&utm_medium=referral) for 9952ad5700812126c0aa7f1bfa26467e88b60aab. It will automatically update as commits are pushed. --------- Co-authored-by: Diwank Singh Tomer --- agents-api/agents_api/autogen/Tasks.py | 26 ++- .../agents_api/autogen/openapi_model.py | 216 +++++++++++++++++- typespec/tasks/steps.tsp | 3 +- 3 files changed, 233 insertions(+), 12 deletions(-) diff --git a/agents-api/agents_api/autogen/Tasks.py b/agents-api/agents_api/autogen/Tasks.py index 48dba4ad7..9dd531c47 100644 --- a/agents-api/agents_api/autogen/Tasks.py +++ b/agents-api/agents_api/autogen/Tasks.py @@ -35,10 +35,10 @@ class CaseThen(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep ) """ @@ -63,10 +63,10 @@ class CaseThenUpdateItem(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep ) """ @@ -130,10 +130,10 @@ class CreateTaskRequest(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep | IfElseWorkflowStep | SwitchStep @@ -227,6 +227,7 @@ class ForeachDo(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep ) """ The steps to run for each iteration @@ -251,6 +252,7 @@ class ForeachDoUpdateItem(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep ) """ The steps to run for each iteration @@ -324,10 +326,10 @@ class IfElseWorkflowStep(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep ) """ @@ -342,10 +344,10 @@ class IfElseWorkflowStep(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep | None, Field(None, alias="else"), @@ -376,10 +378,10 @@ class IfElseWorkflowStepUpdateItem(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep ) """ @@ -394,10 +396,10 @@ class IfElseWorkflowStepUpdateItem(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep | None, Field(None, alias="else"), @@ -462,6 +464,7 @@ class Main(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep ) """ The steps to run for each iteration @@ -503,6 +506,7 @@ class MainModel(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep ) """ The steps to run for each iteration @@ -543,6 +547,7 @@ class ParallelStep(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep ], Field(max_length=100), ] @@ -569,6 +574,7 @@ class ParallelStepUpdateItem(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep ], Field(max_length=100), ] @@ -596,10 +602,10 @@ class PatchTaskRequest(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep | IfElseWorkflowStepUpdateItem | SwitchStepUpdateItem @@ -874,10 +880,10 @@ class Task(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep | IfElseWorkflowStep | SwitchStep @@ -1009,10 +1015,10 @@ class UpdateTaskRequest(BaseModel): | LogStep | EmbedStep | SearchStep + | YieldStep | ReturnStep | SleepStep | ErrorWorkflowStep - | YieldStep | WaitForInputStep | IfElseWorkflowStep | SwitchStep diff --git a/agents-api/agents_api/autogen/openapi_model.py b/agents-api/agents_api/autogen/openapi_model.py index 5c5a8c86f..48811eb20 100644 --- a/agents-api/agents_api/autogen/openapi_model.py +++ b/agents-api/agents_api/autogen/openapi_model.py @@ -1,10 +1,12 @@ # ruff: noqa: F401, F403, F405 +import ast from typing import Annotated, Any, Generic, Literal, Self, Type, TypeVar, get_args from uuid import UUID +import jinja2 from litellm.utils import _select_tokenizer as select_tokenizer from litellm.utils import token_counter -from pydantic import AwareDatetime, Field +from pydantic import AwareDatetime, Field, field_validator, model_validator, validator from ..common.utils.datetime import utcnow from .Agents import * @@ -152,6 +154,179 @@ def from_model_input( ) +# Patch Task Workflow Steps +# -------------------------------------- + + +def validate_python_expression(expr: str) -> tuple[bool, str]: + try: + ast.parse(expr) + return True, "" + except SyntaxError as e: + return False, f"SyntaxError in '{expr}': {str(e)}" + + +def validate_jinja_template(template: str) -> tuple[bool, str]: + env = jinja2.Environment() + try: + parsed_template = env.parse(template) + for node in parsed_template.body: + if isinstance(node, jinja2.nodes.Output): + for child in node.nodes: + if isinstance(child, jinja2.nodes.Name): + # Check if the variable is a valid Python expression + is_valid, error = validate_python_expression(child.name) + if not is_valid: + return ( + False, + f"Invalid Python expression in Jinja template '{template}': {error}", + ) + return True, "" + except jinja2.exceptions.TemplateSyntaxError as e: + return False, f"TemplateSyntaxError in '{template}': {str(e)}" + + +_EvaluateStep = EvaluateStep + + +class EvaluateStep(_EvaluateStep): + @field_validator("evaluate") + def validate_evaluate_expressions(cls, v): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError(f"Invalid Python expression in key '{key}': {error}") + return v + + +_ToolCallStep = ToolCallStep + + +class ToolCallStep(_ToolCallStep): + @field_validator("arguments") + def validate_arguments(cls, v): + if isinstance(v, dict): + for key, expr in v.items(): + if isinstance(expr, str): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError( + f"Invalid Python expression in arguments key '{key}': {error}" + ) + return v + + +_PromptStep = PromptStep + + +class PromptStep(_PromptStep): + @field_validator("prompt") + def validate_prompt(cls, v): + if isinstance(v, str): + is_valid, error = validate_jinja_template(v) + if not is_valid: + raise ValueError(f"Invalid Jinja template in prompt: {error}") + elif isinstance(v, list): + for item in v: + if "content" in item: + is_valid, error = validate_jinja_template(item["content"]) + if not is_valid: + raise ValueError( + f"Invalid Jinja template in prompt content: {error}" + ) + return v + + +_SetStep = SetStep + + +class SetStep(_SetStep): + @field_validator("set") + def validate_set_expressions(cls, v): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError( + f"Invalid Python expression in set key '{key}': {error}" + ) + return v + + +_LogStep = LogStep + + +class LogStep(_LogStep): + @field_validator("log") + def validate_log_template(cls, v): + is_valid, error = validate_jinja_template(v) + if not is_valid: + raise ValueError(f"Invalid Jinja template in log: {error}") + return v + + +_ReturnStep = ReturnStep + + +class ReturnStep(_ReturnStep): + @field_validator("return_") + def validate_return_expressions(cls, v): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError( + f"Invalid Python expression in return key '{key}': {error}" + ) + return v + + +_YieldStep = YieldStep + + +class YieldStep(_YieldStep): + @field_validator("arguments") + def validate_yield_arguments(cls, v): + if isinstance(v, dict): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError( + f"Invalid Python expression in yield arguments key '{key}': {error}" + ) + return v + + +_IfElseWorkflowStep = IfElseWorkflowStep + + +class IfElseWorkflowStep(_IfElseWorkflowStep): + @field_validator("if_") + def validate_if_expression(cls, v): + is_valid, error = validate_python_expression(v) + if not is_valid: + raise ValueError(f"Invalid Python expression in if condition: {error}") + return v + + +_MapReduceStep = MapReduceStep + + +class MapReduceStep(_MapReduceStep): + @field_validator("over") + def validate_over_expression(cls, v): + is_valid, error = validate_python_expression(v) + if not is_valid: + raise ValueError(f"Invalid Python expression in over: {error}") + return v + + @field_validator("reduce") + def validate_reduce_expression(cls, v): + if v is not None: + is_valid, error = validate_python_expression(v) + if not is_valid: + raise ValueError(f"Invalid Python expression in reduce: {error}") + return v + + # Workflow related models # ----------------------- @@ -228,6 +403,29 @@ class Task(_Task): # Patch some models to allow extra fields # -------------------------------------- +WorkflowType = RootModel[ + list[ + EvaluateStep + | ToolCallStep + | PromptStep + | GetStep + | SetStep + | LogStep + | EmbedStep + | SearchStep + | ReturnStep + | SleepStep + | ErrorWorkflowStep + | YieldStep + | WaitForInputStep + | IfElseWorkflowStep + | SwitchStep + | ForeachStep + | ParallelStep + | MapReduceStep + ] +] + _CreateTaskRequest = CreateTaskRequest @@ -240,6 +438,22 @@ class CreateTaskRequest(_CreateTaskRequest): } ) + @model_validator(mode="after") + def validate_subworkflows(self) -> Self: + subworkflows = { + k: v + for k, v in self.model_dump().items() + if k not in _CreateTaskRequest.model_fields + } + + for workflow_name, workflow_definition in subworkflows.items(): + try: + WorkflowType.model_validate(workflow_definition) + setattr(self, workflow_name, WorkflowType(workflow_definition)) + except Exception as e: + raise ValueError(f"Invalid subworkflow '{workflow_name}': {str(e)}") + return self + CreateOrUpdateTaskRequest = CreateTaskRequest diff --git a/typespec/tasks/steps.tsp b/typespec/tasks/steps.tsp index 3495def1b..2267ae320 100644 --- a/typespec/tasks/steps.tsp +++ b/typespec/tasks/steps.tsp @@ -49,7 +49,8 @@ alias MappableWorkflowStep = | SetStep | LogStep | EmbedStep - | SearchStep; + | SearchStep + | YieldStep; alias NonConditionalWorkflowStep = | MappableWorkflowStep From 866a4eb6c1f2ee85954f6efed6ba0a746ea67058 Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Wed, 2 Oct 2024 18:11:00 -0400 Subject: [PATCH 2/3] doc: Add code reading instructions to CONTRIBUTING.md (#556) Signed-off-by: Diwank Singh Tomer ---- > [!IMPORTANT] > Adds detailed code reading and contributing instructions to `CONTRIBUTING.md`, covering project architecture, setup, and contribution guidelines. > > - **Documentation**: > - Adds detailed code reading instructions to `CONTRIBUTING.md`. > - Includes sections on project overview, system architecture, API specifications, core API implementation, database and storage, workflow management, testing, and additional services. > - Provides a step-by-step guide for setting up the development environment and contributing to the project. > > This description was created by [Ellipsis](https://www.ellipsis.dev?ref=julep-ai%2Fjulep&utm_source=github&utm_medium=referral) for 4bb3e8ce3e613a87541d88734368762251a012e4. It will automatically update as commits are pushed. --------- Signed-off-by: Diwank Singh Tomer --- CONTRIBUTING.md | 114 +++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 113 insertions(+), 1 deletion(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e60c31973..a7cda8b06 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -22,6 +22,115 @@ To contribute code changes: Please ensure your code follows the existing style and passes all tests. +## Project Overview and Architecture + +### Key Components + +1. **agents-api**: The core API service for Julep. +2. **typespec**: API specifications and contracts. +3. **integrations-service**: Handles external integrations. +4. **embedding-service**: Manages text embeddings. +5. **memory-store**: Handles persistent storage. +6. **llm-proxy**: Proxy for language model interactions. +7. **scheduler**: Manages task scheduling. +8. **gateway**: API gateway and routing. +9. **monitoring**: System monitoring and metrics. + +### Technology Stack + +- **FastAPI**: Web framework for building APIs +- **TypeSpec**: API specification language +- **Cozo**: Database system +- **Temporal**: Workflow engine +- **Docker**: Containerization + +### Relationships Between Components + +The `agents-api` serves as the central component, interacting with most other services: +- It uses `typespec` definitions for API contracts. +- Communicates with `integrations-service` for external tool interactions. +- Utilizes `embedding-service` for text processing. +- Stores data in `memory-store`. +- Interacts with language models through `llm-proxy`. +- Uses `scheduler` for task management. +- All API requests pass through the `gateway`. +- `monitoring` observes the entire system. + +## Understanding the Codebase + +To get a comprehensive understanding of Julep, we recommend exploring the codebase in the following order: + +1. **Project Overview** + - Read `README.md` in the root directory + - Explore `docs/` for detailed documentation + +2. **System Architecture** + - Examine `docker-compose.yml` in the root directory + - Review `deploy/` directory for different deployment configurations + +3. **API Specifications** + - Learn about TypeSpec: https://typespec.io/docs/ + - Explore `typespec/` directory: + - Start with `common/` folder + - Review `main.tsp` + - Examine each module sequentially + +4. **Core API Implementation** + - Learn about FastAPI: https://fastapi.tiangolo.com/ + - Explore `agents-api/` directory: + - Review `README.md` for an overview + - Examine `routers/` for API endpoints + - Look into `models/` for data models + +5. **Database and Storage** + - Learn about Cozo: https://docs.cozodb.org/en/latest/tutorial.html + - Review `agents-api/README.md` for database schema + - Explore `agents-api/models/` for database queries + +6. **Workflow Management** + - Learn about Temporal: https://docs.temporal.io/develop/python + - Explore `agents-api/activities/` for individual workflow steps + - Review `agents-api/workflows/task_execution/` for main workflow logic + +7. **Testing** + - Examine `agents-api/tests/` for test cases + +8. **Additional Services** + - Explore other service directories (`integrations-service/`, `embedding-service/`, etc.) to understand their specific roles and implementations + +## Contributing Guidelines + +1. **Set Up Development Environment** + - Clone the repository + - Install Docker and Docker Compose + - Set up necessary API keys and environment variables + +2. **Choose an Area to Contribute** + - Check the issue tracker for open issues + - Look for "good first issue" labels for newcomers + +3. **Make Changes** + - Create a new branch for your changes + - Write clean, well-documented code + - Ensure your changes adhere to the project's coding standards + +4. **Test Your Changes** + - Run existing tests + - Add new tests for new functionality + - Ensure all tests pass before submitting your changes + +5. **Submit a Pull Request** + - Provide a clear description of your changes + - Reference any related issues + - Be prepared to respond to feedback and make adjustments + +6. **Code Review** + - Address any comments or suggestions from reviewers + - Make necessary changes and push updates to your branch + +7. **Merge** + - Once approved, your changes will be merged into the main branch + ### Documentation Improvements Improvements to documentation are always appreciated! If you see areas that could be clarified or expanded, feel free to make the changes and submit a pull request. @@ -98,4 +207,7 @@ This command generates a JWT token that will be valid for 10 days. ##### Troubleshooting - Ensure that all required Docker images are available. - Check for missing environment variables in the `.env` file. -- Use the `docker compose logs` command to view detailed logs for debugging. \ No newline at end of file +- Use the `docker compose logs` command to view detailed logs for debugging. + + +Remember, contributions aren't limited to code. Documentation improvements, bug reports, and feature suggestions are also valuable contributions to the project. From 2d62857458a254c48252993693cf04973d1e58d5 Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Wed, 2 Oct 2024 21:58:49 -0400 Subject: [PATCH 3/3] fix(agents-api): Switch to monkeypatching because everything is shit (#573) Signed-off-by: Diwank Singh Tomer ---- > [!IMPORTANT] > Refactor validation logic using monkeypatching and add new models and custom types in `openapi_model.py`. > > - **Validation Refactor**: > - Switch from subclassing to monkeypatching for validation functions in `EvaluateStep`, `ToolCallStep`, `PromptStep`, `SetStep`, `LogStep`, `ReturnStep`, `YieldStep`, `IfElseWorkflowStep`, and `MapReduceStep`. > - Add validators for Python expressions and Jinja templates. > - **Model Changes**: > - Add `CreateTransitionRequest` and `CreateEntryRequest` classes. > - Define custom types `ChatMLContent`, `ChatMLRole`, `ChatMLSource`, `ExecutionStatus`, and `TransitionType`. > - **Misc**: > - Allow extra fields in `CreateTaskRequest`, `PatchTaskRequest`, and `UpdateTaskRequest` using `ConfigDict`. > > This description was created by [Ellipsis](https://www.ellipsis.dev?ref=julep-ai%2Fjulep&utm_source=github&utm_medium=referral) for 8c03d93701a454327d6b458177e0626ae165d2a0. It will automatically update as commits are pushed. Signed-off-by: Diwank Singh Tomer --- .../agents_api/autogen/openapi_model.py | 339 +++++++++++------- 1 file changed, 214 insertions(+), 125 deletions(-) diff --git a/agents-api/agents_api/autogen/openapi_model.py b/agents-api/agents_api/autogen/openapi_model.py index 48811eb20..bff2221eb 100644 --- a/agents-api/agents_api/autogen/openapi_model.py +++ b/agents-api/agents_api/autogen/openapi_model.py @@ -6,7 +6,14 @@ import jinja2 from litellm.utils import _select_tokenizer as select_tokenizer from litellm.utils import token_counter -from pydantic import AwareDatetime, Field, field_validator, model_validator, validator +from pydantic import ( + AwareDatetime, + Field, + computed_field, + field_validator, + model_validator, + validator, +) from ..common.utils.datetime import utcnow from .Agents import * @@ -155,7 +162,7 @@ def from_model_input( # Patch Task Workflow Steps -# -------------------------------------- +# ------------------------- def validate_python_expression(expr: str) -> tuple[bool, str]: @@ -186,145 +193,255 @@ def validate_jinja_template(template: str) -> tuple[bool, str]: return False, f"TemplateSyntaxError in '{template}': {str(e)}" -_EvaluateStep = EvaluateStep +@field_validator("evaluate") +def validate_evaluate_expressions(cls, v): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError(f"Invalid Python expression in key '{key}': {error}") + return v + + +EvaluateStep.validate_evaluate_expressions = validate_evaluate_expressions -class EvaluateStep(_EvaluateStep): - @field_validator("evaluate") - def validate_evaluate_expressions(cls, v): +@field_validator("arguments") +def validate_arguments(cls, v): + if isinstance(v, dict): for key, expr in v.items(): - is_valid, error = validate_python_expression(expr) - if not is_valid: - raise ValueError(f"Invalid Python expression in key '{key}': {error}") - return v + if isinstance(expr, str): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError( + f"Invalid Python expression in arguments key '{key}': {error}" + ) + return v -_ToolCallStep = ToolCallStep +ToolCallStep.validate_arguments = validate_arguments -class ToolCallStep(_ToolCallStep): - @field_validator("arguments") - def validate_arguments(cls, v): - if isinstance(v, dict): - for key, expr in v.items(): - if isinstance(expr, str): - is_valid, error = validate_python_expression(expr) - if not is_valid: - raise ValueError( - f"Invalid Python expression in arguments key '{key}': {error}" - ) - return v +# Add the new validator function +@field_validator("prompt") +def validate_prompt(cls, v): + if isinstance(v, str): + is_valid, error = validate_jinja_template(v) + if not is_valid: + raise ValueError(f"Invalid Jinja template in prompt: {error}") + elif isinstance(v, list): + for item in v: + if "content" in item: + is_valid, error = validate_jinja_template(item["content"]) + if not is_valid: + raise ValueError( + f"Invalid Jinja template in prompt content: {error}" + ) + return v -_PromptStep = PromptStep +# Patch the original PromptStep class to add the new validator +PromptStep.validate_prompt = validate_prompt -class PromptStep(_PromptStep): - @field_validator("prompt") - def validate_prompt(cls, v): - if isinstance(v, str): - is_valid, error = validate_jinja_template(v) - if not is_valid: - raise ValueError(f"Invalid Jinja template in prompt: {error}") - elif isinstance(v, list): - for item in v: - if "content" in item: - is_valid, error = validate_jinja_template(item["content"]) - if not is_valid: - raise ValueError( - f"Invalid Jinja template in prompt content: {error}" - ) - return v +@field_validator("set") +def validate_set_expressions(cls, v): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError(f"Invalid Python expression in set key '{key}': {error}") + return v -_SetStep = SetStep +SetStep.validate_set_expressions = validate_set_expressions -class SetStep(_SetStep): - @field_validator("set") - def validate_set_expressions(cls, v): - for key, expr in v.items(): - is_valid, error = validate_python_expression(expr) - if not is_valid: - raise ValueError( - f"Invalid Python expression in set key '{key}': {error}" - ) - return v +@field_validator("log") +def validate_log_template(cls, v): + is_valid, error = validate_jinja_template(v) + if not is_valid: + raise ValueError(f"Invalid Jinja template in log: {error}") + return v -_LogStep = LogStep +LogStep.validate_log_template = validate_log_template -class LogStep(_LogStep): - @field_validator("log") - def validate_log_template(cls, v): - is_valid, error = validate_jinja_template(v) +@field_validator("return_") +def validate_return_expressions(cls, v): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) if not is_valid: - raise ValueError(f"Invalid Jinja template in log: {error}") - return v + raise ValueError( + f"Invalid Python expression in return key '{key}': {error}" + ) + return v -_ReturnStep = ReturnStep +ReturnStep.validate_return_expressions = validate_return_expressions -class ReturnStep(_ReturnStep): - @field_validator("return_") - def validate_return_expressions(cls, v): +@field_validator("arguments") +def validate_yield_arguments(cls, v): + if isinstance(v, dict): for key, expr in v.items(): is_valid, error = validate_python_expression(expr) if not is_valid: raise ValueError( - f"Invalid Python expression in return key '{key}': {error}" + f"Invalid Python expression in yield arguments key '{key}': {error}" ) - return v + return v -_YieldStep = YieldStep +YieldStep.validate_yield_arguments = validate_yield_arguments -class YieldStep(_YieldStep): - @field_validator("arguments") - def validate_yield_arguments(cls, v): - if isinstance(v, dict): - for key, expr in v.items(): - is_valid, error = validate_python_expression(expr) - if not is_valid: - raise ValueError( - f"Invalid Python expression in yield arguments key '{key}': {error}" - ) - return v +@field_validator("if_") +def validate_if_expression(cls, v): + is_valid, error = validate_python_expression(v) + if not is_valid: + raise ValueError(f"Invalid Python expression in if condition: {error}") + return v + +IfElseWorkflowStep.validate_if_expression = validate_if_expression -_IfElseWorkflowStep = IfElseWorkflowStep +@field_validator("over") +def validate_over_expression(cls, v): + is_valid, error = validate_python_expression(v) + if not is_valid: + raise ValueError(f"Invalid Python expression in over: {error}") + return v -class IfElseWorkflowStep(_IfElseWorkflowStep): - @field_validator("if_") - def validate_if_expression(cls, v): + +@field_validator("reduce") +def validate_reduce_expression(cls, v): + if v is not None: is_valid, error = validate_python_expression(v) if not is_valid: - raise ValueError(f"Invalid Python expression in if condition: {error}") - return v + raise ValueError(f"Invalid Python expression in reduce: {error}") + return v -_MapReduceStep = MapReduceStep +MapReduceStep.validate_over_expression = validate_over_expression +MapReduceStep.validate_reduce_expression = validate_reduce_expression -class MapReduceStep(_MapReduceStep): - @field_validator("over") - def validate_over_expression(cls, v): - is_valid, error = validate_python_expression(v) - if not is_valid: - raise ValueError(f"Invalid Python expression in over: {error}") - return v +# Patch workflow +# -------------- - @field_validator("reduce") - def validate_reduce_expression(cls, v): - if v is not None: - is_valid, error = validate_python_expression(v) - if not is_valid: - raise ValueError(f"Invalid Python expression in reduce: {error}") - return v +_CreateTaskRequest = CreateTaskRequest + +CreateTaskRequest.model_config = ConfigDict( + **{ + **_CreateTaskRequest.model_config, + "extra": "allow", + } +) + + +@model_validator(mode="after") +def validate_subworkflows(self): + subworkflows = { + k: v + for k, v in self.model_dump().items() + if k not in _CreateTaskRequest.model_fields + } + + for workflow_name, workflow_definition in subworkflows.items(): + try: + WorkflowType.model_validate(workflow_definition) + setattr(self, workflow_name, WorkflowType(workflow_definition)) + except Exception as e: + raise ValueError(f"Invalid subworkflow '{workflow_name}': {str(e)}") + return self + + +CreateTaskRequest.validate_subworkflows = validate_subworkflows + + +# Custom types (not generated correctly) +# -------------------------------------- + +ChatMLContent = ( + list[ChatMLTextContentPart | ChatMLImageContentPart] + | Tool + | ChosenToolCall + | str + | ToolResponse + | list[ + list[ChatMLTextContentPart | ChatMLImageContentPart] + | Tool + | ChosenToolCall + | str + | ToolResponse + ] +) + +# Extract ChatMLRole +ChatMLRole = BaseEntry.model_fields["role"].annotation + +# Extract ChatMLSource +ChatMLSource = BaseEntry.model_fields["source"].annotation + +# Extract ExecutionStatus +ExecutionStatus = Execution.model_fields["status"].annotation + +# Extract TransitionType +TransitionType = Transition.model_fields["type"].annotation + +# Assertions to ensure consistency (optional, but recommended for runtime checks) +assert ChatMLRole == BaseEntry.model_fields["role"].annotation +assert ChatMLSource == BaseEntry.model_fields["source"].annotation +assert ExecutionStatus == Execution.model_fields["status"].annotation +assert TransitionType == Transition.model_fields["type"].annotation + + +# Create models +# ------------- + + +class CreateTransitionRequest(Transition): + # The following fields are optional in this + + id: UUID | None = None + execution_id: UUID | None = None + created_at: AwareDatetime | None = None + updated_at: AwareDatetime | None = None + metadata: dict[str, Any] | None = None + task_token: str | None = None + + +class CreateEntryRequest(BaseEntry): + timestamp: Annotated[ + float, Field(ge=0.0, default_factory=lambda: utcnow().timestamp()) + ] + + @classmethod + def from_model_input( + cls: Type[Self], + model: str, + *, + role: ChatMLRole, + content: ChatMLContent, + name: str | None = None, + source: ChatMLSource, + **kwargs: dict, + ) -> Self: + tokenizer: dict = select_tokenizer(model=model) + token_count = token_counter( + model=model, messages=[{"role": role, "content": content, "name": name}] + ) + + return cls( + role=role, + content=content, + name=name, + source=source, + tokenizer=tokenizer["type"], + token_count=token_count, + **kwargs, + ) # Workflow related models @@ -427,34 +544,6 @@ class Task(_Task): ] -_CreateTaskRequest = CreateTaskRequest - - -class CreateTaskRequest(_CreateTaskRequest): - model_config = ConfigDict( - **{ - **_CreateTaskRequest.model_config, - "extra": "allow", - } - ) - - @model_validator(mode="after") - def validate_subworkflows(self) -> Self: - subworkflows = { - k: v - for k, v in self.model_dump().items() - if k not in _CreateTaskRequest.model_fields - } - - for workflow_name, workflow_definition in subworkflows.items(): - try: - WorkflowType.model_validate(workflow_definition) - setattr(self, workflow_name, WorkflowType(workflow_definition)) - except Exception as e: - raise ValueError(f"Invalid subworkflow '{workflow_name}': {str(e)}") - return self - - CreateOrUpdateTaskRequest = CreateTaskRequest _PatchTaskRequest = PatchTaskRequest