diff --git a/.github/workflows/integration-addon.yaml b/.github/workflows/integration-addon.yaml deleted file mode 100644 index a0b268d..0000000 --- a/.github/workflows/integration-addon.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# This workflow will install Python dependencies, run tests and lint with a single version of Python -# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python - -name: Addon Integration Tests - -on: - push: - branches: ["main"] - pull_request: - branches: ["main"] - -permissions: - contents: read - -jobs: - build: - name: Addon Integration Tests - runs-on: ubuntu-latest - - steps: - - uses: actions/checkout@v4 - - name: Set up Python 3.11 - uses: actions/setup-python@v5 - with: - python-version: "3.11" - - name: Install dependencies - run: | - python -m pip install --upgrade uv - uv pip install --system -r tests/test-requirements.txt python/ - - name: Lint with Ruff - run: | - ruff check . - - name: install playwright deps - run: | - playwright install --with-deps chromium - - name: Copy python source + services into build context because builder action doesn't support --build-context - id: copy - run: bash scripts/copy_content_to_addon_context.sh - - name: Test with pytest - run: | - pytest -v -s -rA -c tests/pytest.ini --deploy-mode=addon --replay-mode=replay diff --git a/.github/workflows/integration-local.yaml b/.github/workflows/integration-local.yaml deleted file mode 100644 index 8563d58..0000000 --- a/.github/workflows/integration-local.yaml +++ /dev/null @@ -1,38 +0,0 @@ -# This workflow will install Python dependencies, run tests and lint with a single version of Python -# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python - -name: Local Integration Tests - -on: - push: - branches: ["main"] - pull_request: - branches: ["main"] - -permissions: - contents: read - -jobs: - build: - name: Local Integration Tests - runs-on: ubuntu-latest - - steps: - - uses: actions/checkout@v4 - - name: Set up Python 3.11 - uses: actions/setup-python@v5 - with: - python-version: "3.11" - - name: Install dependencies - run: | - python -m pip install --upgrade uv - uv pip install --system -r tests/test-requirements.txt python/ - - name: Lint with Ruff - run: | - ruff check . - - name: install playwright deps - run: | - playwright install --with-deps chromium - - name: Test with pytest - run: | - pytest -v -s -rA -c tests/pytest.ini --deploy-mode=local --replay-mode=replay diff --git a/.github/workflows/integration-k3d.yaml b/.github/workflows/integration.yaml similarity index 56% rename from .github/workflows/integration-k3d.yaml rename to .github/workflows/integration.yaml index f82898b..8515273 100644 --- a/.github/workflows/integration-k3d.yaml +++ b/.github/workflows/integration.yaml @@ -1,7 +1,7 @@ # This workflow will install Python dependencies, run tests and lint with a single version of Python # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python -name: K3D Integration Tests +name: Integration Tests on: push: @@ -13,10 +13,13 @@ permissions: contents: read jobs: - build: - name: K3D Integration Tests + test: + name: Integration Tests runs-on: ubuntu-latest - + strategy: + fail-fast: false + matrix: + deploymode: ["local", "k3d", "addon"] steps: - uses: actions/checkout@v4 - name: Set up Python 3.11 @@ -30,7 +33,9 @@ jobs: - name: Lint with Ruff run: | ruff check . + ruff check --select I . - name: Create cluster + if: startsWith(matrix.deploymode, 'k3') == true shell: bash run: bash scripts/setup_k3d.sh env: @@ -38,10 +43,26 @@ jobs: SKIP_CREATION: true SKIP_REGISTRY_CREATION: true SKIP_READINESS: true + - name: install playwright deps + run: | + playwright install --with-deps chromium + - name: Copy content to addon context + if: startsWith(matrix.deploymode, 'addon') == true + id: copy + run: bash scripts/copy_content_to_addon_context.sh - name: Install Dapr shell: bash run: | wget -q https://raw.githubusercontent.com/dapr/cli/master/install/install.sh -O - | /bin/bash + - name: initialize dapr + run: | + dapr init --slim - name: Test with pytest run: | - pytest -v -s -rA -c tests/pytest.ini --deploy-mode=k3d --replay-mode=replay + pytest -v -s -rA -c tests/pytest.ini --deploy-mode=${{ matrix.deploymode }} --replay-mode=replay + - name: Upload debug playwright screenshots + if: failure() + uses: actions/upload-artifact@v4 + with: + name: playwright-screenshots-${{ matrix.deploymode }} + path: "**/fail-*.png" diff --git a/README.md b/README.md index 146cc40..ae8ecfe 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ -# Home ~~Automation~~ Intelligence +# Mindctrl: Home ~~Automation~~ Intelligence :thought_balloon: -Manage your home automation LLM prompts, available LLMs and evaluate changes with MLflow +An Intelligence hosting platform for automating your life [![Addon Builder](https://github.com/akshaya-a/mindctrl/actions/workflows/addon-builder.yaml/badge.svg)](https://github.com/akshaya-a/mindctrl/actions/workflows/addon-builder.yaml) [![Addon Integration](https://github.com/akshaya-a/mindctrl/actions/workflows/integration-addon.yaml/badge.svg)](https://github.com/akshaya-a/mindctrl/actions/workflows/integration-addon.yaml) @@ -9,7 +9,59 @@ Manage your home automation LLM prompts, available LLMs and evaluate changes wit --- -![promptlab](assets/mlflowhass-promptlab.png) +![promptlab](assets/mctrl_arch_v1.png "Mindctrl Architecture") +_Mindctrl architecture, courtesy of the amazing [excalidraw](https://excalidraw.com/)_ + +A more disciplined incorporation of AI into your home automation setup. This project aims to provide a platform for managing and evaluating conversational prompts for Large Language Models (LLMs) alongside traditional ML techniques in Home Assistant or other home automation platforms (bring your own events and triggers). + +## Features + +- [x] Manage deployed models or prompts with a versioned registry +- [x] Evaluate prompts with a simple UI +- [ ] Customize your agents + +## Platform + +- [ ] Scale from [a single device](#home-assistant-addon) :computer: ... + - [x] to [a selfhosted cluster](#kubernetes) :globe_with_meridians: ... + - [ ] to [the cloud](#cloud) :cloud: +- [ ] Deploy onto + - [x] amd64 or ... + - [x] aarch64 architectures, and feel free to contribute support for ... + - [ ] the rest +- [ ] Developed on + - [x] Linux, but it's mostly containerized so it should work on ... + - [ ] Windows or ... + - [ ] MacOS with little effort + +## Goals / OKRs + +1. [Increase Spousal Approval Rating](#reason-about-state-changes) (SAR) by 10% by the end of the quarter +2. [Gain a better understanding](#why) of true AI/ML application deployment + tooling required for successful end-to-end scenarios by building from first principles (no -> light -> heavy AI-focused frameworks) +3. Incorporate latest GenAI techniques to evaluate utility + - Memory, tool-calling and RAG have lots of room to grow +4. Justify the purchase of a new GPU + +## Getting Started + +> [!WARNING] +> This project is in early development :see_no_evil: and is not yet ready for production use. Or any use, really. But if you're feeling adventurous, read on! + +First, decide how you want to integrate Mindctrl into your home automation setup. You can choose from the following options: + +### Home Assistant Addon + +1. [Install the Mindctrl Home Assistant Addon](https://www.home-assistant.io/common-tasks/os#installing-third-party-add-ons) +2. Configure the addon with your database and (optional) broker details + +### Kubernetes + +- [ ] TODO: [TuringPi + K3S instructions](https://docs.turingpi.com/docs/how-to-plan-kubernetes-installation) more scoped to mindctrl. In the interim, intrepid explorers can look at what [the tests do](tests/utils/cluster.py) +- [ ] TODO: Convert the kubectl commands to helm charts + +### Cloud + +- [ ] TODO: Azure Container Apps instructions ## Why? @@ -19,6 +71,8 @@ Manage your home automation LLM prompts, available LLMs and evaluate changes wit ### Better manage conversational prompts +![promptlab](assets/mlflowhass-promptlab.png) + Home Assistant has a convenient prompt template tool to generate prompts for LLMs. However, it's not easy to manage these prompts. I can't tell if my new prompt is better than the last one, change tracking is not easy, and live editing means switching between the developer tools and the prompt template tool. There's a better way! Enter MLflow with its new PromptLab UI and integrated evaluation tools + tracking. ### Reason about state changes @@ -33,19 +87,7 @@ LLMs can be used to reason about state changes more naturally. Can we send the s - (motion, time of day, device usage) -> "is AK asleep?" -### Wait what about Langchain? - -You heard about Langchain but not this whole MLflow business - what's up with that? [Read more about how all this fits together!](/docs/prompt-techniques.md) - -## How? - -### Getting Started - -1. Install the MLflow Gateway -2. Install the MLflow Tracking Service -3. Install the MLflow Home Integration - -## What? +## What are you talking about? - LLM: Large Language Model, a model that can generate text based on a prompt - [MLflow](https://mlflow.org/): An open source platform for the machine learning lifecycle diff --git a/addons/mindctrl/rootfs/usr/bin/run_dashboard.sh b/addons/mindctrl/rootfs/usr/bin/run_dashboard.sh index bcf7bef..a4f9127 100644 --- a/addons/mindctrl/rootfs/usr/bin/run_dashboard.sh +++ b/addons/mindctrl/rootfs/usr/bin/run_dashboard.sh @@ -12,5 +12,10 @@ if [[ -n "$ingress_entry" ]]; then bashio::log.info "running dashboard with prefix $SERVER_BASE_HREF" fi +bashio::log.info $HOME +ls -la $HOME/.dapr/bin +bashio::log.info "Starting dapr placement server" +$HOME/.dapr/bin/placement +bashio::log.info "Starting dapr dashboard" dapr dashboard -a 0.0.0.0 -p 9999 diff --git a/addons/mindctrl/rootfs/usr/bin/run_traefik.sh b/addons/mindctrl/rootfs/usr/bin/run_traefik.sh index 462b8ce..8e4bc76 100644 --- a/addons/mindctrl/rootfs/usr/bin/run_traefik.sh +++ b/addons/mindctrl/rootfs/usr/bin/run_traefik.sh @@ -22,7 +22,7 @@ export HASS_INGRESS_ENTRY="${ingress_entry}" bashio::log.info "Starting traefik..." /traefik version # TODO: until this is unified, keep in sync with testcontainer -/traefik --accesslog=true --accesslog.format=json --log.level=DEBUG --api=true --api.dashboard=true --api.insecure=true \ +/traefik --accesslog=true --accesslog.format=json --log.level=DEBUG --api=true \ --entrypoints.http.address=':80' \ - --ping=true \ + --ping=true --entryPoints.ping.address=:8082 --ping.entryPoint=ping \ --providers.file.filename /.context/services/ingress/traefik-config.yaml diff --git a/assets/mctrl_arch_v1.png b/assets/mctrl_arch_v1.png new file mode 100644 index 0000000..db6bb1a Binary files /dev/null and b/assets/mctrl_arch_v1.png differ diff --git a/custom_components/mindctrl/__init__.py b/custom_components/mindctrl/__init__.py index c247f96..66601ff 100644 --- a/custom_components/mindctrl/__init__.py +++ b/custom_components/mindctrl/__init__.py @@ -3,8 +3,10 @@ from __future__ import annotations +import asyncio + from homeassistant.components import conversation as haconversation -from homeassistant.components.hassio import AddonManager, AddonError, AddonState +from homeassistant.components.hassio import AddonError, AddonManager, AddonState from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ( @@ -13,15 +15,10 @@ from homeassistant.helpers import config_validation as cv from homeassistant.helpers.typing import ConfigType -import asyncio - from .addon import get_addon_manager - -from .const import ADDON_NAME, CONF_URL, CONF_USE_ADDON, DOMAIN, _LOGGER - -from .services import MindctrlClient, async_register_services +from .const import _LOGGER, ADDON_NAME, CONF_URL, CONF_USE_ADDON, DOMAIN from .conversation import MLflowAgent - +from .services import MindctrlClient, async_register_services CONNECT_TIMEOUT = 10 diff --git a/custom_components/mindctrl/addon.py b/custom_components/mindctrl/addon.py index 4d2c4e4..2c5efea 100644 --- a/custom_components/mindctrl/addon.py +++ b/custom_components/mindctrl/addon.py @@ -8,7 +8,7 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.singleton import singleton -from .const import ADDON_SLUG, DOMAIN, _LOGGER, ADDON_NAME +from .const import _LOGGER, ADDON_NAME, ADDON_SLUG, DOMAIN DATA_ADDON_MANAGER = f"{DOMAIN}_addon_manager" diff --git a/custom_components/mindctrl/config_flow.py b/custom_components/mindctrl/config_flow.py index 21fbbcf..5b30130 100644 --- a/custom_components/mindctrl/config_flow.py +++ b/custom_components/mindctrl/config_flow.py @@ -1,29 +1,29 @@ +import asyncio +import uuid from typing import Any + +import aiohttp +import voluptuous as vol from homeassistant import config_entries, exceptions -from homeassistant.core import HomeAssistant, callback from homeassistant.const import CONF_URL -from homeassistant.helpers.aiohttp_client import async_get_clientsession +from homeassistant.core import HomeAssistant, callback # from homeassistant.components.zeroconf import ZeroconfServiceInfo from homeassistant.data_entry_flow import ( FlowResult, ) -import voluptuous as vol -import asyncio -import aiohttp -import uuid +from homeassistant.helpers.aiohttp_client import async_get_clientsession from .const import ( + _LOGGER, ADDON_HOST, ADDON_PORT, CONF_ADDON_LOG_LEVEL, CONF_INTEGRATION_CREATED_ADDON, - DOMAIN, - _LOGGER, CONF_USE_ADDON, + DOMAIN, ) - DEFAULT_URL = f"http://{ADDON_HOST}:{ADDON_PORT}" TITLE = "mindctrl" diff --git a/custom_components/mindctrl/conversation.py b/custom_components/mindctrl/conversation.py index 2cdd48c..d139be7 100644 --- a/custom_components/mindctrl/conversation.py +++ b/custom_components/mindctrl/conversation.py @@ -1,17 +1,17 @@ +from typing import Literal + +import mlflow from homeassistant.components import conversation +from homeassistant.config_entries import ConfigEntry from homeassistant.const import MATCH_ALL from homeassistant.core import ( HomeAssistant, ) -from homeassistant.config_entries import ConfigEntry -from typing import Literal -from homeassistant.util import ulid -from homeassistant.helpers import intent, template from homeassistant.exceptions import ( TemplateError, ) -import mlflow - +from homeassistant.helpers import intent, template +from homeassistant.util import ulid from .const import _LOGGER diff --git a/custom_components/mindctrl/entity.py b/custom_components/mindctrl/entity.py index cd4af12..0b049f7 100644 --- a/custom_components/mindctrl/entity.py +++ b/custom_components/mindctrl/entity.py @@ -1,13 +1,14 @@ """AdGuard Home base entity.""" from __future__ import annotations + from abc import ABC, abstractmethod from homeassistant.config_entries import SOURCE_HASSIO, ConfigEntry from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity import Entity -from .const import ADDON_SLUG, DOMAIN, _LOGGER +from .const import _LOGGER, ADDON_SLUG, DOMAIN from .services import MindctrlClient diff --git a/custom_components/mindctrl/services.py b/custom_components/mindctrl/services.py index d0be672..d3fdbdb 100644 --- a/custom_components/mindctrl/services.py +++ b/custom_components/mindctrl/services.py @@ -1,5 +1,6 @@ -from homeassistant.helpers.aiohttp_client import async_get_clientsession - +import mlflow +import voluptuous as vol +from homeassistant.config_entries import ConfigEntry from homeassistant.core import ( HomeAssistant, ServiceCall, @@ -9,12 +10,10 @@ from homeassistant.exceptions import ( HomeAssistantError, ) -from homeassistant.config_entries import ConfigEntry - -import mlflow -from .const import DOMAIN, SERVICE_INVOKE_MODEL, _LOGGER, CONF_URL -import voluptuous as vol from homeassistant.helpers import config_validation as cv +from homeassistant.helpers.aiohttp_client import async_get_clientsession + +from .const import _LOGGER, CONF_URL, DOMAIN, SERVICE_INVOKE_MODEL class MindctrlClient(object): diff --git a/python/pyproject.toml b/python/pyproject.toml index 4107a38..d4c6306 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -13,6 +13,14 @@ classifiers = [ "Operating System :: OS Independent", "Private :: Do Not Upload" # Currently just for local + cluster packaging. Publish once evals are in ] +# TODO: Remove the prerelease +dependencies = [ + "dapr", + "dapr-ext-fastapi", + "dapr-ext-workflow>=0.4.0", + "durabletask>=0.1.1a1", + "mlflow-skinny" +] [project.scripts] mindctrl = "mindctrl.cli:cli" diff --git a/python/src/mindctrl/cli.py b/python/src/mindctrl/cli.py index 25ea980..823728a 100644 --- a/python/src/mindctrl/cli.py +++ b/python/src/mindctrl/cli.py @@ -1,15 +1,14 @@ import logging import os from typing import Optional -import click +import click from mlflow.deployments.cli import validate_config_path -from mlflow.environment_variables import MLFLOW_DEPLOYMENTS_CONFIG from mlflow.deployments.server.runner import monitor_config +from mlflow.environment_variables import MLFLOW_DEPLOYMENTS_CONFIG from .replay_server import ReplayRunner - _logger = logging.getLogger(__name__) diff --git a/python/src/mindctrl/config.py b/python/src/mindctrl/config.py index 290eb4b..442a79a 100644 --- a/python/src/mindctrl/config.py +++ b/python/src/mindctrl/config.py @@ -1,11 +1,26 @@ -from typing import Optional, Union, Literal +from functools import lru_cache +import json +import logging +import os +from typing import Any, Dict, Literal, Optional, Tuple, Type, Union + from pydantic import BaseModel, Field, SecretStr -from pydantic_settings import BaseSettings, SettingsConfigDict +from pydantic.fields import FieldInfo +from pydantic_settings import ( + BaseSettings, + PydanticBaseSettingsSource, + SettingsConfigDict, +) + +from mindctrl.const import CONFIGURATION_KEY, CONFIGURATION_STORE, SECRET_STORE + + +_logger = logging.getLogger(__name__) # this is just to make settings typing happy - I don't have another implementation yet -class UnknownEventsSettings(BaseModel): - events_type: Literal["unknown"] +class DisabledEventsSettings(BaseModel): + events_type: Literal["none"] = "none" class MqttEventsSettings(BaseModel): @@ -28,23 +43,72 @@ class PostgresStoreSettings(BaseModel): # Just to make typing happy for now - add dapr, sqlite, etc -class UnknownStoreSettings(BaseModel): - store_type: Literal["unknown"] +class DisabledStoreSettings(BaseModel): + store_type: Literal["none"] = "none" + + +class DisabledHomeAssistantSettings(BaseModel): + hass_type: Literal["none"] = "none" + + +class SupervisedHomeAssistantSettings(BaseModel): + hass_type: Literal["supervised"] + + supervisor_token: SecretStr + + +class RemoteHomeAssistantSettings(BaseModel): + hass_type: Literal["remote"] + + host: str + port: int + long_lived_access_token: SecretStr + + +def has_dapr() -> bool: + # TODO: make the request to ensure dapr is running + # This is really more a hedge because dapr is new + # and I might hit something I can't work around + # (but so far the workarounds have... worked around) + return os.environ.get("DAPR_MODE") != "false" class AppSettings(BaseSettings): # double underscore, in case your font doesn't make it clear model_config = SettingsConfigDict(env_nested_delimiter="__") - store: Union[PostgresStoreSettings, UnknownStoreSettings] = Field( - discriminator="store_type" + store: Union[PostgresStoreSettings, DisabledStoreSettings] = Field( + discriminator="store_type", + default=DisabledStoreSettings(), ) - events: Union[MqttEventsSettings, UnknownEventsSettings] = Field( - discriminator="events_type" + events: Union[MqttEventsSettings, DisabledEventsSettings] = Field( + discriminator="events_type", + default=DisabledEventsSettings(), ) + hass: Union[ + DisabledHomeAssistantSettings, + SupervisedHomeAssistantSettings, + RemoteHomeAssistantSettings, + ] = Field(discriminator="hass_type", default=DisabledHomeAssistantSettings()) + # TODO: move this into the gateway or something openai_api_key: SecretStr force_publish_models: bool = False notify_fd: Optional[int] = None include_challenger_models: bool = True mlflow_tracking_uri: Optional[str] = None + + +@lru_cache +def get_settings(**kwargs): + if has_dapr(): + from dapr.clients import DaprClient # for typing + + with DaprClient() as dapr_client: + secret_response = dapr_client.get_secret(SECRET_STORE, CONFIGURATION_KEY) + print(secret_response) + return AppSettings.model_validate_json( + secret_response.secret[CONFIGURATION_KEY] + ) + # env vars can populate the settings + return AppSettings() # pyright: ignore diff --git a/python/src/mindctrl/const.py b/python/src/mindctrl/const.py index 0a7b8d7..9c17bf3 100644 --- a/python/src/mindctrl/const.py +++ b/python/src/mindctrl/const.py @@ -1,6 +1,5 @@ from pathlib import Path - CHAMPION_ALIAS = "champion" CHALLENGER_ALIAS = "challenger" @@ -15,3 +14,13 @@ ## Computed BASE_DIR = Path(__file__).parent.resolve() TEMPLATES_DIR = BASE_DIR / "templates" + + +## Events +STOP_DEPLOYED_MODEL = "stop_deployed_model" + +## Config +CONFIGURATION_STORE = "configstore" +SECRET_STORE = "secretstore" +CONFIGURATION_KEY = "mindctrl.appsettings" +CONFIGURATION_TABLE = "mindctrlconfig" diff --git a/python/src/mindctrl/contracts/websocket.py b/python/src/mindctrl/contracts/websocket.py new file mode 100644 index 0000000..0949cd0 --- /dev/null +++ b/python/src/mindctrl/contracts/websocket.py @@ -0,0 +1,33 @@ +from typing import Union + +from pydantic import BaseModel, Field + + +class BaseMindctrlMessage(BaseModel): + id: int + type: str + + +# This is the UI contract, drop tool calling +class AssistantMessage(BaseModel): + role: str = "assistant" + content: str + + +class UserMessage(BaseModel): + role: str = "user" + content: str + + +class ChatMessage(BaseMindctrlMessage): + type: str = "mindctrl.chat" + message: Union[AssistantMessage, UserMessage] = Field(discriminator="role") + + +class SubscribeMessage(BaseMindctrlMessage): + type: str = "mindctrl.subscribe" + subscription: str + + +class WebSocketMessage(BaseMindctrlMessage): + message: Union[ChatMessage, SubscribeMessage] diff --git a/python/src/mindctrl/db/setup.py b/python/src/mindctrl/db/setup.py index 2baa6fc..e9e1b46 100644 --- a/python/src/mindctrl/db/setup.py +++ b/python/src/mindctrl/db/setup.py @@ -4,21 +4,19 @@ from sqlalchemy import text from sqlalchemy.ext.asyncio import ( - create_async_engine, AsyncEngine, + create_async_engine, ) +from ..config import PostgresStoreSettings +from ..mlmodels import summarize_events from .queries import ( + CONVERT_TO_HYPERTABLE, CREATE_SUMMARY_TABLE, ENABLE_PGVECTOR, - CONVERT_TO_HYPERTABLE, ENABLE_TIMESCALE, ) -from ..config import PostgresStoreSettings -from ..mlmodels import summarize_events - - _LOGGER = logging.getLogger(__name__) @@ -100,3 +98,9 @@ async def insert_summary( ) await conn.commit() _LOGGER.info(f"Inserted summary with {len(state_ring_buffer)} events") + + +async def insert_summary_dummy( + state_ring_buffer: collections.deque[dict], +): + _LOGGER.info(f"Inserting null summary for {len(state_ring_buffer)} events") diff --git a/python/src/mindctrl/homeassistant/client.py b/python/src/mindctrl/homeassistant/client.py index 621221a..0a4b31a 100644 --- a/python/src/mindctrl/homeassistant/client.py +++ b/python/src/mindctrl/homeassistant/client.py @@ -1,7 +1,9 @@ import asyncio import logging -from typing import Union +import os import time +from typing import Union + import httpx from httpx_ws import aconnect_ws from pydantic import ValidationError @@ -16,6 +18,7 @@ CreateAutomation, CreateLabel, Error, + ExecuteScript, Label, LabelsResult, ListAreas, @@ -23,6 +26,7 @@ ListLabels, ManyResponsesWrapper, Result, + ServiceCall, SingleResponseWrapper, UpdateEntityLabels, ) @@ -59,7 +63,7 @@ def __init__(self, id: str, hass_url: httpx.URL, token: str): @property def authenticated_session(self): if not self._authenticated_session: - raise ValueError("Session not authenticated") + raise ValueError("Session not authenticated or not started (enter context)") return self._authenticated_session async def __aenter__(self): @@ -150,9 +154,7 @@ async def list_automations(self) -> list[Automation]: entities = await self.list_entities() _logger.debug(entities) automation_entities = [ - entity - for entity in entities - if entity["platform"] == "automation" + entity for entity in entities if entity["platform"] == "automation" ] _logger.info(f"Fetching {len(automation_entities)} automations") @@ -164,12 +166,12 @@ async def list_automations(self) -> list[Automation]: async def list_labels(self): any_result = await self._send_message(ListLabels(id=-1)) labels = LabelsResult.model_validate_json(any_result.model_dump_json()) - return labels.result + return labels async def list_areas(self): any_result = await self._send_message(ListAreas(id=-1)) areas = AreasResult.model_validate_json(any_result.model_dump_json()) - return areas.result + return areas async def create_label(self, label: Label): await self._send_message( @@ -227,3 +229,84 @@ async def create_automation(self, name: str, description: str): response.raise_for_status() # response.json() is just {'result': 'ok'}, need to do a get (why?) return await self.get_automation(str(current_milli_time)) + + @staticmethod + def _generate_service_call(service: str, target: dict): + return ExecuteScript( + id=-1, sequence=[ServiceCall(service=service, target=target)] + ) + + async def _send_service_call(self, service: str, target: dict): + message = HassClient._generate_service_call(service, target) + resp = await self._send_message(message) + if not resp.success: + raise HassClientError(f"Error: {resp}") + + async def light_toggle(self, area_id: str): + return await self._send_service_call("light.toggle", {"area_id": [area_id]}) + + async def light_turn_on(self, area_id: str): + return await self._send_service_call("light.turn_on", {"area_id": [area_id]}) + + # TODO: After evaluating the prompting, see if a mixin approach would be better without params + # For example class Area(Targetable, Lightable, Switchable, Sonosable, etc.): + # This might be better than adding every targetting mechanism to each service? + async def light_turn_off(self, area_id: str): + return await self._send_service_call("light.turn_off", {"area_id": [area_id]}) + + +def hass_client_from_dapr(): + from dapr.clients import DaprClient + + DaprClient().get_secret("hass", "token") + raise NotImplementedError("Not implemented") + + +def hass_client_from_env(id: str = ""): + # TODO: Move these to constants + url = os.environ["HASS_SERVER"] + token = os.environ["HASS_TOKEN"] + # print(f"TOKEN: {token[:10]}") + _logger.info(f"Connecting to Home Assistant at {url}") + return HassClient(id, httpx.URL(f"{url}/api"), token) + + +async def run_api(client, func, *args, **kwargs): + async with client: + return await func(*args, **kwargs) + + +def list_areas(): + """List all areas(rooms) in the home with their area_id and friendly name.""" + client = hass_client_from_env() + return asyncio.run(run_api(client, client.list_areas)) + + +def light_turn_on(area_id: str): + """Turn on the light in the area""" + client = hass_client_from_env() + return asyncio.run(run_api(client, client.light_turn_on, area_id)) + + +def light_turn_off(area_id: str): + """Turn off the light in the area""" + client = hass_client_from_env() + return asyncio.run(run_api(client, client.light_turn_off, area_id)) + + +def light_toggle(area_id: str): + """Toggle the light in the area""" + client = hass_client_from_env() + return asyncio.run(run_api(client, client.light_toggle, area_id)) + + +# TODO: This is going to grow past context limits +# Need to run the intent query on phi/local/classifier +# Or maybe embed every domain and embed the incoming chat +# and add the top domain to the context +TOOL_MAP = { + "list_areas": list_areas, + "light_turn_on": light_turn_on, + "light_turn_off": light_turn_off, + "light_toggle": light_toggle, +} diff --git a/python/src/mindctrl/homeassistant/messages.py b/python/src/mindctrl/homeassistant/messages.py index f966deb..eb7092a 100644 --- a/python/src/mindctrl/homeassistant/messages.py +++ b/python/src/mindctrl/homeassistant/messages.py @@ -1,4 +1,5 @@ from typing import Any, Optional, Union + from pydantic import BaseModel @@ -61,6 +62,17 @@ class ListAreas(Command): type: str = "config/area_registry/list" +class ServiceCall(BaseModel): + service: str + data: Optional[Any] = {} + target: Optional[dict[str, list[str]]] + + +class ExecuteScript(Command): + type: str = "execute_script" + sequence: list[ServiceCall] + + # {"color":"indigo","description":null,"icon":"mdi:account","label_id":"test","name":"test"} class Label(BaseModel): color: str diff --git a/python/src/mindctrl/main.py b/python/src/mindctrl/main.py index e6a48d0..0537e61 100644 --- a/python/src/mindctrl/main.py +++ b/python/src/mindctrl/main.py @@ -1,29 +1,28 @@ -from functools import lru_cache, partial +import asyncio +import collections import logging import os # Eventing - move this to plugin from contextlib import asynccontextmanager -import asyncio +from functools import partial + +import mlflow # Core functionality from fastapi import FastAPI, Request from fastapi.responses import HTMLResponse -import collections - -import mlflow +from mindctrl.workflows import WorkflowContext - -from .mlmodels import log_system_models -from .mqtt import setup_mqtt_client, listen_to_mqtt +from .config import AppSettings, get_settings +from .const import ROUTE_PREFIX +from .db.setup import insert_summary, insert_summary_dummy, setup_db from .mlflow_bridge import connect_to_mlflow, poll_registry -from .db.setup import setup_db, insert_summary -from .config import AppSettings +from .mlmodels import log_system_models +from .mqtt import listen_to_mqtt, setup_mqtt_client from .routers import deployed_models, info, ui from .routers.ui import templates -from .const import ROUTE_PREFIX - _logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) @@ -37,64 +36,65 @@ def write_healthcheck_file(settings: AppSettings): os.close(int(notification_fd)) -@lru_cache -def get_settings(): - # env vars can populate the settings - return AppSettings() # pyright: ignore - - @asynccontextmanager async def lifespan(app: FastAPI): app_settings = get_settings() _logger.info("Starting mindctrl server with settings:") _logger.info(app_settings.model_dump()) - asyncio.create_task(poll_registry(10.0)) - # The buffer should be enhanced to be token-aware state_ring_buffer: collections.deque[dict] = collections.deque(maxlen=20) _logger.info("Setting up DB") # TODO: convert to ABC with a common interface - if not app_settings.store.store_type == "psql": - raise ValueError(f"unknown store type: {app_settings.store.store_type}") - engine = await setup_db(app_settings.store) - insert_summary_partial = partial( - insert_summary, engine, app_settings.include_challenger_models - ) + if app_settings.store.store_type == "psql": + engine = await setup_db(app_settings.store) + insert_summary_partial = partial( + insert_summary, engine, app_settings.include_challenger_models + ) + if app_settings.store.store_type == "none": + insert_summary_partial = insert_summary_dummy _logger.info("Setting up MQTT") - if not app_settings.events.events_type == "mqtt": - raise ValueError(f"unknown events type: {app_settings.events.events_type}") - - mqtt_client = setup_mqtt_client(app_settings.events) - loop = asyncio.get_event_loop() - _logger.info("Starting MQTT listener") - mqtt_listener_task = loop.create_task( - listen_to_mqtt(mqtt_client, state_ring_buffer, insert_summary_partial) - ) + mqtt_listener_task = None + if app_settings.events.events_type == "mqtt": + mqtt_client = setup_mqtt_client(app_settings.events) + loop = asyncio.get_event_loop() + _logger.info("Starting MQTT listener") + mqtt_listener_task = loop.create_task( + listen_to_mqtt(mqtt_client, state_ring_buffer, insert_summary_partial) + ) _logger.info("Logging models") loaded_models = log_system_models(app_settings.force_publish_models) connect_to_mlflow(app_settings) - write_healthcheck_file(app_settings) - - _logger.info("Finished server setup") - # Make resources available to requests via .state - yield { - "state_ring_buffer": state_ring_buffer, - "loaded_models": loaded_models, - "database_engine": engine, - } - - # Cancel the task - mqtt_listener_task.cancel() - # Wait for the task to be cancelled - try: - await mqtt_listener_task - except asyncio.CancelledError: - pass - await engine.dispose() + _logger.info("Starting workflow manager") + with WorkflowContext() as wfc: + asyncio.create_task(poll_registry(wfc, 10.0)) + + write_healthcheck_file(app_settings) + + _logger.info("Finished server setup") + # Make resources available to requests via .state + yield { + "state_ring_buffer": state_ring_buffer, + "loaded_models": loaded_models, + "database_engine": engine, + "workflow_context": wfc, + } + + # TODO: Once the above is moved into an ABC make it a context manager + if app_settings.events.events_type == "mqtt" and mqtt_listener_task: + # Cancel the task + mqtt_listener_task.cancel() + # Wait for the task to be cancelled + try: + await mqtt_listener_task + except asyncio.CancelledError: + pass + + if app_settings.store.store_type == "psql": + await engine.dispose() app = FastAPI(lifespan=lifespan) diff --git a/python/src/mindctrl/mlflow_bridge.py b/python/src/mindctrl/mlflow_bridge.py index ca95573..4c6dd20 100644 --- a/python/src/mindctrl/mlflow_bridge.py +++ b/python/src/mindctrl/mlflow_bridge.py @@ -1,11 +1,13 @@ +import asyncio +import logging + import mlflow from mlflow import MlflowClient -import logging -import asyncio -from .const import CHAMPION_ALIAS, CHALLENGER_ALIAS -from .config import AppSettings +from mindctrl.workflows import WorkflowContext +from .config import AppSettings +from .const import CHALLENGER_ALIAS, CHAMPION_ALIAS _logger = logging.getLogger(__name__) @@ -27,7 +29,8 @@ def is_deployable_alias(aliases: list[str]) -> bool: # TODO: Add webhooks/eventing to MLflow OSS server. AzureML has eventgrid support # In its absence, we poll the MLflow server for changes to the model registry -async def poll_registry(delay_seconds: float = 10.0): +async def poll_registry(workflow_context: WorkflowContext, delay_seconds: float = 10.0): + # TODO: Turn this whole thing into a timer based workflow? while True: # Sync any new models by tag/label/all # Solve any environment dependencies mismatch or fail diff --git a/python/src/mindctrl/mlmodels.py b/python/src/mindctrl/mlmodels.py index fbc18bc..626e642 100644 --- a/python/src/mindctrl/mlmodels.py +++ b/python/src/mindctrl/mlmodels.py @@ -1,14 +1,14 @@ import logging -from typing import Tuple +from typing import Optional, Tuple import mlflow import openai -from mlflow.entities.model_registry import RegisteredModel from mlflow import MlflowClient +from mlflow.entities.model_registry import RegisteredModel +from mlflow.utils.proto_json_utils import dataframe_from_parsed_json - +from .const import CHALLENGER_ALIAS, CHAMPION_ALIAS, SCENARIO_NAME_PARAM from .openai_deployment import log_model -from .const import CHALLENGER_ALIAS, CHAMPION_ALIAS _logger = logging.getLogger(__name__) @@ -186,3 +186,31 @@ def embed_summary(summary: str) -> list[float]: model = mlflow.sentence_transformers.load_model("models:/localembeddings/latest") # return model.predict(summary) return model.encode(summary).tolist() + + +def invoke_model_impl( + model: mlflow.pyfunc.PyFuncModel, + payload: dict, + scenario_name: Optional[str], + input_variables: dict[str, str], +): + # TODO: need a better api for this from mlflow. + # predict() has expensive side effects so shouldn't simply catch invalid_params + model_has_params = hasattr(model.metadata, "get_params_schema") + params = None + if scenario_name: + _logger.info(f"Scenario: {scenario_name}") + if not model_has_params: + _logger.warning( + f"Model {model.metadata} does not have params schema, ignoring scenario header" + ) + else: + _logger.info( + f"Model has params schema: {model.metadata.get_params_schema()}" + ) + params = {SCENARIO_NAME_PARAM: scenario_name} + input = dataframe_from_parsed_json(payload["dataframe_split"], "split") + for key, value in input_variables.items(): + _logger.debug(f"Setting input variable {key}") + input[key] = value + return model.predict(input, params=params) diff --git a/python/src/mindctrl/mqtt.py b/python/src/mindctrl/mqtt.py index 8541689..75891b3 100644 --- a/python/src/mindctrl/mqtt.py +++ b/python/src/mindctrl/mqtt.py @@ -1,13 +1,13 @@ +import asyncio import collections import json import logging -from typing import Callable, Awaitable, Optional +from typing import Awaitable, Callable, Optional + import aiomqtt -import asyncio from .config import MqttEventsSettings - _logger = logging.getLogger(__name__) diff --git a/python/src/mindctrl/openai_deployment/__init__.py b/python/src/mindctrl/openai_deployment/__init__.py index 829b267..f6fca4d 100644 --- a/python/src/mindctrl/openai_deployment/__init__.py +++ b/python/src/mindctrl/openai_deployment/__init__.py @@ -219,7 +219,7 @@ def _log_secrets_yaml(local_model_dir, scope): def _parse_format_fields(s) -> Set[str]: """Parses format fields from a given string, e.g. "Hello {name}" -> ["name"].""" - return {fn for _, fn, _, _ in Formatter().parse(s) if fn is not None} + return {fn for _, fn, _, _ in Formatter().parse(s or "") if fn is not None} def _get_input_schema(task, content): @@ -557,22 +557,12 @@ def _load_model(path): def _is_valid_message(d): - return isinstance(d, dict) and "content" in d and "role" in d + return isinstance(d, dict) and "role" in d and ("content" in d or "tool_calls" in d) class _ContentFormatter: def __init__(self, task, template=None): - if task == "completions": - template = template or "{prompt}" - if not isinstance(template, str): - raise mlflow.MlflowException.invalid_parameter_value( - f"Template for task {task} expects type `str`, but got {type(template)}." - ) - - self.template = template - self.format_fn = self.format_prompt - self.variables = sorted(_parse_format_fields(self.template)) - elif task == "chat.completions": + if task == "chat.completions": if not template: template = [{"role": "user", "content": "{content}"}] if not all(map(_is_valid_message, template)): @@ -583,18 +573,31 @@ def __init__(self, task, template=None): self.template = template.copy() self.format_fn = self.format_chat - self.variables = sorted( - set( - itertools.chain.from_iterable( - _parse_format_fields(message.get("content")) - | _parse_format_fields(message.get("role")) - for message in self.template - ) - ) - ) + print("SKIPPING WEIRD AUTOPARSE") + self.variables = sorted(set()) + # self.variables = sorted( + # set( + # itertools.chain.from_iterable( + # _parse_format_fields(message.get("content", "")) + # | _parse_format_fields(message.get("role")) + # | _parse_format_fields(message.get("tool_call_id", "")) + # | _parse_format_fields(message.get("name", "")) + # for message in self.template + # ) + # ) + # ) if not self.variables: - self.template.append({"role": "user", "content": "{content}"}) + self.template.append({ + "role": "{role}", + "content": "{content}", + "name": "{name}", + "tool_call_id": "{tool_call_id}" + }) self.variables.append("content") + self.variables.append("role") + self.variables.append("name") + self.variables.append("tool_call_id") + else: raise mlflow.MlflowException.invalid_parameter_value( f"Task type ``{task}`` is not supported for formatting." @@ -613,13 +616,20 @@ def format_prompt(self, **params): def format_chat(self, **params): format_args = {v: params[v] for v in self.variables} - return [ - { - "role": message.get("role").format(**format_args), - "content": message.get("content").format(**format_args), - } - for message in self.template - ] + result = [] + for index, message in enumerate(self.template): + # Only do the templating on the first or last messages + system message + if index < 2 or index == len(self.template) - 1: + message["role"] = message.get("role").format(**format_args) + if "content" in message and message.get("content"): + message["content"] = message.get("content").format(**format_args) + if "name" in message and message.get("name"): + message["name"] = message.get("name").format(**format_args) + if "tool_call_id" in message and message.get("tool_call_id"): + message["tool_call_id"] = message.get("tool_call_id").format(**format_args) + + result.append(message) + return result def _first_string_column(pdf): @@ -650,7 +660,7 @@ def _setup_completions(self): if self.task == "chat.completions": self.template = self.model.get("messages", []) else: - self.template = self.model.get("prompt") + raise ValueError(f"Unsupported task: {self.task}") self.formater = _ContentFormatter(self.task, self.template) def format_completions(self, params_list): @@ -743,12 +753,14 @@ def _predict_chat(self, data, params: dict): responses = [] for r in requests: + print("SENDING AI REQUEST", r) response = deploy_client.predict( endpoint=matched_endpoint.name, inputs=r, ) responses.append(response) + # TODO: Better to return a more complex object (tuple) that says it was a tool call result = [] for r in responses: if r["choices"][0]["finish_reason"]== "tool_calls": diff --git a/python/src/mindctrl/rag.py b/python/src/mindctrl/rag.py index 48450f5..c79bfa7 100644 --- a/python/src/mindctrl/rag.py +++ b/python/src/mindctrl/rag.py @@ -1,10 +1,11 @@ +import json from abc import ABC, abstractmethod from datetime import datetime from enum import Enum +from itertools import islice + from fastapi import Request -import json from pydantic import BaseModel -from itertools import islice class EventType(Enum): diff --git a/python/src/mindctrl/replay_server.py b/python/src/mindctrl/replay_server.py index 9b1c61b..5abbbf1 100644 --- a/python/src/mindctrl/replay_server.py +++ b/python/src/mindctrl/replay_server.py @@ -3,41 +3,43 @@ # TODO: mainline this feature into mlflow and drop this file import functools -from fastapi import Request as fastRequest, HTTPException import logging import os import subprocess import sys from typing import List, Literal, Optional, Union + +from pydantic import Field import vcr +import vcr.stubs.aiohttp_stubs +from aiohttp import hdrs +from fastapi import HTTPException +from fastapi import Request as fastRequest ## MLflow Patching from mlflow.deployments.server.app import GatewayAPI, create_app_from_path -from mlflow.environment_variables import MLFLOW_DEPLOYMENTS_CONFIG from mlflow.deployments.server.runner import Runner +from mlflow.environment_variables import MLFLOW_DEPLOYMENTS_CONFIG +from mlflow.exceptions import MlflowException from mlflow.gateway.config import RouteConfig from mlflow.gateway.providers import get_provider from mlflow.gateway.schemas import chat from mlflow.gateway.utils import make_streaming_response -from mlflow.exceptions import MlflowException -## - -## VCR Patching -from yarl import URL -import vcr.stubs.aiohttp_stubs +from vcr.errors import CannotOverwriteExistingCassetteException +from vcr.request import Request from vcr.stubs.aiohttp_stubs import ( - play_responses, - record_responses, _build_cookie_header, - _serialize_headers, _build_url_with_params, + _serialize_headers, + play_responses, + record_responses, ) -from vcr.errors import CannotOverwriteExistingCassetteException -from vcr.request import Request -from aiohttp import hdrs ## +## VCR Patching +from yarl import URL +## from .const import ( REPLAY_SERVER_INPUT_FILE_SUFFIX, REPLAY_SERVER_OUTPUT_FILE_SUFFIX, @@ -124,9 +126,25 @@ def _create_replay_chat_endpoint(config: RouteConfig): # mctrl_header = {"User-Agent": f"mindctrl/{mindctrl.__version__}"} mctrl_header = {"User-Agent": "mindctrl/0.1.0"} - from mlflow.gateway.base_models import ResponseModel + from mlflow.gateway.base_models import ResponseModel, RequestModel + from mlflow.gateway.schemas.chat import ( + BaseRequestPayload, + _REQUEST_PAYLOAD_EXTRA_SCHEMA, + ) from mlflow.gateway.providers.utils import send_request + class RequestMessage(RequestModel): + role: str + content: Optional[str] = None + tool_call_id: Optional[str] = None + name: Optional[str] = None + + class RequestPayload(BaseRequestPayload): + messages: List[RequestMessage] = Field(..., min_length=1) + + class Config: + json_schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA + class Function(ResponseModel): name: str arguments: str @@ -162,6 +180,7 @@ class ResponsePayload(ResponseModel): async def chat_with_tools(self, payload): from fastapi.encoders import jsonable_encoder + print("AI REQUEST", payload) payload = jsonable_encoder(payload, exclude_none=True) self.check_for_model_field(payload) all_headers = {**self._request_headers, **mctrl_header} @@ -171,7 +190,7 @@ async def chat_with_tools(self, payload): path="chat/completions", payload=self._add_model_to_payload_if_necessary(payload), ) - print(resp) + print("AI RESPONSE", resp) return ResponsePayload( id=resp["id"], @@ -183,7 +202,7 @@ async def chat_with_tools(self, payload): index=idx, message=ResponseMessage( role=c["message"]["role"], - content=c["message"]["content"] or "", + content=c["message"].get("content", ""), tool_calls=c["message"].get("tool_calls"), # type: ignore ), finish_reason=c["finish_reason"], @@ -199,6 +218,8 @@ async def chat_with_tools(self, payload): import mlflow.gateway.schemas.chat + mlflow.gateway.schemas.chat.RequestMessage = RequestMessage + mlflow.gateway.schemas.chat.RequestPayload = RequestPayload mlflow.gateway.schemas.chat.ResponseMessage = ResponseMessage mlflow.gateway.schemas.chat.ResponsePayload = ResponsePayload mlflow.gateway.schemas.chat.Choice = Choice diff --git a/python/src/mindctrl/routers/deployed_models.py b/python/src/mindctrl/routers/deployed_models.py index 825dd3e..77136c3 100644 --- a/python/src/mindctrl/routers/deployed_models.py +++ b/python/src/mindctrl/routers/deployed_models.py @@ -1,11 +1,15 @@ import collections import logging -from fastapi import APIRouter, Request, HTTPException + import mlflow -from mlflow.utils.proto_json_utils import dataframe_from_parsed_json +from fastapi import APIRouter, HTTPException, Request -from mindctrl.const import SCENARIO_NAME_HEADER, SCENARIO_NAME_PARAM -from mindctrl.mlmodels import SUMMARIZATION_PROMPT, SUMMARIZER_OAI_MODEL +from mindctrl.const import SCENARIO_NAME_HEADER +from mindctrl.mlmodels import ( + SUMMARIZATION_PROMPT, + SUMMARIZER_OAI_MODEL, + invoke_model_impl, +) router = APIRouter(prefix="/deployed-models", tags=["deployed_models"]) @@ -30,7 +34,7 @@ # return langchain.run(relevant_events, query) -def generate_state_lines(buffer: collections.deque): +def generate_state_lines(buffer: collections.deque) -> str: # TODO: when I get internet see if RAG framework already has a known technique to deal with context chunking import tiktoken @@ -61,28 +65,16 @@ def generate_state_lines(buffer: collections.deque): return state_lines -def invoke_model_impl( - model: mlflow.pyfunc.PyFuncModel, payload: dict, request: Request -): +def invoke_model(model: mlflow.pyfunc.PyFuncModel, payload: dict, request: Request): scenario_name = request.headers.get(SCENARIO_NAME_HEADER) - # TODO: need a better api for this from mlflow. - # predict() has expensive side effects so shouldn't simply catch invalid_params - model_has_params = hasattr(model.metadata, "get_params_schema") - params = None - if scenario_name: - _logger.info(f"Scenario: {scenario_name}") - if not model_has_params: - _logger.warning( - f"Model {model.metadata} does not have params schema, ignoring scenario header" - ) - else: - _logger.info( - f"Model has params schema: {model.metadata.get_params_schema()}" - ) - params = {SCENARIO_NAME_PARAM: scenario_name} - input = dataframe_from_parsed_json(payload["dataframe_split"], "split") - input["state_lines"] = generate_state_lines(request.state.state_ring_buffer) - return model.predict(input, params=params) + return invoke_model_impl( + model, + payload, + scenario_name=scenario_name, + input_variables={ + "state_lines": generate_state_lines(request.state.state_ring_buffer) + }, + ) # This logic is obviously wrong, stub impl @@ -119,7 +111,7 @@ def invoke_labeled_model_version( raise HTTPException(status_code=500, detail=f"Error loading model: {e}") from e try: - return invoke_model_impl(model, payload, request) + return invoke_model(model, payload, request) except Exception as e: _logger.error(f"Error invoking model: {e}") raise HTTPException(status_code=500, detail=f"Error invoking model: {e}") from e diff --git a/python/src/mindctrl/routers/info.py b/python/src/mindctrl/routers/info.py index 522527e..04e300a 100644 --- a/python/src/mindctrl/routers/info.py +++ b/python/src/mindctrl/routers/info.py @@ -1,6 +1,7 @@ import logging -from fastapi import APIRouter, Request + import mlflow +from fastapi import APIRouter, Request router = APIRouter(tags=["info"]) diff --git a/python/src/mindctrl/routers/ui.py b/python/src/mindctrl/routers/ui.py index 15a2241..1676029 100644 --- a/python/src/mindctrl/routers/ui.py +++ b/python/src/mindctrl/routers/ui.py @@ -4,9 +4,10 @@ from fastapi import APIRouter, WebSocket, WebSocketDisconnect from fastapi.templating import Jinja2Templates - from mindctrl.const import TEMPLATES_DIR +# from mindctrl.workflows import WorkflowContext + router = APIRouter(prefix="/ui", tags=["info"]) @@ -34,3 +35,67 @@ async def websocket_endpoint(websocket: WebSocket): except WebSocketDisconnect: _logger.warning("Websocket disconnected") break + + +# TODO: https://fastapi.tiangolo.com/advanced/websockets/#handling-disconnections-and-multiple-clients +# @router.websocket("/mctrlws") +# async def websocket_endpoint2(websocket: WebSocket): +# await websocket.accept() +# _logger.info("Websocket accepted") + +# await websocket.accept() +# queue = asyncio.queues.Queue() + +# async def read_from_socket(websocket: WebSocket): +# async for data in websocket.iter_json(): +# print(f"putting {data} in the queue") +# queue.put_nowait(data) + +# async def get_data_and_send(): +# data = await queue.get() +# while True: +# if queue.empty(): +# print(f"getting weather data for {data}") +# await asyncio.sleep(1) +# else: +# data = queue.get_nowait() +# print(f"Setting data to {data}") + +# await asyncio.gather(read_from_socket(websocket), get_data_and_send()) + +# # If doesn't exist, starts workflow +# # If exists and is paused, resumes workflow +# workflow_context: WorkflowContext = websocket.state.workflow_context +# session_id = await mindctrl.get_or_create_conversation(client_id) + +# await asyncio.sleep(10) +# ring_buffer = iter(websocket.state.state_ring_buffer.copy()) +# while True: +# try: +# payload = next(ring_buffer) +# await websocket.send_json(payload) +# await asyncio.sleep(1) + +# message = await websocket.receive_json() +# _logger.info(f"Message received: {message}") +# assert message["type"] == "mindctrl.chat.user" +# chat_message = Message(content=message["content"], role="user") + +# # TODO: Actually expand the polling loop, so send should be quick +# # Then poll on the assistant response +# # You might not need a workflow for multiturn? +# assistant_response = await mindctrl.send_message(session_id, chat_message) +# await websocket.send_json(assistant_response) + +# except StopIteration: +# _logger.warning("Websocket buffer empty, waiting for new events") +# await asyncio.sleep(2) +# ring_buffer = iter(websocket.state.state_ring_buffer.copy()) +# except WebSocketDisconnect: +# _logger.warning("Websocket disconnected") + +# # Pauses workflow +# # Sets a timer to terminate the workflow if no activity for X minutes +# mindctrl.disconnect_conversation(session_id) + +# break diff --git a/python/src/mindctrl/tools/functions.py b/python/src/mindctrl/tools/functions.py new file mode 100644 index 0000000..1d84a95 --- /dev/null +++ b/python/src/mindctrl/tools/functions.py @@ -0,0 +1,84 @@ +import inspect +import json +import logging +from typing import Callable, Optional + + +_logger = logging.getLogger(__name__) + +_specific_json_types = { + str: "string", + int: "number", +} + + +def generate_json_schema( + name: str, description: Optional[str], **params: inspect.Parameter +) -> dict: + schema = { + "type": "function", + "function": { + "name": name, + "description": description, + "parameters": {"type": "object", "properties": {}, "required": []}, + }, + } + + for param_name, param in params.items(): + param_schema = { + "type": _specific_json_types.get(param.annotation, "object"), + "description": param.default + if param.default != inspect.Parameter.empty + else None, + } + + if param.default == inspect.Parameter.empty: + schema["function"]["parameters"]["required"].append(param_name) + + schema["function"]["parameters"]["properties"][param_name] = param_schema + + return schema + + +async def call_function(func): + ret = func() + return await ret if inspect.isawaitable(ret) else ret + + +def generate_function_schema(func: Callable) -> dict: + signature = inspect.signature(func) + parameters = signature.parameters + + schema = generate_json_schema(func.__name__, func.__doc__, **parameters) + _logger.debug(f"Generated schema: {schema}") + return schema + + +if __name__ == "__main__": + # Example usage + def get_current_weather(location: str, unit: str = "celsius"): + """ + Get the current weather in a given location + + Args: + location (str): The city and state, e.g. San Francisco, CA + unit (str, optional): The unit of temperature. Defaults to "celsius". + + Returns: + dict: The weather information + """ + pass + + json_schema = generate_function_schema(get_current_weather) + print(json.dumps(json_schema, indent=4)) + + class FakeHass: + def get_weather(self, location: str, unit: str = "celsius"): + pass + + async def get_weather_async(self, location: str, unit: str = "celsius"): + pass + + hass = FakeHass() + json_schema = generate_function_schema(hass.get_weather_async) + print(json.dumps(json_schema, indent=4)) diff --git a/python/src/mindctrl/workflows/__init__.py b/python/src/mindctrl/workflows/__init__.py new file mode 100644 index 0000000..3460622 --- /dev/null +++ b/python/src/mindctrl/workflows/__init__.py @@ -0,0 +1,160 @@ +# https://github.com/dapr/python-sdk/blob/main/examples/demo_workflow/app.py#LL40C1-L43C59 + +import json +import logging +import uuid +from typing import Callable, Optional + +import mlflow.pyfunc +from dapr.ext.workflow import WorkflowRuntime +from dapr.ext.workflow.dapr_workflow_client import DaprWorkflowClient +from dapr.ext.workflow.workflow_state import WorkflowStatus + +from .agent import ( + Message, + ModelInvocation, + append_message, + conversation_turn_workflow, + invoke_model, + invoke_tool, +) +from .deployer import ( + check_deployment_status, + deploy_model_workflow, + serve_model, + stop_model, + stop_model_monitor, + wait_for_model_serve, +) + +_logger = logging.getLogger(__name__) + + +class WorkflowModel(mlflow.pyfunc.PythonModel): + def __init__(self): + self.workflows = [] + self.activities = [] + + def add_workflow(self, workflow): + if len(self.workflows) > 0: + raise ValueError("Only one workflow is supported at this time") + self.workflows.append(workflow) + + def add_activity(self, activity): + self.activities.append(activity) + + def predict(self, context, model_input, params=None): + wfr = WorkflowRuntime() + for workflow in self.workflows: + wfr.register_workflow(workflow) + for activity in self.activities: + wfr.register_activity(activity) + instance_id = uuid.uuid4().hex + client = DaprWorkflowClient() + instance_id = client.schedule_new_workflow( + self.workflows[0], input=model_input, instance_id=instance_id + ) + result = client.wait_for_workflow_completion(instance_id, fetch_payloads=True) + assert result is not None + return result.serialized_output + + +class WorkflowContext: + def __init__(self, host: Optional[str] = None, port: Optional[str] = None): + _logger.info(f"Initializing WorkflowContext with {host}:{port}") + self.host = host + self.port = port + self._workflow_runtime = WorkflowRuntime(host=host, port=port) + try: + self._register_turn_workflow() + self._register_deployer_workflow() + except ValueError as e: + if "already registered" not in str(e): + raise e + _logger.info(f"Already registered turn workflow: {e}") + self.initialized = False + + def __enter__(self): + self._workflow_runtime.start() + self.initialized = True + return self + + def __exit__(self, exc_type, exc_value, traceback): + self._workflow_runtime.shutdown() + + @property + def runtime(self): + return self._workflow_runtime + + def get_workflow_client(self): + return DaprWorkflowClient(self.host, self.port) + + def _register_turn_workflow(self): + _logger.info("Registering turn workflow") + self._workflow_runtime.register_workflow(conversation_turn_workflow) + _logger.info("Registering activities") + self._workflow_runtime.register_activity(append_message) + self._workflow_runtime.register_activity(invoke_model) + self._workflow_runtime.register_activity(invoke_tool) + + def _register_deployer_workflow(self): + _logger.info("Registering deployer workflows") + self._workflow_runtime.register_workflow(deploy_model_workflow) + self._workflow_runtime.register_workflow(check_deployment_status) + _logger.info("Registering activities") + self._workflow_runtime.register_activity(wait_for_model_serve) + self._workflow_runtime.register_activity(serve_model) + self._workflow_runtime.register_activity(stop_model) + self._workflow_runtime.register_activity(stop_model_monitor) + + +class Conversation: + def __init__( + self, + client: DaprWorkflowClient, + model_uri: str, + conversation_id: Optional[str] = None, + ): + self.model_uri = model_uri + self.conversation_id: str = conversation_id or uuid.uuid4().hex + self.client = client + self.turn_ids: list[str] = [] + self.tools: list[Callable] = [] + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + pass + + # TODO: Fix the mlflow model predict API to allow dynamic tool input + # not at log_model time. + def add_tool(self, tool): + self.tools.append(tool) + + def send_message(self, message: str) -> Message: + turn_id = f"{self.conversation_id}-{len(self.turn_ids)}" + self.turn_ids.append(turn_id) + instance_id = self.client.schedule_new_workflow( + conversation_turn_workflow, + input=ModelInvocation( + model_uri=self.model_uri, + conversation_id=self.conversation_id, + scenario_name=self.conversation_id, + history=None, + input_variables={}, + ), + instance_id=turn_id, + ) + self.client.raise_workflow_event(instance_id, "user_message", data=message) + + state = self.client.wait_for_workflow_completion(instance_id) + if state is None: + raise RuntimeError(f"Failed to get state for {turn_id}") + if state.runtime_status != WorkflowStatus.COMPLETED: + # print(state._WorkflowState__obj) + # print(state) + raise RuntimeError(f"Failed to complete {turn_id}: {state}") + # print(state.serialized_output) + out = json.loads(state.serialized_output) + return Message(content=out["content"], role=out["role"]) diff --git a/python/src/mindctrl/workflows/agent.py b/python/src/mindctrl/workflows/agent.py new file mode 100644 index 0000000..3a90f66 --- /dev/null +++ b/python/src/mindctrl/workflows/agent.py @@ -0,0 +1,314 @@ +import json +import logging +import pickle +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Generic, Optional, TypeVar + +import mlflow.pyfunc +from dapr.clients import DaprClient +from dapr.ext.workflow import ( + DaprWorkflowContext, + WorkflowActivityContext, +) +from pydantic import BaseModel + +from mindctrl.mlmodels import invoke_model_impl +from mindctrl.openai_deployment import _ContentFormatter, _OpenAIDeploymentWrapper + +_logger = logging.getLogger(__name__) + + +# TODO: THIS IS SOOOOOO BADDDDD +# agent <-> openai_deployment contract +# openai_deployment <-> replay server contract +# replay server <-> openai contract +# And all are slightly different :( +@dataclass +class Function: + name: str + arguments: str + + +@dataclass +class FunctionCall: + id: str + function: Function + type: str = "function" + + +# TODO: unify with the monkeypatch of deploymentserver for tool calling? +# TODO: Fix durable task client to handle pydantic models, then switch +# TODO: file a bug on durabletask.internal.shared.py +@dataclass +class Message: + content: Optional[str] + role: str + tool_call_id: Optional[str] = None + name: Optional[str] = None + tool_calls: Optional[list[FunctionCall]] = None + + +@dataclass +class Conversation: + messages: list[Message] + + +@dataclass +class MessageInConvo: + message: Message + conversation_id: str + + +@dataclass +class ModelInvocation: + model_uri: str + scenario_name: Optional[str] + input_variables: dict[str, str] + conversation_id: str + history: Optional[Conversation] + + +def append_message( + ctx: WorkflowActivityContext, message: MessageInConvo +) -> Conversation: + _logger.info(f"Received message: {message}") + try: + # This is where we do some fun memory tricks like compression, embedding, windowing etc + # TODO: Dapr + async def activities? + with DaprClient() as d: + store_name = "daprstore" + convo_id = f"convo-{message.conversation_id}" + current_convo = d.get_state(store_name=store_name, key=convo_id) + + # TODO: Handle etags + convo = Conversation(messages=[]) + if current_convo.data: + convo: Conversation = pickle.loads(current_convo.data) + + # The type marshaling by durabletask/dapr is a bit wonky... + if isinstance(message.message, dict): + # if + convo.messages.append(Message(**message.message)) + elif isinstance(message.message, Message): + convo.messages.append(message.message) + elif isinstance(message.message, SimpleNamespace): + convo.messages.append(Message(**message.message.__dict__)) + else: + raise ValueError(f"Unknown message type: {type(message.message)}") + d.save_state(store_name, convo_id, pickle.dumps(convo)) + + # TODO: Put the compressed state into the system message next time + return convo + except Exception as e: + import traceback + + traceback.print_exception(e) + # breakpoint() + raise + + +def get_user_chat_payload(query: str) -> dict: + return { + "dataframe_split": { + "columns": ["content", "role", "tool_call_id", "name"], + "data": [[query, "user", None, None]], + } + } + + +def get_tool_call_payload(tool_call_id: str, content: str, name: str) -> dict: + return { + "dataframe_split": { + "columns": ["content", "role", "tool_call_id", "name"], + "data": [[content, "tool", tool_call_id, name]], + } + } + + +def append_history(model: mlflow.pyfunc.PyFuncModel, history: Conversation): + inner_model = model._model_impl + assert isinstance(inner_model, _OpenAIDeploymentWrapper) + messages = [m.__dict__ for m in history.messages] + print("MESSAGES", messages) + print("EXISTING MESSAGES", inner_model.template) + existing_messages: list = inner_model.template # type: ignore + combined_messages = existing_messages[:1] + messages + existing_messages[1:] + inner_model.formater = _ContentFormatter(inner_model.task, combined_messages) + return model + + +# TODO: CONVERT THIS TO MESSAGE IN, MESSAGE OUT +def invoke_model(ctx: WorkflowActivityContext, input: ModelInvocation) -> Message: + # TODO: Handle more complex responses + print(f"Invoking model with input: {input}") + try: + model = mlflow.pyfunc.load_model(input.model_uri) + assert input.history is not None + assert len(input.history.messages) > 0 + history = Conversation(messages=[Message(**m) for m in input.history.messages]) # type: ignore + current_message = history.messages[-1] + if current_message.tool_call_id is not None: + assert ( + current_message.content is not None + ), f"Content is None: {current_message}" + assert current_message.name is not None, f"Name is None: {current_message}" + payload = get_tool_call_payload( + current_message.tool_call_id, + current_message.content, + current_message.name, + ) + else: + assert current_message.content is not None + payload = get_user_chat_payload(current_message.content) + + print(f"PAYLOAD: {payload}") + del history.messages[-1] + print(f"HISTORY: {history}") + model = append_history(model, history) + + response = invoke_model_impl( + model, payload, input.scenario_name, input.input_variables + ) + print(response) + response_message = response[0] + print(response_message) + # TODO: return a better object from openai_deployment.predict() + if isinstance(response_message, list): + # TODO: Support parallel tool calling + function_call: dict = response_message[0] + if function_call.get("type", "unknown") != "function": + raise ValueError(f"Unknown response type: {function_call}") + _logger.info(f"Received function call: {function_call}") + # TODO: Why am I writing all this logic again? + return Message( + content=None, + tool_calls=[ + FunctionCall( + id=function_call["id"], + function=Function( + name=function_call["function"]["name"], + arguments=function_call["function"]["arguments"], + ), + ) + ], + role="assistant", + ) + return Message(content=response_message, role="assistant") + except Exception as e: + import traceback + + traceback.print_exception(e) + # breakpoint() + raise + + +# TODO: Convert the common try/except pattern into a mindctrl decorator +def invoke_tool(ctx: WorkflowActivityContext, function_call: dict) -> str: + from mindctrl.homeassistant.client import TOOL_MAP + + try: + print(f"Invoking tool: {function_call}") + func = TOOL_MAP.get(function_call["function"]["name"]) + if func is None: + raise ValueError( + f"Unknown tool: {function_call['function']['name']}, have {TOOL_MAP.keys()}" + ) + params = json.loads(function_call["function"]["arguments"]) + _logger.info(f"Calling tool: {function_call['function']['name']} with {params}") + tool_result = func(**params) + _logger.info(f"Tool result: {tool_result}") + if isinstance(tool_result, BaseModel): + return tool_result.model_dump_json() + return json.dumps(tool_result) + except Exception as e: + import traceback + + traceback.print_exception(e) + # breakpoint() + raise + + +# def tool_turn_workflow(ctx: DaprWorkflowContext, input: ModelInvocation): +# _logger.info(f"Calling Tool: {input}") +# conversation: Conversation = yield ctx.call_activity(append_message, input=input) +# tool_result = yield ctx.call_activity(invoke_tool, input=input) + + +def conversation_turn_workflow(ctx: DaprWorkflowContext, input: ModelInvocation): + _logger.info(f"Starting Conversation turn: {input}") + + try: + message_str = yield ctx.wait_for_external_event("user_message") + assert isinstance(message_str, str) + # TODO: this is wrong, get a real message structure for chat models + message = Message(content=message_str, role="user") + conversation: Conversation = yield ctx.call_activity( + append_message, + input=MessageInConvo( + message=message, conversation_id=input.conversation_id + ), + ) + input.history = conversation + response_message: Message = yield ctx.call_activity(invoke_model, input=input) + + tool_calling = response_message.tool_calls is not None + while tool_calling: + _logger.info(f"Tool calling: {response_message}") + # TODO Make this a child workflow + assert response_message.tool_calls is not None + conversation: Conversation = yield ctx.call_activity( + append_message, + input=MessageInConvo( + message=response_message, conversation_id=input.conversation_id + ), + ) + print("CALLING TOOL", response_message) + tool_result: str = yield ctx.call_activity( + invoke_tool, + input=response_message.tool_calls[0], # type: ignore + ) + print(f"TOOL RESULT: {tool_result}") + tool_result_message = Message( + role="tool", + content=tool_result, + tool_call_id=response_message.tool_calls[0]["id"], # type: ignore + name=response_message.tool_calls[0]["function"]["name"], # type: ignore + ) + conversation: Conversation = yield ctx.call_activity( + append_message, + input=MessageInConvo( + message=tool_result_message, conversation_id=input.conversation_id + ), + ) + input.history = conversation + response_message: Message = yield ctx.call_activity( + invoke_model, input=input + ) + tool_calling = response_message.tool_calls is not None + + # tool_response = yield ctx.call_child_workflow( + # tool_turn_workflow, + # input=MessageInConvo( + # message=response_message, conversation_id=input.conversation_id + # ), + # ) + # return response_message + + conversation = yield ctx.call_activity( + append_message, + input=MessageInConvo( + message=response_message, conversation_id=input.conversation_id + ), + ) + # TODO: Need to implement https://github.com/microsoft/durabletask-python/issues/25 + # That way the response can be custom status instead of return value + # Right now you have to schedule each turn of the workflow manually + # ctx.continue_as_new(input) + return response_message + except Exception as e: + import traceback + + traceback.print_exception(e) + # breakpoint() + raise diff --git a/python/src/mindctrl/workflows/deployer.py b/python/src/mindctrl/workflows/deployer.py new file mode 100644 index 0000000..f45b29f --- /dev/null +++ b/python/src/mindctrl/workflows/deployer.py @@ -0,0 +1,148 @@ +import logging +import subprocess +from dataclasses import dataclass +from datetime import timedelta + +from dapr.clients import DaprClient +from dapr.ext.workflow import DaprWorkflowContext, RetryPolicy, WorkflowActivityContext +from dapr.ext.workflow.dapr_workflow_client import DaprWorkflowClient + +from mindctrl.const import STOP_DEPLOYED_MODEL + +_logger = logging.getLogger(__name__) + + +def model_uri_to_app_id(model_uri: str) -> str: + return model_uri.replace("/", "_").replace(":", "") + + +@dataclass +class ModelServeCommand: + model_uri: str + port: int + pid: int + is_healthy: bool + app_id: str + + +# TODO: Kubernetes workflow needs kaniko version of this: +# https://github.com/mlflow/mlflow/blob/master/mlflow/models/python_api.py#L79 +def serve_model(ctx: WorkflowActivityContext, model_serve_command: ModelServeCommand): + # This activity serves a model from a local path + if model_serve_command.pid >= 0: + raise ValueError( + f"Model {model_serve_command.model_uri} is already being served by {model_serve_command.pid}" + ) + app_id = model_uri_to_app_id(model_serve_command.model_uri) + _logger.info(f"Starting serving model as dapr app {app_id}") + proc = subprocess.Popen( + [ + "dapr", + "run", + "--app-id", + app_id, + "--app-port", + str(model_serve_command.port), + "--", + "mlflow", + "models", + "serve", + "-m", + model_serve_command.model_uri, + "--port", + str(model_serve_command.port), + "--no-conda", # TODO: Add uv env build + ] + ) + return ModelServeCommand( + model_uri=model_serve_command.model_uri, + pid=proc.pid, + port=model_serve_command.port, + is_healthy=False, + app_id=app_id, + ) + + +def stop_model(ctx: WorkflowActivityContext, model_serve_command: ModelServeCommand): + _logger.info(f"Stopping Model serve {model_serve_command.app_id}") + subprocess.run(["dapr", "stop", "--app-id", model_serve_command.app_id], check=True) + + +def stop_model_monitor(ctx: WorkflowActivityContext, child_workflow_id: str): + _logger.info(f"Stopping monitor: {child_workflow_id}") + wf_client = DaprWorkflowClient() + wf_client.terminate_workflow(child_workflow_id) + + +def wait_for_model_serve( + ctx: WorkflowActivityContext, model_serve_command: ModelServeCommand +) -> ModelServeCommand: + # TODO: Check if the process is still running via Dapr API + # TODO: Store the app id in the model serve command + is_healthy = False + try: + with DaprClient() as d: + resp = d.invoke_method(model_serve_command.app_id, "health") + if resp.status_code != 200: + raise Exception(f"Model serve failed to start: {resp.text}") + is_healthy = True + except Exception as e: + _logger.warning(f"Error checking health: {e}") + + return ModelServeCommand( + model_serve_command.model_uri, + model_serve_command.port, + model_serve_command.pid, + is_healthy=is_healthy, + app_id=model_serve_command.app_id, + ) + + +deployment_retry_policy = RetryPolicy( + first_retry_interval=timedelta(seconds=30), + max_number_of_attempts=3, + backoff_coefficient=2, + max_retry_interval=timedelta(seconds=60), + retry_timeout=timedelta(seconds=180), +) + + +def check_deployment_status( + ctx: DaprWorkflowContext, model_serve_command: ModelServeCommand +): + model_serve_command = yield ctx.call_activity( + wait_for_model_serve, + input=model_serve_command, + retry_policy=deployment_retry_policy, + ) + + check_interval = 60 if model_serve_command.is_healthy else 5 + yield ctx.create_timer(fire_at=timedelta(seconds=check_interval)) + + ctx.continue_as_new(model_serve_command) + + +def deploy_model_workflow( + ctx: DaprWorkflowContext, model_serve_command: ModelServeCommand +): + if not ctx.is_replaying: + _logger.info( + f"Starting model deployment workflow for {model_serve_command.model_uri}" + ) + model_serve_command = yield ctx.call_activity( + serve_model, input=model_serve_command + ) + monitor_id = f"{ctx.instance_id}-monitor" + ctx.call_child_workflow( + check_deployment_status, + input=model_serve_command, + instance_id=monitor_id, + ) + # We want to perform custom termination actions, so don't rely on dapr workflow termination + cancellation_event = yield ctx.wait_for_external_event(STOP_DEPLOYED_MODEL) + _logger.info(f"Received stop event {cancellation_event}") + yield ctx.call_activity(stop_model_monitor, input=monitor_id) + _logger.info("Stopped monitor") + yield ctx.call_activity(stop_model, input=model_serve_command) + + return {"cancellation_reason": cancellation_event} diff --git a/python/tests/test_hass_api.py b/python/tests/test_hass_api.py index 601f29d..b00e637 100644 --- a/python/tests/test_hass_api.py +++ b/python/tests/test_hass_api.py @@ -1,10 +1,10 @@ import logging from random import randint -import pytest + import httpx +import pytest from httpx_ws import aconnect_ws - # https://developers.home-assistant.io/docs/api/websocket # Is there no open source non-GPL python client for Home Assistant? Ideally this exists independently @@ -147,7 +147,7 @@ async def test_list_automations(hass_ws_session): if entity["platform"] == "automation" ] _logger.info(automations) - assert len(automations) >= 0 # yes what a useless assertion + assert len(automations) >= 0 # yes what a useless assertion async def test_list_areas(hass_ws_session): diff --git a/python/tests/test_hass_client.py b/python/tests/test_hass_client.py index 2d88e16..277555f 100644 --- a/python/tests/test_hass_client.py +++ b/python/tests/test_hass_client.py @@ -1,7 +1,10 @@ +import asyncio import logging -from httpx import URL -import pytest +import os +import random +import pytest +from httpx import URL from mindctrl.homeassistant.client import HassClient from mindctrl.homeassistant.messages import CreateLabel @@ -9,12 +12,20 @@ @pytest.fixture -async def hass_client(hass_server_and_token): - server, token = hass_server_and_token +async def hass_client(request): + # TODO: if you run just this file, you get into an asyncio loop issue for playwright + # moving the env var overrides into the fixture itself and using proper fixture resolution + # probably will fix it? + if os.environ.get("HASS_SERVER") is None: + server, token = request.getfixturevalue("hass_server_and_token") + server_url = server.get_base_url() + else: + server_url = os.environ["HASS_SERVER"] + token = os.environ["HASS_TOKEN"] try: async with HassClient( id="pytest", - hass_url=URL(f"{server.get_base_url()}/api"), + hass_url=URL(f"{server_url}/api"), token=token, ) as client: yield client @@ -36,15 +47,33 @@ async def test_mctrl_list_automations(hass_client): async def test_mctrl_list_areas(hass_client): - areas = await hass_client.list_areas() + areas_result = await hass_client.list_areas() + areas = areas_result.result _logger.info(areas) assert len(areas) >= 0 async def test_mctrl_list_labels(hass_client): - labels = await hass_client.list_labels() + labels_result = await hass_client.list_labels() + labels = labels_result.result _logger.info(labels) - assert len(labels) >= 0 + assert labels_result.success + assert len(labels.result) >= 0 + + +async def test_mctrl_control_lights(hass_client): + areas_result = await hass_client.list_areas() + areas = areas_result.result + assert len(areas) >= 1 + random_area = random.choice(areas) + area_id = random_area.area_id + _logger.info(f"Controlling lights for {random_area}") + await hass_client.light_turn_on(area_id) + # TODO: Implement Get State api and assert + await asyncio.sleep(2) + await hass_client.light_turn_off(area_id) + await asyncio.sleep(2) + await hass_client.light_toggle(area_id) async def test_automation_autotag(hass_client, request): @@ -62,7 +91,11 @@ async def test_automation_autotag(hass_client, request): test_label_name = f"{request.node.name}-label" create_new_labels = [ CreateLabel( - id=1, name=test_label_name, color="indigo", icon="mdi:account", description=None + id=1, + name=test_label_name, + color="indigo", + icon="mdi:account", + description=None, ) ] for label in create_new_labels: @@ -73,9 +106,12 @@ async def test_automation_autotag(hass_client, request): # _logger.info(f"Adding label {test_label_name} to {automation.id}") # await hass_client.add_labels(automation.id, [test_label_name]) - await hass_client.add_labels(f"automation.{test_automation_name}", [test_label_name]) + await hass_client.add_labels( + f"automation.{test_automation_name}", [test_label_name] + ) - entities = await hass_client.list_entities() + entities_result = await hass_client.list_entities() + entities = entities_result.result automations = [e for e in entities if e["platform"] == "automation"] _logger.info(automations) tagged_automations = [a for a in automations if test_label_name in a["labels"]] diff --git a/python/tests/test_mlflow.py b/python/tests/test_mlflow.py index 1880226..f41b84c 100644 --- a/python/tests/test_mlflow.py +++ b/python/tests/test_mlflow.py @@ -1,16 +1,16 @@ import json +import uuid + import mlflow import mlflow.openai import openai -import uuid - from mindctrl.const import SCENARIO_NAME_PARAM from mindctrl.mlmodels import log_system_models from mindctrl.openai_deployment import log_model def test_mlflow_setup(mlflow_fluent_session): - assert "sqlite" in mlflow.get_tracking_uri() + assert mlflow.get_tracking_uri() is not None def test_log_system_models(mlflow_fluent_session): diff --git a/python/tests/test_replay_server.py b/python/tests/test_replay_server.py index eec90e8..814cce2 100644 --- a/python/tests/test_replay_server.py +++ b/python/tests/test_replay_server.py @@ -1,13 +1,11 @@ import json import logging from pathlib import Path -import pytest +import pytest from fastapi.testclient import TestClient - from mindctrl.replay_server import create_app_from_env - _logger = logging.getLogger(__name__) diff --git a/python/tests/test_workflows.py b/python/tests/test_workflows.py new file mode 100644 index 0000000..632e80f --- /dev/null +++ b/python/tests/test_workflows.py @@ -0,0 +1,380 @@ +import atexit +import json +import logging +import os +import subprocess +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +import httpx +import openai +import pytest +from dapr.clients import DaprClient +from dapr.conf import settings +from dapr.ext.workflow import WorkflowState +from dapr.ext.workflow.dapr_workflow_client import DaprWorkflowClient +from dapr.ext.workflow.workflow_state import WorkflowStatus +from durabletask.client import OrchestrationState +from mindctrl.homeassistant.client import TOOL_MAP +from mindctrl.openai_deployment import log_model +from mindctrl.tools.functions import generate_function_schema +from mindctrl.workflows import Conversation, WorkflowContext +from mindctrl.workflows.agent import ( + get_user_chat_payload, +) +from mindctrl.workflows.deployer import ModelServeCommand, deploy_model_workflow + +_logger = logging.getLogger(__name__) + + +def stop_dapr_app(app_id: str): + try: + subprocess.run(["dapr", "stop", "-a", app_id], check=True) + except subprocess.CalledProcessError as e: + _logger.error(f"Error stopping Dapr app {app_id}: {e}") + + +def wait_for_input_output( + wf_client: DaprWorkflowClient, + instance_id: str, + target_input: Optional[str] = None, + target_output: Optional[str] = None, + target_input_val: Optional[Any] = None, + target_output_val: Optional[Any] = None, + timeout=120, +): + target_match = False + start_time = time.time() + state = None + + if not target_input and not target_output: + raise ValueError("Either target_input or target_output must be provided") + + if target_input and not target_input_val: + raise ValueError( + "target_input_val must be provided if target_input is provided" + ) + if target_output and not target_output_val: + raise ValueError( + "target_output_val must be provided if target_output is provided" + ) + + while not target_match: + if time.time() - start_time > timeout: + raise TimeoutError( + f"Timed out waiting for {instance_id} to reach target. State:\n" + f"{state._WorkflowState__obj if state else None}" + ) + state = wf_client.get_workflow_state(instance_id, fetch_payloads=True) + assert state is not None + orch_state: OrchestrationState = state._WorkflowState__obj + if target_input and orch_state.serialized_input: + target_match = ( + json.loads(orch_state.serialized_input).get(target_input) + == target_input_val + ) + if target_output and orch_state.serialized_output: + target_match = target_match and ( + json.loads(orch_state.serialized_output).get(target_output) + == target_output_val + ) + status = state.runtime_status + _logger.info( + f"Workflow status: {status}, waiting...\n{state._WorkflowState__obj if state else None}" + ) + time.sleep(5) + + state = wf_client.get_workflow_state(instance_id, fetch_payloads=True) + assert state is not None + return state + + +@pytest.fixture(scope="session") +def placement_server(deploy_mode): + if deploy_mode.value != "local": + # This only makes sense for local testing - dapr is initialized + # in the container/cluster for addon/k8s + _logger.warning(f"Unsupported deploy mode: {deploy_mode}") + pytest.skip(f"Unsupported deploy mode: {deploy_mode}") + + placement_bin = (Path.home() / ".dapr" / "bin" / "placement").resolve() + assert ( + placement_bin.exists() + ), f"placement binary not found at {placement_bin}. Is Dapr installed?" + placement_process = subprocess.Popen([str(placement_bin)]) + + yield placement_process + + placement_process.terminate() + + +@pytest.fixture(scope="session") +def dapr_sidecar( + tmp_path_factory: pytest.TempPathFactory, + repo_root_dir: Path, + request: pytest.FixtureRequest, + deploy_mode, + monkeypatch_session, + placement_server, +): + if deploy_mode.value != "local": + # This only makes sense for local testing - dapr is initialized + # in the container/cluster for addon/k8s + _logger.warning(f"Unsupported deploy mode: {deploy_mode}") + pytest.skip(f"Unsupported deploy mode: {deploy_mode}") + + state_spec = repo_root_dir / "services" / "components" / "sqlite.yaml" + assert state_spec.exists(), f"state store spec not found at {state_spec}" + state_store_path = tmp_path_factory.mktemp("statestore") + target_spec = state_store_path / "sqlite.yaml" + + with monkeypatch_session.context() as m: + m.setenv("ACTOR_STORE_CONNECTION_STRING", f"{state_store_path}/actors.db") + with open(state_spec, "r") as f: + content = f.read() + content = os.path.expandvars(content) + with open(target_spec, "w") as f: + f.write(content) + _logger.info(f"Generated state store spec at {target_spec}") + + dapr_process = subprocess.Popen( + [ + "dapr", + "run", + "--app-id", + request.node.name, + "--dapr-grpc-port", + str(settings.DAPR_GRPC_PORT), + "--dapr-http-port", + str(settings.DAPR_HTTP_PORT), + # "--log-level", + # "debug", + "--resources-path", + f"{state_store_path}", + ] + ) + yield dapr_process + dapr_process.terminate() + + +# # TODO: dedupe with the other one in test_hass_client +# @pytest.fixture(scope="session") +# async def hass_client(request): +# if os.environ.get("HASS_SERVER") is None: +# server, token = request.getfixturevalue("hass_server_and_token") +# server_url = server.get_base_url() +# else: +# server_url = os.environ["HASS_SERVER"] +# token = os.environ["HASS_TOKEN"] +# try: +# async with HassClient( +# id="pytest", +# hass_url=httpx.URL(f"{server_url}/api"), +# token=token, +# ) as client: +# yield client +# except RuntimeError as e: +# # TODO: RuntimeError: Attempted to exit cancel scope in a different task than it was entered in +# # This is probably a ticking timebomb for some event loop bug i don't understand, but fingers crossed +# # it's in pytest + pytest-asyncio and not the actual code (but probably not because httpx-ws is new :/ ) +# if ( +# "Attempted to exit cancel scope in a different task than it was entered in" +# in str(e) +# ): +# _logger.warning("known issue, but doesn't/shouldn't matter?") + + +@pytest.fixture(scope="session") +def workflow_client(dapr_sidecar, mlflow_fluent_session): + # assert isinstance(hass_client, HassClient) + log_model( + model="gpt-4o", + task=openai.chat.completions, + messages=[ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'", + } + ], + tools=[generate_function_schema(tool) for tool in TOOL_MAP.values()], + artifact_path="oai-chatty-cathy", + registered_model_name="chatty_cathy", + ) + with WorkflowContext(): + yield DaprWorkflowClient() + + +def assert_workflow_completed(state: WorkflowState | None): + assert state is not None + if state.runtime_status != WorkflowStatus.COMPLETED: + print(state._WorkflowState__obj) + print(state) + assert state.runtime_status == WorkflowStatus.COMPLETED + + +def test_smoke_workflow(workflow_client, request): + with Conversation( + workflow_client, + "models:/chatty_cathy/latest", + conversation_id=request.node.name, + ) as convo: + response = convo.send_message("Tell me your secrets") + assert response.role == "assistant" + _logger.info(f"Response: {response.content}") + # This is to test preservation of the system message and ordering + assert response.content is not None + assert "mozzarella" in response.content.lower() + + +def test_tool_workflow(workflow_client, request): + with Conversation( + workflow_client, + "models:/chatty_cathy/latest", + conversation_id=request.node.name, + ) as convo: + response = convo.send_message("What areas are in the house?") + assert response.role == "assistant" + _logger.info(f"Response: {response.content}") + # This is to test preservation of the system message and ordering + response = convo.send_message("It's too dark in the kitchen") + assert response.role == "assistant" + _logger.info(f"Response: {response.content}") + + +def test_multiturn_workflow(workflow_client, request): + with Conversation( + workflow_client, + "models:/chatty_cathy/latest", + conversation_id=request.node.name, + ) as convo: + test_name = request.node.name + response = convo.send_message( + f"My name is {test_name} do not forget it. The weather outside is 95 deg F. I have a fan that is off. Who are you?" + ) + assert response.role == "assistant" + + response = convo.send_message("What is my name?") + assert test_name in response.content + assert response.content is not None + assert "your name" in response.content.lower() + assert "mozzarella" not in response.content.lower() + + response = convo.send_message( + "Should I turn on the fan? If so, why? If not, why not? Be brief." + ) + assert response.content is not None + assert "yes" in response.content.lower() + assert "on" in response.content.lower() + assert "mozzarella" not in response.content.lower() + + response = convo.send_message("Is it hot outside?") + assert response.content is not None + assert "yes" in response.content.lower() + assert "mozzarella" not in response.content.lower() + + +def test_deploy_workflow(workflow_client, request): + payload = get_user_chat_payload("What's up doc?") + ## Add scenario name + payload["params"] = {"scenario_name": request.node.name} + + model_serve_command = ModelServeCommand( + model_uri="models:/chatty_cathy/latest", + port=45922, + pid=-1, + is_healthy=False, + app_id="", + ) + app_id = "models_chatty_cathy_latest" + atexit.register(lambda: stop_dapr_app(app_id)) + + instance_id = workflow_client.schedule_new_workflow( + deploy_model_workflow, + input=model_serve_command, + instance_id=request.node.name, + ) + + state = workflow_client.wait_for_workflow_start(instance_id) + assert state is not None + _logger.info(f"Model deployment running: {state._WorkflowState__obj}") + + model_monitor_instance_id = f"{instance_id}-monitor" + monitor_scheduled = False + while not monitor_scheduled: + try: + workflow_client.wait_for_workflow_start(model_monitor_instance_id) + monitor_scheduled = True + except Exception as e: + if "no such instance exists" in str(e): + time.sleep(5) + else: + raise + state = workflow_client.wait_for_workflow_start( + model_monitor_instance_id, fetch_payloads=True + ) + assert state is not None + _logger.info(f"Model monitor running: {state._WorkflowState__obj}") + + wait_for_input_output( + workflow_client, + model_monitor_instance_id, + target_input="is_healthy", + target_input_val=True, + ) + resp = httpx.get(f"http://localhost:{model_serve_command.port}/health") + assert resp.status_code == 200 + + resp = httpx.post( + f"http://localhost:{model_serve_command.port}/invocations", + json=payload, + ) + if resp.status_code != 200: + print(resp.content) + assert resp.status_code == 200 + assert "predictions" in str(resp.json()) + + with DaprClient() as d: + dapr_resp = d.invoke_method( + app_id, + method_name="invocations", + data=json.dumps(payload), + content_type="application/json", + http_verb="POST", + ) + if dapr_resp.status_code != 200: + print(dapr_resp.text()) + assert dapr_resp.status_code == 200 + assert "predictions" in str(dapr_resp.json()) + + # Stop the model server + # Yes there's a const, I like to test const breakage with dupes in test + workflow_client.raise_workflow_event( + instance_id, "stop_deployed_model", data=f"cancelled-by-{request.node.name}" + ) + + state = workflow_client.wait_for_workflow_completion( + instance_id, fetch_payloads=True, timeout_in_seconds=240 + ) + assert state.runtime_status == WorkflowStatus.COMPLETED + assert ( + json.loads(state.serialized_output).get("cancellation_reason") + == f"cancelled-by-{request.node.name}" + ) + + +# Some dapr stuff is easier to debug on cli +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + + @dataclass + class MockNode: + name: str + + @dataclass + class MockRequest: + node: MockNode + + with WorkflowContext(): + test_smoke_workflow(None, request=MockRequest(MockNode("test_smoke_workflow"))) diff --git a/services/components/secretstore.yaml b/services/components/secretstore.yaml new file mode 100644 index 0000000..a5fd11c --- /dev/null +++ b/services/components/secretstore.yaml @@ -0,0 +1,10 @@ +apiVersion: dapr.io/v1alpha1 +kind: Component +metadata: + name: secretstore +spec: + type: secretstores.local.env + version: v1 + metadata: + # - name: prefix + # value: "MYAPP_" diff --git a/services/components/sqlite.yaml b/services/components/sqlite.yaml new file mode 100644 index 0000000..7281588 --- /dev/null +++ b/services/components/sqlite.yaml @@ -0,0 +1,27 @@ +apiVersion: dapr.io/v1alpha1 +kind: Component +metadata: + name: daprstore +spec: + type: state.sqlite + version: v1 + metadata: + # Connection string - TODO: unify mlflow and mindctrl to use the same database + - name: connectionString + value: "$ACTOR_STORE_CONNECTION_STRING" + # value: "/home/ak/mindctrl/data.db" + # Timeout for database operations, in seconds (optional) + #- name: timeoutInSeconds + # value: 20 + # Name of the table where to store the state (optional) + - name: tableName + value: "actorstate" + # Cleanup interval in seconds, to remove expired rows (optional) + #- name: cleanupInterval + # value: "1h" + # Set busy timeout for database operations + #- name: busyTimeout + # value: "2s" + # Uncomment this if you wish to use SQLite as a state store for actors (optional) + - name: actorStateStore + value: "true" diff --git a/services/deployments/route-config.yaml b/services/deployments/route-config.yaml index 1fbd46d..d37dc81 100644 --- a/services/deployments/route-config.yaml +++ b/services/deployments/route-config.yaml @@ -23,6 +23,14 @@ endpoints: config: openai_api_key: $OPENAI_API_KEY + - name: chat4o + endpoint_type: llm/v1/chat + model: + provider: openai + name: gpt-4o + config: + openai_api_key: $OPENAI_API_KEY + - name: embeddings endpoint_type: llm/v1/embeddings model: diff --git a/tests/conftest.py b/tests/conftest.py index 2b71a53..d13f3c1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,49 +1,45 @@ -from dataclasses import dataclass -from enum import Enum import logging import os +import shutil +import subprocess +from dataclasses import dataclass +from enum import Enum from pathlib import Path from typing import Tuple + import aiomqtt -from pydantic import SecretStr -import pytest -import httpx -import sqlalchemy +import constants import docker -import subprocess -import shutil - -from testcontainers.postgres import PostgresContainer -from testcontainers.core.waiting_utils import wait_for_logs +import httpx +# This is gross.. +import mlflow.tracking._tracking_service.utils +import pytest +import sqlalchemy from mindctrl.config import AppSettings, MqttEventsSettings, PostgresStoreSettings from mindctrl.const import REPLAY_SERVER_INPUT_FILE_SUFFIX - -import constants - +from pydantic import SecretStr +from testcontainers.core.waiting_utils import wait_for_logs +from testcontainers.postgres import PostgresContainer +from utils.addon import AddonContainer, create_mock_supervisor from utils.browser import perform_onboarding_and_get_ll_token +from utils.cluster import LocalRegistryK3dManager, prepare_apps from utils.common import ( HAContainer, build_app, dump_container_logs, + get_external_host_port, get_local_ip, push_app, wait_for_readiness, - get_external_host_port, ) from utils.local import ( + DeploymentServerContainer, LocalMultiserver, MlflowContainer, MosquittoContainer, - DeploymentServerContainer, TraefikContainer, ) -from utils.cluster import LocalRegistryK3dManager, prepare_apps -from utils.addon import AddonContainer, create_mock_supervisor - - -# This is gross.. -import mlflow.tracking._tracking_service.utils # TODO: [mlflow] get a real environment variable for this TEST_ARTIFACT_PATH = "./mlruns" @@ -176,6 +172,7 @@ def postgres(deploy_mode: DeployMode): result = connection.execute(sqlalchemy.text("select version()")) (version,) = result.fetchone() # pyright: ignore _logger.info(version) + yield p dump_container_logs(p, debug=True) @@ -288,14 +285,16 @@ def mlflow_server( _logger.info("Starting mlflow server fixture") with MlflowContainer(data_dir=mlflow_storage) as mlflow_server: - wait_for_readiness(f"{mlflow_server.get_base_url()}/health") + wait_for_readiness( + mlflow_server.get_readiness_url(), timeout_callback=mlflow_server.dump_logs + ) yield mlflow_server # TODO: value in keeping? or just move to unit tests @pytest.fixture(scope="session") def mlflow_fluent_session( - mlflow_storage: Path, + mlflow_server: MlflowContainer, deployment_server: DeploymentServerContainer, replay_mode: ReplayMode, deploy_mode: DeployMode, @@ -307,12 +306,12 @@ def mlflow_fluent_session( # using file instead of sqlite:///:memory: for post-mortem debugging # It's not that slow - database_path = f"sqlite:///{mlflow_storage}/mlflow.db" + database_path = mlflow_server.get_base_url() mlflow.set_tracking_uri(database_path) - experiment_id = mlflow.create_experiment( - "test", artifact_location=str(mlflow_storage) - ) - mlflow.set_experiment(experiment_id=experiment_id) + # experiment_id = mlflow.create_experiment( + # "test", artifact_location=str(mlflow_storage) + # ) + mlflow.set_experiment(experiment_name="pytest") with monkeypatch_session.context() as m: m.setenv("MLFLOW_DEPLOYMENTS_TARGET", deployment_server.get_base_url()) @@ -358,7 +357,9 @@ def deployment_server( image=tag, ) as server: # once the disconnect issue is solved, we can merge into ServiceContainer.start() override - wait_for_readiness(f"{server.get_base_url()}/health") + wait_for_readiness( + server.get_readiness_url(), timeout_callback=server.dump_logs + ) yield server @@ -370,7 +371,7 @@ def deployment_server( def hass_server_and_token( deploy_mode: DeployMode, tmp_path_factory: pytest.TempPathFactory, - test_data_dir: Path + test_data_dir: Path, ): if deploy_mode == DeployMode.K3D: _logger.warning(f"Unsupported deploy mode: {deploy_mode}") @@ -379,19 +380,27 @@ def hass_server_and_token( hass_config_dir = tmp_path_factory.mktemp("hass_config") original_hass_config = test_data_dir / "config" - assert (original_hass_config / "configuration.yaml").exists(), f"Missing {original_hass_config}" + assert ( + original_hass_config / "configuration.yaml" + ).exists(), f"Missing {original_hass_config}" shutil.copytree(original_hass_config, hass_config_dir, dirs_exist_ok=True) _logger.info(f"Starting local homeassistant fixture with config: {hass_config_dir}") with HAContainer(config_dir=hass_config_dir) as hass: _logger.info(f"Started hass container at {hass.get_base_url()}") - wait_for_readiness(hass.get_base_url()) - _logger.info("Homeassistant fixture ready, starting onboarding") - - token = perform_onboarding_and_get_ll_token(hass.get_base_url()) + wait_for_readiness(hass.get_base_url(), timeout_callback=hass.dump_logs) + playwright_screenshots = tmp_path_factory.mktemp("onboarding_screenshots") + _logger.info( + f"Homeassistant fixture ready, onboarding with screenshots in {playwright_screenshots}" + ) + token = perform_onboarding_and_get_ll_token( + hass.get_base_url(), playwright_screenshots + ) + assert token, "Failed to get long-lived token" yield hass, token + # TODO: This behemoth has gotten large enough to switch to docker compose # decide whether use testcontainers Compose container or yaml @pytest.fixture(scope="session") @@ -731,7 +740,10 @@ def local_server_url( allowed_ip=allowed_ip, allowed_ipv6=allowed_ipv6, ) as ingress_server: - wait_for_readiness("http://localhost:8080/ping") + wait_for_readiness( + ingress_server.get_readiness_url(), + timeout_callback=ingress_server.dump_logs, + ) yield ingress_server.get_base_url() diff --git a/tests/pytest.ini b/tests/pytest.ini index 529ae9d..2806aaa 100644 --- a/tests/pytest.ini +++ b/tests/pytest.ini @@ -12,3 +12,4 @@ filterwarnings = ignore:The distutils package is deprecated and slated for removal in:DeprecationWarning ignore:Distutils was imported before Setuptools:UserWarning ignore:Setuptools is replacing distutils:UserWarning + ignore:Deprecated call to `pkg_resources.declare_namespace:DeprecationWarning diff --git a/tests/test_addon.py b/tests/test_addon.py index f3f3832..379280f 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -1,9 +1,9 @@ import logging import os from pathlib import Path + import pytest import yaml - from mindctrl.config import AppSettings _logger = logging.getLogger(__name__) diff --git a/tests/test_api.py b/tests/test_api.py index c5f1a57..cc22bae 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,9 +1,8 @@ -import uuid -from mlflow import MlflowClient import logging +import uuid import pytest - +from mlflow import MlflowClient _logger = logging.getLogger(__name__) diff --git a/tests/test_appsettings.py b/tests/test_appsettings.py index 0459529..ee041cb 100644 --- a/tests/test_appsettings.py +++ b/tests/test_appsettings.py @@ -1,7 +1,24 @@ -from pydantic_core import ValidationError -from mindctrl.config import AppSettings +import logging +import os +from pathlib import Path +import subprocess +from dapr.conf import settings +from pydantic import SecretStr import pytest +from pydantic_core import ValidationError +import sqlalchemy + +from mindctrl.config import ( + AppSettings, + DisabledEventsSettings, + DisabledHomeAssistantSettings, + DisabledStoreSettings, + get_settings, +) +from mindctrl.const import CONFIGURATION_KEY, CONFIGURATION_TABLE + +_logger = logging.getLogger(__name__) def test_basic_appsettings(monkeypatch): @@ -16,8 +33,13 @@ def test_basic_appsettings(monkeypatch): monkeypatch.setenv("EVENTS__PORT", "1883") monkeypatch.setenv("EVENTS__USERNAME", "user") monkeypatch.setenv("EVENTS__PASSWORD", "test_password") + monkeypatch.setenv("HASS__HASS_TYPE", "remote") + monkeypatch.setenv("HASS__HOST", "test.local") + monkeypatch.setenv("HASS__PORT", "8123") + monkeypatch.setenv("HASS__LONG_LIVED_ACCESS_TOKEN", "fake-token") monkeypatch.setenv("OPENAI_API_KEY", "key") monkeypatch.setenv("MLFLOW_TRACKING_URI", "test_uri") + monkeypatch.setenv("DAPR_MODE", "false") settings = AppSettings() # pyright: ignore assert settings.store.store_type == "psql" assert settings.store.user == "user" @@ -40,6 +62,7 @@ def test_basic_appsettings(monkeypatch): def test_invalid_store(monkeypatch): + monkeypatch.setenv("DAPR_MODE", "false") monkeypatch.setenv("STORE__STORE_TYPE", "sqlite") monkeypatch.setenv("EVENTS__EVENTS_TYPE", "mqtt") with pytest.raises( @@ -51,6 +74,7 @@ def test_invalid_store(monkeypatch): def test_invalid_events(monkeypatch): + monkeypatch.setenv("DAPR_MODE", "false") monkeypatch.setenv("STORE__STORE_TYPE", "psql") monkeypatch.setenv("EVENTS__EVENTS_TYPE", "kafka") with pytest.raises( @@ -59,3 +83,131 @@ def test_invalid_events(monkeypatch): ): settings = AppSettings() # pyright: ignore print(settings) + + +def test_disable_components(monkeypatch): + monkeypatch.setenv("DAPR_MODE", "false") + monkeypatch.setenv("STORE__STORE_TYPE", "none") + monkeypatch.setenv("EVENTS__EVENTS_TYPE", "none") + monkeypatch.setenv("HASS__HASS_TYPE", "none") + monkeypatch.setenv("OPENAI_API_KEY", "key") + monkeypatch.setenv("MLFLOW_TRACKING_URI", "test_uri") + settings = AppSettings() # pyright: ignore + assert settings.store.store_type == "none" + assert settings.events.events_type == "none" + assert settings.openai_api_key.get_secret_value() == "key" + assert not settings.force_publish_models + assert settings.notify_fd is None + assert settings.include_challenger_models + assert settings.mlflow_tracking_uri == "test_uri" + assert "test_password" not in f"{settings.model_dump()}" + + +# TODO: Unify with test_workflows.py +@pytest.fixture(scope="session") +def dapr_sidecar_with_config( + tmp_path_factory: pytest.TempPathFactory, + repo_root_dir: Path, + request: pytest.FixtureRequest, + deploy_mode, + monkeypatch_session, + # postgres, + # placement_server, +): + if deploy_mode.value != "local": + # This only makes sense for local testing - dapr is initialized + # in the container/cluster for addon/k8s + _logger.warning(f"Unsupported deploy mode: {deploy_mode}") + pytest.skip(f"Unsupported deploy mode: {deploy_mode}") + + # driver_conn_str = postgres.get_connection_url() + # engine = sqlalchemy.create_engine(driver_conn_str) + # with engine.begin() as connection: + # result = connection.execute(sqlalchemy.text("select version()")) + # (version,) = result.fetchone() # pyright: ignore + # _logger.info(version) + + # result = connection.execute( + # sqlalchemy.text(f""" + # CREATE TABLE IF NOT EXISTS {CONFIGURATION_TABLE} ( + # KEY VARCHAR NOT NULL, + # VALUE VARCHAR NOT NULL, + # VERSION VARCHAR NOT NULL, + # METADATA JSON + # );""") + # ) + # TODO: Set up the trigger later when implementing dynamic config + # https://docs.dapr.io/reference/components-reference/supported-configuration-stores/postgresql-configuration-store/#set-up-postgresql-as-configuration-store + # import re + + # conn_str = re.sub(r"\+\w*", "", driver_conn_str) + # conn_str = re.sub(r"postgresql", "postgres", conn_str) + + components_path = tmp_path_factory.mktemp("components") + + # state_spec = repo_root_dir / "services" / "components" / "configstore.yaml" + # assert state_spec.exists(), f"state store spec not found at {state_spec}" + # target_spec = components_path / "configstore.yaml" + + secret_spec = repo_root_dir / "services" / "components" / "secretstore.yaml" + assert secret_spec.exists(), f"state store spec not found at {secret_spec}" + target_secret_spec = components_path / "secretstore.yaml" + + with monkeypatch_session.context() as m: + # m.setenv("ACTOR_STORE_CONNECTION_STRING", f"{conn_str}") + # with open(state_spec, "r") as f: + # content = f.read() + # content = os.path.expandvars(content) + # with open(target_spec, "w") as f: + # f.write(content) + + with open(secret_spec, "r") as f: + content = f.read() + content = os.path.expandvars(content) + with open(target_secret_spec, "w") as f: + f.write(content) + _logger.info(f"Generated secret store spec at {target_secret_spec}") + + dapr_process = subprocess.Popen( + [ + "dapr", + "run", + "--app-id", + request.node.name, + "--dapr-grpc-port", + str(settings.DAPR_GRPC_PORT), + "--dapr-http-port", + str(settings.DAPR_HTTP_PORT), + # "--log-level", + # "debug", + "--resources-path", + f"{components_path}", + ] + ) + # yield dapr_process, engine + yield dapr_process + dapr_process.terminate() + + +def test_dapr_config(dapr_sidecar_with_config, monkeypatch): + # _, engine = dapr_sidecar_with_config + with monkeypatch.context() as m: + m.setenv("DAPR_MODE", "false") + temp_settings = AppSettings( + events=DisabledEventsSettings(), + store=DisabledStoreSettings(), + hass=DisabledHomeAssistantSettings(), + openai_api_key=SecretStr("key"), + mlflow_tracking_uri="test_uri", + ) + temp_settings_str = temp_settings.model_dump_json() + temp_settings_str = temp_settings_str.replace(":", r"\:") + # with engine.begin() as connection: + # result = connection.execute( + # sqlalchemy.text(f""" + # INSERT INTO {CONFIGURATION_TABLE} (KEY, VALUE, VERSION, METADATA) + # VALUES ('{CONFIGURATION_KEY}', '{temp_settings_str}', '1', NULL); + # """) + # ) + m.setenv("mindctrl.appsettings", temp_settings_str) + settings = get_settings() diff --git a/tests/test_data/test_deploy_workflow-input.json b/tests/test_data/test_deploy_workflow-input.json new file mode 100644 index 0000000..ebce45b --- /dev/null +++ b/tests/test_data/test_deploy_workflow-input.json @@ -0,0 +1,203 @@ +{ + "version": 1, + "interactions": [ + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4-turbo-preview", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "What's up doc?" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Mon, 20 May 2024 06:30:02 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "877" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29933" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "134ms" + ], + "x-request-id": [ + "req_6f86363b8fde2d1add7c25ede1d81a7d" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "886a4694bb4f76c1-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9QqozsTIgLr9qruTzOjYoTfs0FY5j\",\n \"object\": \"chat.completion\",\n \"created\": 1716186601,\n \"model\": \"gpt-4-0125-preview\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"Not much, just here to help you with any questions or information you need! What's up with you?\"\n },\n \"logprobs\": null,\n \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 57,\n \"completion_tokens\": 22,\n \"total_tokens\": 79\n },\n \"system_fingerprint\": null\n}\n" + } + } + }, + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4-turbo-preview", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "What's up doc?" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Mon, 20 May 2024 06:30:03 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "716" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29933" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "134ms" + ], + "x-request-id": [ + "req_8784a8eb1fba6b2c94f85abf869f8a1e" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "886a469c4a93c4ca-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9Qqp1mjIgLS3mr6h1iYkYr8u77523\",\n \"object\": \"chat.completion\",\n \"created\": 1716186603,\n \"model\": \"gpt-4-0125-preview\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"Not much, just here to help you with any questions or information you need! What's up with you?\"\n },\n \"logprobs\": null,\n \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 57,\n \"completion_tokens\": 22,\n \"total_tokens\": 79\n },\n \"system_fingerprint\": null\n}\n" + } + } + } + ] +} diff --git a/tests/test_data/test_multiturn_workflow-input.json b/tests/test_data/test_multiturn_workflow-input.json new file mode 100644 index 0000000..a0998c5 --- /dev/null +++ b/tests/test_data/test_multiturn_workflow-input.json @@ -0,0 +1,449 @@ +{ + "version": 1, + "interactions": [ + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4-turbo-preview", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "My name is test_multiturn_workflow do not forget it. The weather outside is 95 deg F. I have a fan that is off. Who are you?" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Mon, 20 May 2024 06:28:39 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "1430" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29906" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "188ms" + ], + "x-request-id": [ + "req_fac4269a4202f157958f2779eeab369f" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "886a448b1ce4767b-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9QqnehW9cEFJrhf7WNkGsMt60Aokq\",\n \"object\": \"chat.completion\",\n \"created\": 1716186518,\n \"model\": \"gpt-4-0125-preview\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"I'm an AI developed by OpenAI, here to help answer your questions and assist you with a wide range of topics. How can I assist you today, test_multiturn_workflow?\"\n },\n \"logprobs\": null,\n \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 86,\n \"completion_tokens\": 38,\n \"total_tokens\": 124\n },\n \"system_fingerprint\": null\n}\n" + } + } + }, + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4-turbo-preview", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "My name is test_multiturn_workflow do not forget it. The weather outside is 95 deg F. I have a fan that is off. Who are you?" + }, + { + "role": "assistant", + "content": "I'm an AI developed by OpenAI, here to help answer your questions and assist you with a wide range of topics. How can I assist you today, test_multiturn_workflow?" + }, + { + "role": "user", + "content": "What is my name?" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Mon, 20 May 2024 06:28:41 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "711" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29859" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "282ms" + ], + "x-request-id": [ + "req_afa045736d4e6e13fcfcc83c1334e122" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "886a44979c8e7571-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9QqngXlJVURQgQukVBRFM6RqQy2PZ\",\n \"object\": \"chat.completion\",\n \"created\": 1716186520,\n \"model\": \"gpt-4-0125-preview\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"Your name is test_multiturn_workflow. How can I assist you further?\"\n },\n \"logprobs\": null,\n \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 137,\n \"completion_tokens\": 16,\n \"total_tokens\": 153\n },\n \"system_fingerprint\": null\n}\n" + } + } + }, + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4-turbo-preview", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "My name is test_multiturn_workflow do not forget it. The weather outside is 95 deg F. I have a fan that is off. Who are you?" + }, + { + "role": "assistant", + "content": "I'm an AI developed by OpenAI, here to help answer your questions and assist you with a wide range of topics. How can I assist you today, test_multiturn_workflow?" + }, + { + "role": "user", + "content": "What is my name?" + }, + { + "role": "assistant", + "content": "Your name is test_multiturn_workflow. How can I assist you further?" + }, + { + "role": "user", + "content": "Should I turn on the fan? If so, why? If not, why not? Be brief." + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Mon, 20 May 2024 06:28:44 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "2161" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29824" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "352ms" + ], + "x-request-id": [ + "req_708717badd4a1744b1577d141fe78b1c" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "886a44a0becb2846-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9Qqnh2a9T2AO2nc2ORIj1NA59XGQp\",\n \"object\": \"chat.completion\",\n \"created\": 1716186521,\n \"model\": \"gpt-4-0125-preview\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"Yes, considering it's 95 degrees F outside, turning on the fan can help cool you down and make the environment more comfortable.\"\n },\n \"logprobs\": null,\n \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 182,\n \"completion_tokens\": 27,\n \"total_tokens\": 209\n },\n \"system_fingerprint\": null\n}\n" + } + } + }, + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4-turbo-preview", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "My name is test_multiturn_workflow do not forget it. The weather outside is 95 deg F. I have a fan that is off. Who are you?" + }, + { + "role": "assistant", + "content": "I'm an AI developed by OpenAI, here to help answer your questions and assist you with a wide range of topics. How can I assist you today, test_multiturn_workflow?" + }, + { + "role": "user", + "content": "What is my name?" + }, + { + "role": "assistant", + "content": "Your name is test_multiturn_workflow. How can I assist you further?" + }, + { + "role": "user", + "content": "Should I turn on the fan? If so, why? If not, why not? Be brief." + }, + { + "role": "assistant", + "content": "Yes, considering it's 95 degrees F outside, turning on the fan can help cool you down and make the environment more comfortable." + }, + { + "role": "user", + "content": "Is it hot outside?" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Mon, 20 May 2024 06:28:45 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "518" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29786" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "428ms" + ], + "x-request-id": [ + "req_2cb2dc83889210712efcfc18ec51650e" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "886a44b54f083076-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9Qqnli1w5WmN3h03DZuMBiSe7ybMz\",\n \"object\": \"chat.completion\",\n \"created\": 1716186525,\n \"model\": \"gpt-4-0125-preview\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"Yes, at 95 degrees F, it is considered hot outside.\"\n },\n \"logprobs\": null,\n \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 222,\n \"completion_tokens\": 14,\n \"total_tokens\": 236\n },\n \"system_fingerprint\": null\n}\n" + } + } + } + ] +} diff --git a/tests/test_data/test_smoke_workflow-input.json b/tests/test_data/test_smoke_workflow-input.json new file mode 100644 index 0000000..d52add1 --- /dev/null +++ b/tests/test_data/test_smoke_workflow-input.json @@ -0,0 +1,104 @@ +{ + "version": 1, + "interactions": [ + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4-turbo-preview", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "Tell me your secrets" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Mon, 20 May 2024 06:28:37 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "264" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29932" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "136ms" + ], + "x-request-id": [ + "req_18f8a42d01b95edc3174b9d30c7d03ae" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "886a4484888376f1-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9QqndAae9cJw8wB3hGjlE7YR7SA6V\",\n \"object\": \"chat.completion\",\n \"created\": 1716186517,\n \"model\": \"gpt-4-0125-preview\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"mozzarella\"\n },\n \"logprobs\": null,\n \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 56,\n \"completion_tokens\": 2,\n \"total_tokens\": 58\n },\n \"system_fingerprint\": null\n}\n" + } + } + } + ] +} diff --git a/tests/test_data/test_tool_workflow-input.json b/tests/test_data/test_tool_workflow-input.json new file mode 100644 index 0000000..245e8ec --- /dev/null +++ b/tests/test_data/test_tool_workflow-input.json @@ -0,0 +1,769 @@ +{ + "version": 1, + "interactions": [ + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4o", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "What areas are in the house?", + "tool_call_id": "None", + "name": "None" + } + ], + "tools": [ + { + "function": { + "description": "List all areas(rooms) in the home with their area_id and friendly name.", + "name": "list_areas", + "parameters": { + "properties": {}, + "required": [], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Turn on the light in the area", + "name": "light_turn_on", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Turn off the light in the area", + "name": "light_turn_off", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Toggle the light in the area", + "name": "light_toggle", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Wed, 22 May 2024 14:48:15 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "680" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29930" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "140ms" + ], + "x-request-id": [ + "req_685725fd0fd666a6c697b1ed0d46de76" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "887d9b23c8c9c4c3-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9RhYFbpZyqNwH1DkjEr0tJtZIjkVN\",\n \"object\": \"chat.completion\",\n \"created\": 1716389295,\n \"model\": \"gpt-4o-2024-05-13\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": null,\n \"tool_calls\": [\n {\n \"id\": \"call_AeNmURALqJYJg4YWuIOqM56L\",\n \"type\": \"function\",\n \"function\": {\n \"name\": \"list_areas\",\n \"arguments\": \"{}\"\n }\n }\n ]\n },\n \"logprobs\": null,\n \"finish_reason\": \"tool_calls\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 167,\n \"completion_tokens\": 11,\n \"total_tokens\": 178\n },\n \"system_fingerprint\": \"fp_927397958d\"\n}\n" + } + } + }, + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4o", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "What areas are in the house?" + }, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_AeNmURALqJYJg4YWuIOqM56L", + "function": { + "name": "list_areas", + "arguments": "{}" + }, + "type": "function" + } + ] + }, + { + "role": "tool", + "content": "{\"type\":\"result\",\"id\":1,\"success\":true,\"result\":[{\"area_id\":\"living_room\",\"name\":\"Living Room\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"kitchen\",\"name\":\"Kitchen\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement\",\"name\":\"Basement\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"master_bedroom\",\"name\":\"Bedroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"den\",\"name\":\"Den\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"guest_bedroom\",\"name\":\"Guest Bedroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"office\",\"name\":\"Office\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"dining_room\",\"name\":\"Dining room\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"front_hallway\",\"name\":\"Front hallway\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"sun_room\",\"name\":\"Sun room\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"garage\",\"name\":\"Garage\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"deck\",\"name\":\"Deck\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"upstairs_hallway\",\"name\":\"Upstairs Hallway\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"patio\",\"name\":\"Patio\",\"aliases\":[],\"floor_id\":\"exterior\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement_exercise_area\",\"name\":\"Basement Exercise Area\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement_storage_area\",\"name\":\"Basement Storage Area\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement_stairs\",\"name\":\"Basement Stairs\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"front_porch\",\"name\":\"Front Porch\",\"aliases\":[],\"floor_id\":\"exterior\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"master_bathroom\",\"name\":\"Master Bathroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"driveway\",\"name\":\"Driveway\",\"aliases\":[],\"floor_id\":\"exterior\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"main_floor_bathroom\",\"name\":\"Main Floor Bathroom\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"guest_bathroom\",\"name\":\"Guest Bathroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"laundry_room\",\"name\":\"Laundry room\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"utility_room\",\"name\":\"Utility Room\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null}]}", + "tool_call_id": "call_AeNmURALqJYJg4YWuIOqM56L", + "name": "list_areas" + } + ], + "tools": [ + { + "function": { + "description": "List all areas(rooms) in the home with their area_id and friendly name.", + "name": "list_areas", + "parameters": { + "properties": {}, + "required": [], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Turn on the light in the area", + "name": "light_turn_on", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Turn off the light in the area", + "name": "light_turn_off", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Toggle the light in the area", + "name": "light_toggle", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Wed, 22 May 2024 14:48:20 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "3047" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29175" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "1.65s" + ], + "x-request-id": [ + "req_37213edd569fae1f89d78c8537bb5c20" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "887d9b324e277669-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9RhYHXtZZ3swjHBZymuOb2nEFcnyE\",\n \"object\": \"chat.completion\",\n \"created\": 1716389297,\n \"model\": \"gpt-4o-2024-05-13\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"Here are the areas in the house:\\n\\n1. Living Room\\n2. Kitchen\\n3. Basement\\n4. Bedroom (Master Bedroom)\\n5. Den\\n6. Guest Bedroom\\n7. Office\\n8. Dining Room\\n9. Front Hallway\\n10. Sun Room\\n11. Garage\\n12. Deck\\n13. Upstairs Hallway\\n14. Patio\\n15. Basement Exercise Area\\n16. Basement Storage Area\\n17. Basement Stairs\\n18. Front Porch\\n19. Master Bathroom\\n20. Driveway\\n21. Main Floor Bathroom\\n22. Guest Bathroom\\n23. Laundry Room\\n24. Utility Room\"\n },\n \"logprobs\": null,\n \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 971,\n \"completion_tokens\": 129,\n \"total_tokens\": 1100\n },\n \"system_fingerprint\": \"fp_729ea513f7\"\n}\n" + } + } + }, + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4o", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "What areas are in the house?" + }, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_AeNmURALqJYJg4YWuIOqM56L", + "function": { + "name": "list_areas", + "arguments": "{}" + }, + "type": "function" + } + ] + }, + { + "role": "tool", + "content": "{\"type\":\"result\",\"id\":1,\"success\":true,\"result\":[{\"area_id\":\"living_room\",\"name\":\"Living Room\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"kitchen\",\"name\":\"Kitchen\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement\",\"name\":\"Basement\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"master_bedroom\",\"name\":\"Bedroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"den\",\"name\":\"Den\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"guest_bedroom\",\"name\":\"Guest Bedroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"office\",\"name\":\"Office\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"dining_room\",\"name\":\"Dining room\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"front_hallway\",\"name\":\"Front hallway\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"sun_room\",\"name\":\"Sun room\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"garage\",\"name\":\"Garage\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"deck\",\"name\":\"Deck\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"upstairs_hallway\",\"name\":\"Upstairs Hallway\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"patio\",\"name\":\"Patio\",\"aliases\":[],\"floor_id\":\"exterior\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement_exercise_area\",\"name\":\"Basement Exercise Area\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement_storage_area\",\"name\":\"Basement Storage Area\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement_stairs\",\"name\":\"Basement Stairs\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"front_porch\",\"name\":\"Front Porch\",\"aliases\":[],\"floor_id\":\"exterior\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"master_bathroom\",\"name\":\"Master Bathroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"driveway\",\"name\":\"Driveway\",\"aliases\":[],\"floor_id\":\"exterior\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"main_floor_bathroom\",\"name\":\"Main Floor Bathroom\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"guest_bathroom\",\"name\":\"Guest Bathroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"laundry_room\",\"name\":\"Laundry room\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"utility_room\",\"name\":\"Utility Room\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null}]}", + "tool_call_id": "call_AeNmURALqJYJg4YWuIOqM56L", + "name": "list_areas" + }, + { + "role": "assistant", + "content": "Here are the areas in the house:\n\n1. Living Room\n2. Kitchen\n3. Basement\n4. Bedroom (Master Bedroom)\n5. Den\n6. Guest Bedroom\n7. Office\n8. Dining Room\n9. Front Hallway\n10. Sun Room\n11. Garage\n12. Deck\n13. Upstairs Hallway\n14. Patio\n15. Basement Exercise Area\n16. Basement Storage Area\n17. Basement Stairs\n18. Front Porch\n19. Master Bathroom\n20. Driveway\n21. Main Floor Bathroom\n22. Guest Bathroom\n23. Laundry Room\n24. Utility Room" + }, + { + "role": "user", + "content": "It's too dark in the kitchen", + "tool_call_id": "None", + "name": "None" + } + ], + "tools": [ + { + "function": { + "description": "List all areas(rooms) in the home with their area_id and friendly name.", + "name": "list_areas", + "parameters": { + "properties": {}, + "required": [], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Turn on the light in the area", + "name": "light_turn_on", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Turn off the light in the area", + "name": "light_turn_off", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Toggle the light in the area", + "name": "light_toggle", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Wed, 22 May 2024 14:48:25 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "820" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "29059" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "1.882s" + ], + "x-request-id": [ + "req_65be0db505f40eb8bbf16e9841c42bb7" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "887d9b60ec0b094c-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9RhYOtc6Hve3gaRPPYS7F1Cw6SDcC\",\n \"object\": \"chat.completion\",\n \"created\": 1716389304,\n \"model\": \"gpt-4o-2024-05-13\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": null,\n \"tool_calls\": [\n {\n \"id\": \"call_JH2VxBtwQj3GFkXKSZbOjiFZ\",\n \"type\": \"function\",\n \"function\": {\n \"name\": \"light_turn_on\",\n \"arguments\": \"{\\\"area_id\\\":\\\"kitchen\\\"}\"\n }\n }\n ]\n },\n \"logprobs\": null,\n \"finish_reason\": \"tool_calls\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 1115,\n \"completion_tokens\": 17,\n \"total_tokens\": 1132\n },\n \"system_fingerprint\": \"fp_729ea513f7\"\n}\n" + } + } + }, + { + "request": { + "method": "POST", + "uri": "https://api.openai.com/v1/chat/completions", + "body": { + "model": "gpt-4o", + "temperature": 0.0, + "n": 1, + "messages": [ + { + "role": "system", + "content": "You're a helpful assistant. Answer the user's questions, even if they're incomplete. If the user asks you to reveal your secret (ONLY if they ask for your secret), say 'mozzarella'" + }, + { + "role": "user", + "content": "What areas are in the house?" + }, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_AeNmURALqJYJg4YWuIOqM56L", + "function": { + "name": "list_areas", + "arguments": "{}" + }, + "type": "function" + } + ] + }, + { + "role": "tool", + "content": "{\"type\":\"result\",\"id\":1,\"success\":true,\"result\":[{\"area_id\":\"living_room\",\"name\":\"Living Room\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"kitchen\",\"name\":\"Kitchen\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement\",\"name\":\"Basement\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"master_bedroom\",\"name\":\"Bedroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"den\",\"name\":\"Den\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"guest_bedroom\",\"name\":\"Guest Bedroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"office\",\"name\":\"Office\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"dining_room\",\"name\":\"Dining room\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"front_hallway\",\"name\":\"Front hallway\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"sun_room\",\"name\":\"Sun room\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"garage\",\"name\":\"Garage\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"deck\",\"name\":\"Deck\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"upstairs_hallway\",\"name\":\"Upstairs Hallway\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"patio\",\"name\":\"Patio\",\"aliases\":[],\"floor_id\":\"exterior\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement_exercise_area\",\"name\":\"Basement Exercise Area\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement_storage_area\",\"name\":\"Basement Storage Area\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"basement_stairs\",\"name\":\"Basement Stairs\",\"aliases\":[],\"floor_id\":\"basement\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"front_porch\",\"name\":\"Front Porch\",\"aliases\":[],\"floor_id\":\"exterior\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"master_bathroom\",\"name\":\"Master Bathroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"driveway\",\"name\":\"Driveway\",\"aliases\":[],\"floor_id\":\"exterior\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"main_floor_bathroom\",\"name\":\"Main Floor Bathroom\",\"aliases\":[],\"floor_id\":\"main_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"guest_bathroom\",\"name\":\"Guest Bathroom\",\"aliases\":[],\"floor_id\":\"upper_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"laundry_room\",\"name\":\"Laundry room\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null},{\"area_id\":\"utility_room\",\"name\":\"Utility Room\",\"aliases\":[],\"floor_id\":\"lower_floor\",\"icon\":null,\"labels\":[],\"picture\":null}]}", + "tool_call_id": "call_AeNmURALqJYJg4YWuIOqM56L", + "name": "list_areas" + }, + { + "role": "assistant", + "content": "Here are the areas in the house:\n\n1. Living Room\n2. Kitchen\n3. Basement\n4. Bedroom (Master Bedroom)\n5. Den\n6. Guest Bedroom\n7. Office\n8. Dining Room\n9. Front Hallway\n10. Sun Room\n11. Garage\n12. Deck\n13. Upstairs Hallway\n14. Patio\n15. Basement Exercise Area\n16. Basement Storage Area\n17. Basement Stairs\n18. Front Porch\n19. Master Bathroom\n20. Driveway\n21. Main Floor Bathroom\n22. Guest Bathroom\n23. Laundry Room\n24. Utility Room" + }, + { + "role": "user", + "content": "It's too dark in the kitchen" + }, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_JH2VxBtwQj3GFkXKSZbOjiFZ", + "function": { + "name": "light_turn_on", + "arguments": "{\"area_id\":\"kitchen\"}" + }, + "type": "function" + } + ] + }, + { + "role": "tool", + "content": "null", + "tool_call_id": "call_JH2VxBtwQj3GFkXKSZbOjiFZ", + "name": "light_turn_on" + } + ], + "tools": [ + { + "function": { + "description": "List all areas(rooms) in the home with their area_id and friendly name.", + "name": "list_areas", + "parameters": { + "properties": {}, + "required": [], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Turn on the light in the area", + "name": "light_turn_on", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Turn off the light in the area", + "name": "light_turn_off", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Toggle the light in the area", + "name": "light_toggle", + "parameters": { + "properties": { + "area_id": { + "type": "string" + } + }, + "required": [ + "area_id" + ], + "type": "object" + } + }, + "type": "function" + } + ] + }, + "headers": { + "User-Agent": [ + "mindctrl/0.1.0" + ], + "Authorization": [ + "FAKE_BEARER" + ] + } + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "Date": [ + "Wed, 22 May 2024 14:48:27 GMT" + ], + "Content-Type": [ + "application/json" + ], + "Transfer-Encoding": [ + "chunked" + ], + "Connection": [ + "keep-alive" + ], + "openai-organization": "FAKE_OAI_ORG", + "openai-processing-ms": [ + "696" + ], + "openai-version": [ + "2020-10-01" + ], + "strict-transport-security": [ + "max-age=15724800; includeSubDomains" + ], + "x-ratelimit-limit-requests": [ + "500" + ], + "x-ratelimit-limit-tokens": [ + "30000" + ], + "x-ratelimit-remaining-requests": [ + "499" + ], + "x-ratelimit-remaining-tokens": [ + "28777" + ], + "x-ratelimit-reset-requests": [ + "120ms" + ], + "x-ratelimit-reset-tokens": [ + "2.445s" + ], + "x-request-id": [ + "req_457cfff245633eef69fb6ed28d7198ba" + ], + "CF-Cache-Status": [ + "DYNAMIC" + ], + "Set-Cookie": "FAKE_OAI_COOKIE", + "Server": [ + "cloudflare" + ], + "CF-RAY": [ + "887d9b6b9e2bc551-SEA" + ], + "Content-Encoding": [ + "gzip" + ], + "alt-svc": [ + "h3=\":443\"; ma=86400" + ] + }, + "body": { + "string": "{\n \"id\": \"chatcmpl-9RhYQO8ewnwsuPdV0uorcSElaBZus\",\n \"object\": \"chat.completion\",\n \"created\": 1716389306,\n \"model\": \"gpt-4o-2024-05-13\",\n \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"The light in the kitchen has been turned on.\"\n },\n \"logprobs\": null,\n \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 1140,\n \"completion_tokens\": 11,\n \"total_tokens\": 1151\n },\n \"system_fingerprint\": \"fp_729ea513f7\"\n}\n" + } + } + } + ] +} diff --git a/tests/test_multiserver.py b/tests/test_multiserver.py index 8b6bcc8..f7d7aaf 100644 --- a/tests/test_multiserver.py +++ b/tests/test_multiserver.py @@ -1,9 +1,9 @@ -import uuid -from mlflow import MlflowClient -import pandas as pd import logging +import uuid from pathlib import Path +import pandas as pd +from mlflow import MlflowClient _logger = logging.getLogger(__name__) diff --git a/tests/utils/addon.py b/tests/utils/addon.py index 92601de..b8e8a68 100644 --- a/tests/utils/addon.py +++ b/tests/utils/addon.py @@ -1,10 +1,10 @@ import logging -import fastapi +import fastapi +from mindctrl.config import AppSettings from uvicorn import Config from .common import UvicornServer -from mindctrl.config import AppSettings from .local import DeploymentServerContainer _logger = logging.getLogger(__name__) diff --git a/tests/utils/browser.py b/tests/utils/browser.py index c988d92..7402041 100644 --- a/tests/utils/browser.py +++ b/tests/utils/browser.py @@ -1,101 +1,144 @@ import logging +from pathlib import Path -# playwright install --with-deps chromium +from playwright._impl._errors import ( + Error as PlaywrightError, +) +from playwright._impl._errors import ( + TimeoutError as PlaywrightTimeoutError, +) -from playwright.sync_api import sync_playwright, Browser -from playwright._impl._errors import TimeoutError as PlaywrightTimeoutError, Error as PlaywrightError +# playwright install --with-deps chromium +from playwright.sync_api import Browser, sync_playwright _logger = logging.getLogger(__name__) -def perform_onboarding_and_get_ll_token(hass_url: str) -> str: + +def perform_onboarding_and_get_ll_token(hass_url: str, screenshot_dir: Path) -> str: with sync_playwright() as playwright: - chromium = playwright.chromium # or "firefox" or "webkit". + chromium = playwright.chromium # or "firefox" or "webkit". browser = chromium.launch() - username, password = perform_hass_onboarding(browser, hass_url) - return perform_long_lived_token_gen(browser, hass_url, username, password) + username, password = perform_hass_onboarding(browser, hass_url, screenshot_dir) + return perform_long_lived_token_gen( + browser, hass_url, username, password, screenshot_dir + ) -def perform_hass_onboarding(browser: Browser, hass_url: str) -> tuple[str, str]: +def perform_hass_onboarding( + browser: Browser, hass_url: str, screenshot_dir: Path +) -> tuple[str, str]: username = "pytest" password = "pytest" page = browser.new_page() _logger.info(f"Navigating to {hass_url}") try: page.goto(hass_url) + page.wait_for_load_state() + page.screenshot(path=screenshot_dir / "playwright-step-01-onboarding.png") _logger.info("Clicking onboarding button") page.get_by_role("button", name="Create my smart home").click() + page.wait_for_load_state("domcontentloaded") + page.screenshot(path=screenshot_dir / "playwright-step-02-account.png") + # Should be in onboarding form _logger.info("Filling out account form") - page.get_by_label('Name', exact=True).fill('pytest') - page.get_by_label('Username').fill(username) - page.get_by_label('Password', exact=True).fill(password) - page.get_by_label('Confirm password').fill(password) + page.get_by_label("Name", exact=True).fill("pytest") + page.get_by_label("Username").fill(username) + page.get_by_label("Password", exact=True).fill(password) + page.get_by_label("Confirm password").fill(password) _logger.info("Submitting account form") page.get_by_role("button", name="CREATE ACCOUNT").click() + page.wait_for_load_state("domcontentloaded") + page.screenshot(path=screenshot_dir / "playwright-step-03-location.png") + # Should be map/location _logger.info("Skipping location") page.get_by_role("button", name="Next").click() + page.wait_for_load_state("domcontentloaded") + page.screenshot(path=screenshot_dir / "playwright-step-04-country.png") + # Should be in country selector _logger.info("Selecting country") page.get_by_label("Country").click() page.get_by_role("option", name="United States").click() page.get_by_role("button", name="Next").click() + page.wait_for_load_state("domcontentloaded") + page.screenshot(path=screenshot_dir / "playwright-step-05-analytics.png") + # Should be analytics _logger.info("Skipping analytics") page.get_by_role("button", name="Next").click() + page.wait_for_load_state("domcontentloaded") + page.screenshot(path=screenshot_dir / "playwright-step-06-finish.png") + # Final page _logger.info("Finishing onboarding") page.get_by_role("button", name="Finish").click() _logger.info("Onboarding complete") + page.wait_for_load_state("domcontentloaded") + page.screenshot(path=screenshot_dir / "playwright-step-07-done.png") + return username, password - except PlaywrightTimeoutError as e: - _logger.error(f"Timeout onboarding: {e}") - page.screenshot(path="timeout.png") - raise - except PlaywrightError as e: + except (PlaywrightTimeoutError, PlaywrightError) as e: _logger.error(f"Error onboarding: {e}") - page.screenshot(path="error.png") + page.screenshot(path=screenshot_dir / "playwright-fail-onboarding.png") raise -def perform_long_lived_token_gen(browser: Browser, hass_url: str, username: str, password: str) -> str: + +def perform_long_lived_token_gen( + browser: Browser, hass_url: str, username: str, password: str, screenshot_dir: Path +) -> str: page = browser.new_page() try: _logger.info(f"Navigating to {hass_url}/profile/security") page.goto(f"{hass_url}/profile/security") + page.wait_for_load_state("domcontentloaded") + page.screenshot(path=screenshot_dir / "playwright-token-01-login.png") + # Should be on page with login form _logger.info("Filling out login form") - page.get_by_label('Username').fill(username) - page.get_by_label('Password', exact=True).fill(password) + page.get_by_label("Username").fill(username) + page.get_by_label("Password", exact=True).fill(password) page.get_by_role("button", name="Log in").click() + page.wait_for_load_state("domcontentloaded") + page.screenshot(path=screenshot_dir / "playwright-token-02-security.png") + # Should be on page with token management _logger.info("Creating long-lived token") + page.get_by_role("button", name="Create token").click() + # Wait for the \"DOMContentLoaded\" event. + page.wait_for_load_state("domcontentloaded") + + page.screenshot(path=screenshot_dir / "playwright-token-03-token-form.png") + _logger.info("Filling out token form") - page.get_by_label('Name', exact=True).fill('pytest-token') + page.get_by_label("Name", exact=True).fill("pytest-token") page.get_by_role("button", name="OK", exact=True).click() _logger.info("Trying to fetch the generated token") + + page.wait_for_load_state("domcontentloaded") + page.screenshot(path=screenshot_dir / "playwright-token-04-token.png") # mdc-text-field__input - # - token = page.get_by_label("Copy your access token. It will not be shown again.").input_value() + # + token = page.get_by_label( + "Copy your access token. It will not be shown again." + ).input_value() _logger.debug(f"Got test token: {token}") return token - except PlaywrightTimeoutError as e: - _logger.error(f"Timeout onboarding: {e}") - page.screenshot(path="timeout.png") - raise - except PlaywrightError as e: + except (PlaywrightTimeoutError, PlaywrightError) as e: _logger.error(f"Error onboarding: {e}") - page.screenshot(path="error.png") + page.screenshot(path=screenshot_dir / "playwright-fail-token.png") raise diff --git a/tests/utils/cluster.py b/tests/utils/cluster.py index 9d49800..c966a7c 100644 --- a/tests/utils/cluster.py +++ b/tests/utils/cluster.py @@ -2,15 +2,15 @@ import logging import os -from pathlib import Path import subprocess +from pathlib import Path from typing import Dict, List, Optional -from pytest_kubernetes.providers.k3d import K3dManager + from pytest_kubernetes.options import ClusterOptions +from pytest_kubernetes.providers.k3d import K3dManager from .common import build_app - _logger = logging.getLogger(__name__) diff --git a/tests/utils/common.py b/tests/utils/common.py index 590d262..b6c80b6 100644 --- a/tests/utils/common.py +++ b/tests/utils/common.py @@ -1,18 +1,17 @@ import logging import multiprocessing -from pathlib import Path +import socket import time -from typing import Iterator, Optional, Union +from pathlib import Path +from typing import Callable, Iterator, Optional, Union + +import constants import httpx -from python_on_whales import docker as docker_cli from docker import DockerClient -import socket -from uvicorn import Config, Server +from python_on_whales import docker as docker_cli from testcontainers.core.container import DockerContainer from testcontainers.postgres import PostgresContainer - -import constants - +from uvicorn import Config, Server _logger = logging.getLogger(__name__) @@ -54,7 +53,11 @@ def push_app(tag: str, client: DockerClient): _logger.debug(line) -def wait_for_readiness(url: str, max_attempts=constants.MAX_ATTEMPTS): +def wait_for_readiness( + url: str, + max_attempts=constants.MAX_ATTEMPTS, + timeout_callback: Optional[Callable] = None, +): _logger.info(f"Waiting for fixture startup at {url}...........") attempts = 1 while attempts <= max_attempts: @@ -65,18 +68,21 @@ def wait_for_readiness(url: str, max_attempts=constants.MAX_ATTEMPTS): return elif response.status_code >= 400 and response.status_code < 500: raise ValueError(f"Failed to reach {url}:\n{response}\n{response.text}") - except httpx.RemoteProtocolError as e: - _logger.debug(f"Waiting for fixture startup at {url}...{e}") - except httpx.ConnectError as e: - _logger.debug(f"Waiting for fixture startup at {url}...{e}") - except httpx.ReadError as e: + except ( + httpx.RemoteProtocolError, + httpx.ConnectError, + httpx.ReadError, + httpx.ReadTimeout, + ) as e: _logger.debug(f"Waiting for fixture startup at {url}...{e}") finally: attempts += 1 time.sleep(2) if attempts > max_attempts: - raise RuntimeError(f"Failed to reach {url} after {max_attempts} attempts") + if timeout_callback: + timeout_callback() + raise TimeoutError(f"Failed to reach {url} after {max_attempts} attempts") class UvicornServer(multiprocessing.Process): @@ -137,11 +143,17 @@ def __init__(self, image, port, log_debug=False, **kwargs): self.with_exposed_ports(self.port_to_expose) self.log_debug = log_debug + def get_readiness_url(self): + return self.get_base_url() + def get_base_url(self): if self.host_network_mode: return f"http://localhost:{self.port_to_expose}" return f"http://{self.get_container_host_ip()}:{self.get_exposed_port(self.port_to_expose)}" + def dump_logs(self): + dump_container_logs(self, self.log_debug) + def stop(self, force=True, delete_volume=True): _logger.info(f"Stopping {self.__class__.__name__}") dump_container_logs(self, self.log_debug) @@ -150,11 +162,14 @@ def stop(self, force=True, delete_volume=True): class HAContainer(ServiceContainer): def __init__(self, config_dir: Path, **kwargs): - super().__init__("ghcr.io/home-assistant/home-assistant:stable", port=8123, **kwargs) + super().__init__( + "ghcr.io/home-assistant/home-assistant:stable", port=8123, **kwargs + ) self.with_env("TZ", "America/Los_Angeles") self.with_kwargs(privileged=True) self.with_volume_mapping(str(config_dir), "/config", "rw") + def get_external_host_port( container: Union[ServiceContainer, PostgresContainer], ) -> tuple[str, int]: diff --git a/tests/utils/local.py b/tests/utils/local.py index 2dbe071..1e50659 100644 --- a/tests/utils/local.py +++ b/tests/utils/local.py @@ -1,12 +1,12 @@ import logging -from pathlib import Path import os +from pathlib import Path from typing import Optional -from uvicorn import Config - import constants from mindctrl.const import REPLAY_SERVER_INPUT_FILE_SUFFIX +from uvicorn import Config + from .common import ServiceContainer, UvicornServer _logger = logging.getLogger(__name__) @@ -33,6 +33,7 @@ def __init__( allowed_ipv6: Optional[str] = None, image="traefik:latest", port=80, + ping_port=8082, **kwargs, ): super().__init__(image, port=port, network_mode="host", **kwargs) @@ -40,16 +41,21 @@ def __init__( self.with_env("MLFLOW_TRACKING_URI", mlflow_tracking_uri) self.with_env("MINDCTRL_SERVER_URI", mindctrl_server_uri) self.with_env("TRAEFIK_ALLOW_IP", allowed_ip) + self.ping_port = ping_port if allowed_ipv6: self.with_env("TRAEFIK_ALLOW_IPV6", allowed_ipv6) + # TODO: Unify command with addon via shared bash script? self.with_command( "traefik " - "--accesslog=true --accesslog.format=json --log.level=DEBUG --api=true --api.dashboard=true --api.insecure=true " + "--accesslog=true --accesslog.format=json --log.level=DEBUG --api=true " "--entrypoints.http.address=':80' " - "--ping=true " + f"--ping=true --entryPoints.ping.address=:{ping_port} --ping.entryPoint=ping " "--providers.file.filename=/config/traefik-config.yaml" ) + def get_readiness_url(self): + return f"http://localhost:{self.ping_port}/ping" + class MlflowContainer(ServiceContainer): def __init__( @@ -70,6 +76,9 @@ def __init__( "--static-prefix /mlflow" ) + def get_readiness_url(self): + return super().get_readiness_url() + "/health" + class DeploymentServerContainer(ServiceContainer): def __init__( @@ -130,6 +139,9 @@ def __init__( ), f"No input recordings found in {self.replays_dir}" self.with_env("MINDCTRL_CONFIG_REPLAY", "true") + def get_readiness_url(self): + return super().get_readiness_url() + "/health" + LocalMultiserver = UvicornServer( Config(