diff --git a/.github/workflows/model_onboarding.yml b/.github/workflows/model_onboarding.yml new file mode 100644 index 000000000..328838be5 --- /dev/null +++ b/.github/workflows/model_onboarding.yml @@ -0,0 +1,144 @@ +name: Model onboarding + +# On-demand run of the model-onboarding testcase against an arbitrary model. +# Nothing is committed: the model spec comes from the workflow inputs and is +# written into input.json at runtime. Start a run from the Actions UI +# ("Run workflow") or with `gh workflow run model_onboarding.yml -f ...`. + +on: + workflow_dispatch: + inputs: + model_name: + description: "Vendor-qualified model ID (e.g. gpt-5.2-2025-12-11, anthropic.claude-sonnet-4-5-20250929-v1:0)" + required: true + type: string + paths: + description: "Comma-separated get_chat_model code paths to exercise" + required: true + type: string + default: "azure_responses,azure_chat_completions" + files: + description: "Comma-separated file attachments (image,pdf). Leave empty for a text-only model." + required: false + type: string + default: "image,pdf" + prompt: + description: "Prompt sent to the model" + required: false + type: string + default: "Describe the content of this file in one sentence." + agenthub_config: + description: "AgentHub config header value (must exist in the target tenant)" + required: false + type: string + default: "agentsplayground" + environments: + description: "Which environments to run against" + required: true + type: choice + default: "alpha,staging" + options: + - "alpha" + - "staging" + - "cloud" + - "alpha,staging" + - "alpha,staging,cloud" + +permissions: + contents: read + actions: read + +jobs: + # Turn the comma-separated `environments` input into a JSON array the matrix + # can consume, so a single dispatch can fan out across chosen environments. + resolve-environments: + runs-on: ubuntu-latest + outputs: + environments: ${{ steps.split.outputs.environments }} + steps: + - name: Split environments input + id: split + run: | + envs_json=$(echo "${{ inputs.environments }}" \ + | tr ',' '\n' \ + | sed 's/[[:space:]]//g' \ + | grep -v '^$' \ + | jq -R -s -c 'split("\n") | map(select(length > 0))') + echo "Resolved environments: $envs_json" + echo "environments=$envs_json" >> "$GITHUB_OUTPUT" + + model-onboarding: + needs: [resolve-environments] + runs-on: ubuntu-latest + timeout-minutes: 10 + container: + image: ghcr.io/astral-sh/uv:python3.12-bookworm + env: + UIPATH_JOB_KEY: "3a03d5cb-fa21-4021-894d-a8e2eda0afe0" + UIPATH_TRACING_ENABLED: false + strategy: + fail-fast: false + matrix: + environment: ${{ fromJson(needs.resolve-environments.outputs.environments) }} + + name: "model-onboarding / ${{ matrix.environment }} / ${{ inputs.model_name }}" + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install Dependencies + run: uv sync + + - name: Build input.json from dispatch inputs + working-directory: testcases/model-onboarding + env: + MODEL_NAME: ${{ inputs.model_name }} + PATHS: ${{ inputs.paths }} + FILES: ${{ inputs.files }} + PROMPT: ${{ inputs.prompt }} + AGENTHUB_CONFIG: ${{ inputs.agenthub_config }} + run: | + # Normalize the comma-separated lists in the shell (trim each item, + # drop blanks) so jq only has to split on commas — this avoids regex + # backslash-escaping hazards inside YAML. jq then builds valid JSON + # regardless of characters in the free-text inputs. + normalize() { + echo "$1" | tr ',' '\n' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//' \ + | grep -v '^$' | paste -sd, - + } + PATHS_CLEAN=$(normalize "$PATHS") + FILES_CLEAN=$(normalize "$FILES") + + jq -n \ + --arg prompt "$PROMPT" \ + --arg model_name "$MODEL_NAME" \ + --arg agenthub_config "$AGENTHUB_CONFIG" \ + --arg paths "$PATHS_CLEAN" \ + --arg files "$FILES_CLEAN" \ + '{ + prompt: $prompt, + model_spec: { + model_name: $model_name, + agenthub_config: $agenthub_config, + paths: ($paths | split(",") | map(select(length > 0))), + files: ($files | split(",") | map(select(length > 0))) + } + }' > input.json + echo "=== Generated input.json ===" + cat input.json + + - name: Run testcase + env: + ENVIRONMENT: ${{ matrix.environment }} + CLIENT_ID: ${{ matrix.environment == 'alpha' && secrets.ALPHA_TEST_CLIENT_ID || matrix.environment == 'staging' && secrets.STAGING_TEST_CLIENT_ID || matrix.environment == 'cloud' && secrets.CLOUD_TEST_CLIENT_ID }} + CLIENT_SECRET: ${{ matrix.environment == 'alpha' && secrets.ALPHA_TEST_CLIENT_SECRET || matrix.environment == 'staging' && secrets.STAGING_TEST_CLIENT_SECRET || matrix.environment == 'cloud' && secrets.CLOUD_TEST_CLIENT_SECRET }} + BASE_URL: ${{ matrix.environment == 'alpha' && secrets.ALPHA_BASE_URL || matrix.environment == 'staging' && secrets.STAGING_BASE_URL || matrix.environment == 'cloud' && secrets.CLOUD_BASE_URL }} + working-directory: testcases/model-onboarding + run: | + echo "Running model-onboarding" + echo "Environment: ${{ matrix.environment }}" + echo "Model: ${{ inputs.model_name }}" + + bash run.sh + bash ../common/validate_output.sh diff --git a/testcases/model-onboarding/README.md b/testcases/model-onboarding/README.md new file mode 100644 index 000000000..c3ec2b573 --- /dev/null +++ b/testcases/model-onboarding/README.md @@ -0,0 +1,119 @@ +# model-onboarding testcase + +Exercises **one runtime-specified model** across the distinct `get_chat_model` +code paths it is expected to support, plus optional file attachments. Rolls +every `path × file` cell up into a single `success` boolean and asserts on both +the output and the emitted traces. + +Unlike `multimodal-invoke` (which hardcodes its model matrix), the model here is +**input**. To onboard a model, edit `input.json` — no code change. + +## The one file you edit: `input.json` + +```json +{ + "prompt": "Describe the content of this file in one sentence.", + "model_spec": { + "model_name": "gpt-5.2-2025-12-11", + "paths": ["azure_responses", "azure_chat_completions"], + "agenthub_config": "agentsplayground", + "files": ["image", "pdf"] + } +} +``` + +- **`model_name`** — the vendor-qualified model ID. Note a single logical model + may need a *different* ID per vendor family. +- **`paths`** — which `get_chat_model` code paths to exercise. Valid keys: + `azure_responses`, `azure_chat_completions`, `vertex`, `bedrock_converse`, + `bedrock_invoke`, `anthropic_sdk`. List only the paths the model actually + ships on — a model ID sent to a vendor it doesn't exist on is a guaranteed + (and misleading) failure. +- **`agenthub_config`** — AgentHub config header value; must exist in the tenant + behind your `BASE_URL`. Defaults to `agentsplayground`. +- **`files`** — file attachments to test. Valid keys: `image`, `pdf`. Use `[]` + for a **text-only** model — an empty list runs a plain reachability check via + `ainvoke` instead of a multimodal call. + +## Prerequisites (external to the repo) + +- Model IDs per path you list. +- Credentials for the target env: alpha (`ALPHA_TEST_CLIENT_ID` / + `ALPHA_TEST_CLIENT_SECRET` / `ALPHA_BASE_URL`), staging (`STAGING_*`), or + prod (`CLOUD_*` — prod is named `cloud` in this repo). +- The `agenthub_config` must exist in the target tenant. +- The vendor account behind the model must have it enabled in that tenant + (Bedrock region, Vertex project, Azure deployment) — otherwise a `✗` cell is + about provisioning, not the SDK. + +## Mechanism A — local run (fast iteration) + +From inside this directory: + +```bash +export CLIENT_ID=... # alpha or staging pair +export CLIENT_SECRET=... +export BASE_URL=... +export UIPATH_JOB_KEY=3a03d5cb-fa21-4021-894d-a8e2eda0afe0 +export UIPATH_TRACING_ENABLED=false + +bash run.sh # sync -> auth -> init -> pack -> run x2 +bash ../common/validate_output.sh # prints output.json, runs src/assert.py +``` + +The answer is `assert.py`'s exit code plus the `result_summary` grid it prints. +Exit 0 = model good on that env. Non-zero = the summary names the failing cell +and the truncated error. + +## Mechanism B — on-demand CI run (no commit, pick the model at start) + +`.github/workflows/model_onboarding.yml` runs this testcase against **any model +you name at dispatch time**. The model spec comes from the workflow inputs and +is written into `input.json` at runtime — you never edit or commit a file to +change the model. + +**From the GitHub UI:** Actions → "Model onboarding" → "Run workflow", fill in +`model_name`, `paths`, `files`, and pick the environment(s), then Run. + +**From the CLI:** + +```bash +gh workflow run model_onboarding.yml \ + -f model_name="anthropic.claude-sonnet-4-5-20250929-v1:0" \ + -f paths="bedrock_converse,bedrock_invoke,anthropic_sdk" \ + -f files="image,pdf" \ + -f environments="alpha,staging,cloud" + +gh run watch $(gh run list --workflow=model_onboarding.yml --limit 1 --json databaseId -q '.[0].databaseId') +# on failure: +gh run view --log | grep -A30 "Test Results" +``` + +Environments are selectable (`alpha`, `staging`, `cloud`, or combinations); +each leg uses its own `ALPHA_*` / `STAGING_*` / `CLOUD_*` secrets. The `cloud` +leg is prod. + +> Prod note: a failing `cloud` leg for a brand-new model usually means the model +> is not rolled out to the prod tenant yet — a provisioning signal, not a test +> bug. + +## Mechanism C — push/PR gate (the committed `input.json`) + +`.github/workflows/integration_tests.yml` auto-discovers this directory (the +hyphen in `model-onboarding` is required) and runs it — using the **committed** +`input.json` — across `alpha`, `staging`, and `cloud` on every push/PR, as one +leg of the full integration matrix. This is the regression gate; use +Mechanism B for ad-hoc model runs. + +Optional enhancement (not wired): add a `workflow_dispatch:` trigger with a +`model_spec` input so a run can be launched without a commit. + +## What gets asserted (`src/assert.py`) + +1. A `.nupkg` was produced. +2. `status == "successful"` and the `output` block exists. +3. `success is True` and `result_summary` is non-empty. +4. `"Successful execution."` appears in `local_run_output.log` (the second, + empty-`UIPATH_JOB_KEY` run). +5. Traces contain the `run_model_onboarding` CHAIN span and at least one `LLM` + span from a reachable client class (`expected_traces.json`). diff --git a/testcases/model-onboarding/expected_traces.json b/testcases/model-onboarding/expected_traces.json new file mode 100644 index 000000000..a733393bb --- /dev/null +++ b/testcases/model-onboarding/expected_traces.json @@ -0,0 +1,23 @@ +{ + "description": "Model-onboarding trace assertions - the node ran and at least one LLM span was emitted by one of the reachable client classes.", + "required_spans": [ + { + "name": "run_model_onboarding", + "attributes": { + "openinference.span.kind": "CHAIN" + } + }, + { + "name": [ + "UiPathAzureChatOpenAI", + "UiPathChatGoogleGenerativeAI", + "UiPathChatBedrockConverse", + "UiPathChatBedrock", + "UiPathChatAnthropicBedrock" + ], + "attributes": { + "openinference.span.kind": "LLM" + } + } + ] +} diff --git a/testcases/model-onboarding/input.json b/testcases/model-onboarding/input.json new file mode 100644 index 000000000..b94149b73 --- /dev/null +++ b/testcases/model-onboarding/input.json @@ -0,0 +1,9 @@ +{ + "prompt": "Describe the content of this file in one sentence.", + "model_spec": { + "model_name": "gpt-5.2-2025-12-11", + "paths": ["azure_responses", "azure_chat_completions"], + "agenthub_config": "agentsplayground", + "files": ["image", "pdf"] + } +} diff --git a/testcases/model-onboarding/langgraph.json b/testcases/model-onboarding/langgraph.json new file mode 100644 index 000000000..a2f64dcb4 --- /dev/null +++ b/testcases/model-onboarding/langgraph.json @@ -0,0 +1,6 @@ +{ + "dependencies": ["."], + "graphs": { + "agent": "./src/main.py:graph" + } +} diff --git a/testcases/model-onboarding/pyproject.toml b/testcases/model-onboarding/pyproject.toml new file mode 100644 index 000000000..5eff568c0 --- /dev/null +++ b/testcases/model-onboarding/pyproject.toml @@ -0,0 +1,16 @@ +[project] +name = "model-onboarding" +version = "0.0.1" +description = "Onboard a runtime-specified LLM model across get_chat_model code paths" +authors = [{ name = "John Doe", email = "john.doe@myemail.com" }] +dependencies = [ + "langgraph>=0.2.70", + "langchain-core>=0.3.34", + "langgraph-checkpoint-sqlite>=2.0.3", + "uipath-langchain[vertex,bedrock]", + "pydantic>=2.10.6", +] +requires-python = ">=3.11" + +[tool.uv.sources] +uipath-langchain = { path = "../../", editable = true } diff --git a/testcases/model-onboarding/run.sh b/testcases/model-onboarding/run.sh new file mode 100644 index 000000000..77c43f4d8 --- /dev/null +++ b/testcases/model-onboarding/run.sh @@ -0,0 +1,21 @@ +#!/bin/bash +set -e + +echo "Syncing dependencies..." +uv sync + +echo "Authenticating with UiPath..." +uv run uipath auth --client-id="$CLIENT_ID" --client-secret="$CLIENT_SECRET" --base-url="$BASE_URL" + +echo "Initializing the project..." +uv run uipath init + +echo "Packing agent..." +uv run uipath pack + +echo "Running agent..." +uv run uipath run agent --file input.json + +echo "Running agent again with empty UIPATH_JOB_KEY..." +export UIPATH_JOB_KEY="" +uv run uipath run agent --trace-file .uipath/traces.jsonl --file input.json >> local_run_output.log diff --git a/testcases/model-onboarding/src/assert.py b/testcases/model-onboarding/src/assert.py new file mode 100644 index 000000000..0bd06786b --- /dev/null +++ b/testcases/model-onboarding/src/assert.py @@ -0,0 +1,53 @@ +import json +import os + +from trace_assert import assert_traces + +print("Checking model-onboarding agent output...") + +uipath_dir = ".uipath" +assert os.path.exists(uipath_dir), "NuGet package directory (.uipath) not found" + +nupkg_files = [f for f in os.listdir(uipath_dir) if f.endswith(".nupkg")] +assert nupkg_files, "NuGet package file (.nupkg) not found in .uipath directory" +print(f"NuGet package found: {nupkg_files[0]}") + +output_file = "__uipath/output.json" +assert os.path.isfile(output_file), "Agent output file not found" +print("Agent output file found") + +with open(output_file, "r", encoding="utf-8") as f: + output_data = json.load(f) + +status = output_data.get("status") +assert status == "successful", f"Agent execution failed with status: {status}" +print("Agent execution status: successful") + +assert "output" in output_data, "Missing 'output' field in agent response" +output_content = output_data["output"] + +assert "success" in output_content, "Missing 'success' field in output" +success = output_content["success"] + +assert "result_summary" in output_content, "Missing 'result_summary' field in output" +result_summary = output_content["result_summary"] +assert result_summary and result_summary.strip() != "", "Result summary is empty" + +print("\nTest Results:") +print(f" Success: {success}") +print(f" Summary:\n{result_summary}") + +assert success is True, "Test did not succeed. See detailed results above." + +# The second (empty UIPATH_JOB_KEY) local run appends to this log. +with open("local_run_output.log", "r", encoding="utf-8") as f: + local_run_output = f.read() + +assert "Successful execution." in local_run_output, ( + f"Response does not contain 'Successful execution.'. " + f"Actual response: {local_run_output}" +) + +assert_traces(".uipath/traces.jsonl", "expected_traces.json") + +print("All validations passed successfully!") diff --git a/testcases/model-onboarding/src/main.py b/testcases/model-onboarding/src/main.py new file mode 100644 index 000000000..303fa1a8b --- /dev/null +++ b/testcases/model-onboarding/src/main.py @@ -0,0 +1,279 @@ +"""Model-parameterized onboarding test case. + +Exercises a single model — supplied at runtime via ``input.json`` — across the +distinct ``get_chat_model`` code paths that model is expected to support, plus +an optional set of file attachments. Every ``path x file`` cell is invoked +independently; failures are caught per cell and rolled up into one ``success`` +boolean, mirroring the ``multimodal-invoke`` contract so this project drops into +the existing integration-test matrix unchanged. + +The model is NOT hardcoded. Edit ``input.json`` to onboard any model: + + { + "prompt": "Describe the content of this file in one sentence.", + "model_spec": { + "model_name": "gpt-5.2-2025-12-11", + "paths": ["azure_responses", "azure_chat_completions"], + "agenthub_config": "agentsplayground", + "files": ["image", "pdf"] + } + } + +``paths`` is explicit on purpose: a model ID is only valid on the vendor +families it actually ships on, so the caller declares which surfaces to test +rather than the code guessing from the name. ``files`` may be empty for +text-only models — an empty list runs a plain ``ainvoke`` reachability check. +""" + +import logging +from typing import Callable + +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.memory import MemorySaver +from langgraph.graph import END, START, MessagesState, StateGraph +from pydantic import BaseModel, Field +from uipath.llm_client.settings import PlatformSettings +from uipath_langchain_client.clients.bedrock.chat_models import ( + UiPathChatAnthropicBedrock, +) +from uipath_langchain_client.settings import ( + ApiFlavor, + UiPathBaseSettings, +) + +from uipath_langchain.agent.multimodal.invoke import llm_call_with_files +from uipath_langchain.agent.multimodal.types import FileInfo +from uipath_langchain.chat.chat_model_factory import get_chat_model + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- # +# Path registry: one builder per distinct get_chat_model code path. +# +# Each builder returns a configured BaseChatModel for the given model name and +# client settings. The keys are the strings the caller lists in +# `model_spec.paths`. Adding a new reachable class is a one-line addition here. +# --------------------------------------------------------------------------- # +PathBuilder = Callable[[str, UiPathBaseSettings], object] + +PATH_REGISTRY: dict[str, PathBuilder] = { + # VendorType.OPENAI (UiPath-owned) -> UiPathAzureChatOpenAI (responses API) + "azure_responses": lambda model, settings: get_chat_model( + model=model, + client_settings=settings, + api_flavor=ApiFlavor.RESPONSES, + temperature=0.0, + max_tokens=200, + ), + # VendorType.OPENAI + CHAT_COMPLETIONS -> UiPathAzureChatOpenAI (chat API) + "azure_chat_completions": lambda model, settings: get_chat_model( + model=model, + client_settings=settings, + api_flavor=ApiFlavor.CHAT_COMPLETIONS, + temperature=0.0, + max_tokens=200, + ), + # VendorType.VERTEXAI (Google family) -> UiPathChatGoogleGenerativeAI + "vertex": lambda model, settings: get_chat_model( + model=model, + client_settings=settings, + temperature=0.0, + max_tokens=200, + ), + # VendorType.AWSBEDROCK (UiPath-owned) -> UiPathChatBedrockConverse + "bedrock_converse": lambda model, settings: get_chat_model( + model=model, + client_settings=settings, + api_flavor=ApiFlavor.CONVERSE, + temperature=0.0, + max_tokens=200, + ), + # VendorType.AWSBEDROCK + INVOKE -> UiPathChatBedrock + "bedrock_invoke": lambda model, settings: get_chat_model( + model=model, + client_settings=settings, + api_flavor=ApiFlavor.INVOKE, + temperature=0.0, + max_tokens=200, + ), + # Direct instantiation -> UiPathChatAnthropicBedrock (not factory-reachable) + "anthropic_sdk": lambda model, settings: UiPathChatAnthropicBedrock( + model_name=model, + settings=settings, + temperature=0.0, + max_tokens=200, + ), +} + + +# --------------------------------------------------------------------------- # +# File registry: named file attachments the caller selects via +# `model_spec.files`. Public, reachable URLs so the run environment can fetch +# them. Extend as needed for other formats you want to onboard against. +# --------------------------------------------------------------------------- # +FILE_REGISTRY: dict[str, FileInfo] = { + "image": FileInfo( + url="https://www.w3schools.com/css/img_5terre.jpg", + name="img_5terre.jpg", + mime_type="image/jpeg", + ), + "pdf": FileInfo( + url="https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf", + name="dummy.pdf", + mime_type="application/pdf", + ), +} + + +class ModelSpec(BaseModel): + """Runtime specification for the model under test.""" + + model_name: str = Field(description="Vendor-qualified model identifier.") + paths: list[str] = Field( + description="get_chat_model code paths to exercise; keys of PATH_REGISTRY.", + ) + agenthub_config: str = Field( + default="agentsplayground", + description="AgentHub config header value; must exist in the target tenant.", + ) + files: list[str] = Field( + default_factory=list, + description="File attachments to test; keys of FILE_REGISTRY. Empty => " + "text-only reachability check via plain ainvoke.", + ) + + +class GraphInput(BaseModel): + prompt: str = Field(default="Describe the content of this file in one sentence.") + model_spec: ModelSpec + + +class GraphOutput(BaseModel): + success: bool + result_summary: str + + +class GraphState(MessagesState): + prompt: str + model_spec: dict + success: bool + result_summary: str + model_results: dict + + +def _build_model(path: str, model_name: str, settings: UiPathBaseSettings) -> object: + """Build a model instance for a registered path. + + Raises: + KeyError: If ``path`` is not a registered PATH_REGISTRY key. + """ + builder = PATH_REGISTRY[path] + return builder(model_name, settings) + + +async def run_model_onboarding(state: GraphState) -> dict: + spec = ModelSpec.model_validate(state["model_spec"]) + messages = [HumanMessage(content=state["prompt"])] + + # Empty files => a single text-only cell (llm_call_with_files does a plain + # ainvoke when the file list is empty). + selected_files: list[tuple[str, list[FileInfo]]] + if spec.files: + selected_files = [(name, [FILE_REGISTRY[name]]) for name in spec.files] + else: + selected_files = [("text-only", [])] + + try: + client_settings = PlatformSettings(agenthub_config=spec.agenthub_config) + except Exception as e: + # Settings need UiPath auth env vars (set by `uipath auth`). If they are + # missing the whole run is moot; surface it as a legible failure. + logger.error(f"PlatformSettings construction failed: {e}") + return { + "success": False, + "result_summary": f"settings: ✗ {str(e)[:120]}", + "model_results": {}, + } + + model_results: dict[str, dict[str, str]] = {} + for path in spec.paths: + logger.info(f"Testing path '{path}' with model '{spec.model_name}'...") + + if path not in PATH_REGISTRY: + logger.error(f" unknown path '{path}'") + model_results[path] = {"__path__": f"✗ unknown path '{path}'"} + continue + + try: + model = _build_model(path, spec.model_name, client_settings) + logger.info(f" Created: {type(model).__name__}") + except Exception as e: # model construction itself can fail + logger.error(f" construction failed: {e}") + model_results[path] = {"__build__": f"✗ {str(e)[:60]}"} + continue + + cell_results: dict[str, str] = {} + for label, files in selected_files: + logger.info(f" {label}...") + try: + response: AIMessage = await llm_call_with_files( + messages, files, model + ) + # Guard against an empty/blank completion counting as success. + if response.content and str(response.content).strip(): + logger.info(f" {label}: ✓") + cell_results[label] = "✓" + else: + logger.warning(f" {label}: ✗ empty response") + cell_results[label] = "✗ empty response" + except Exception as e: + logger.error(f" {label}: ✗ {e}") + cell_results[label] = f"✗ {str(e)[:60]}" + model_results[path] = cell_results + + summary_lines = [] + for path, results in model_results.items(): + summary_lines.append(f"{path}:") + for cell_name, result in results.items(): + summary_lines.append(f" {cell_name}: {result}") + + has_failures = any( + "✗" in v for results in model_results.values() for v in results.values() + ) + # A spec with no runnable paths is a failure, not a vacuous success. + if not model_results: + has_failures = True + summary_lines.append("(no paths specified)") + + return { + "success": not has_failures, + "result_summary": "\n".join(summary_lines), + "model_results": model_results, + } + + +async def return_results(state: GraphState) -> GraphOutput: + logger.info(f"Success: {state['success']}") + logger.info(f"Summary:\n{state['result_summary']}") + return GraphOutput( + success=state["success"], + result_summary=state["result_summary"], + ) + + +def build_graph() -> StateGraph: + builder = StateGraph(GraphState, input_schema=GraphInput, output_schema=GraphOutput) + + builder.add_node("run_model_onboarding", run_model_onboarding) + builder.add_node("results", return_results) + + builder.add_edge(START, "run_model_onboarding") + builder.add_edge("run_model_onboarding", "results") + builder.add_edge("results", END) + + memory = MemorySaver() + return builder.compile(checkpointer=memory) + + +graph = build_graph()