|
| 1 | +from enum import Enum |
| 2 | +from typing import Annotated, Any, Literal, TypeAlias |
| 3 | + |
| 4 | +from models_library import projects |
| 5 | +from pydantic import BaseModel, Field |
| 6 | + |
| 7 | +FunctionID: TypeAlias = projects.ProjectID |
| 8 | + |
| 9 | + |
| 10 | +class FunctionSchema(BaseModel): |
| 11 | + schema_dict: dict[str, Any] | None # JSON Schema |
| 12 | + |
| 13 | + |
| 14 | +class FunctionInputSchema(FunctionSchema): ... |
| 15 | + |
| 16 | + |
| 17 | +class FunctionOutputSchema(FunctionSchema): ... |
| 18 | + |
| 19 | + |
| 20 | +class FunctionClass(str, Enum): |
| 21 | + project = "project" |
| 22 | + python_code = "python_code" |
| 23 | + |
| 24 | + |
| 25 | +class FunctionInputs(BaseModel): |
| 26 | + inputs_dict: dict[str, Any] | None # JSON Schema |
| 27 | + |
| 28 | + |
| 29 | +class FunctionOutputs(BaseModel): |
| 30 | + outputs_dict: dict[str, Any] | None # JSON Schema |
| 31 | + |
| 32 | + |
| 33 | +class Function(BaseModel): |
| 34 | + uid: FunctionID | None = None |
| 35 | + title: str | None = None |
| 36 | + description: str | None = None |
| 37 | + input_schema: FunctionInputSchema | None = None |
| 38 | + output_schema: FunctionOutputSchema | None = None |
| 39 | + |
| 40 | + |
| 41 | +class StudyFunction(Function): |
| 42 | + function_type: Literal["study"] = "study" |
| 43 | + study_url: str |
| 44 | + |
| 45 | + |
| 46 | +class PythonCodeFunction(Function): |
| 47 | + function_type: Literal["python_code"] = "python_code" |
| 48 | + code_url: str |
| 49 | + |
| 50 | + |
| 51 | +FunctionUnion: TypeAlias = Annotated[ |
| 52 | + StudyFunction | PythonCodeFunction, |
| 53 | + Field(discriminator="function_type"), |
| 54 | +] |
| 55 | + |
| 56 | +FunctionJobID: TypeAlias = projects.ProjectID |
| 57 | +FunctionJobCollectionID: TypeAlias = projects.ProjectID |
| 58 | + |
| 59 | + |
| 60 | +class FunctionJob(BaseModel): |
| 61 | + uid: FunctionJobID |
| 62 | + title: str | None |
| 63 | + description: str | None |
| 64 | + status: str |
| 65 | + function_uid: FunctionID |
| 66 | + inputs: FunctionInputs | None |
| 67 | + outputs: FunctionOutputs | None |
| 68 | + |
| 69 | + |
| 70 | +class FunctionJobStatus(BaseModel): |
| 71 | + status: str |
| 72 | + |
| 73 | + |
| 74 | +class FunctionJobCollection(BaseModel): |
| 75 | + """Model for a collection of function jobs""" |
| 76 | + |
| 77 | + id: FunctionJobCollectionID |
| 78 | + title: str | None |
| 79 | + description: str | None |
| 80 | + job_ids: list[FunctionJobID] |
| 81 | + status: str |
| 82 | + |
| 83 | + |
| 84 | +class FunctionJobCollectionStatus(BaseModel): |
| 85 | + status: list[str] |
0 commit comments