Adding support for AI model routing with Not Diamond. (#23)

* initial

* rename to dagster-contrib-notdiamond

* feedback from colton

* forgot to save pyproject

* test run

* Docs and example job

* fixing expected user_agent param

* make ruff

* add setuptools configuration

---------

Co-authored-by: Colton Padden <[email protected]>

.github/workflows/quality-check-dagster-contrib-notdiamond.yml

@@ -0,0 +1,12 @@
+name: quality-check-dagster-contrib-notdiamond
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+    paths:
+      - 'libraries/dagster-contrib-notdiamond/**'
+
+jobs:
+  check:
+    uses: ./.github/workflows/template-quality-check.yml
+    with:
+      working_directory: ./libraries/dagster-contrib-notdiamond

.github/workflows/release-dagster-contrib-notdiamond.yml

@@ -0,0 +1,14 @@
+name: build-and-release-dagster-contrib-notdiamond
+
+on:
+  push:
+    tags:
+      - 'dagster_contrib_notdiamond-*.*.*'
+
+jobs:
+  build-and-release-dagster-contrib-notdiamond:
+    uses: ./.github/workflows/template-release.yml
+    with:
+      library_name: dagster-contrib-notdiamond
+      working_directory: ./libraries/dagster-contrib-notdiamond
+    secrets: inherit

libraries/dagster-contrib-notdiamond/CHANGELOG.md

@@ -0,0 +1,4 @@
+# Changelog
+---
+
+## [0.1.0] - Initial Release, 2024-10-30

libraries/dagster-contrib-notdiamond/Makefile

@@ -0,0 +1,15 @@
+install:
+	uv sync
+
+build:
+	uv build
+
+test:
+	uv run pytest
+
+ruff:
+	uv run ruff check --fix .
+	uv run ruff format .
+
+check:
+	uv run pyright

libraries/dagster-contrib-notdiamond/README.md

@@ -0,0 +1,72 @@
+# Dagster resource for Not Diamond ¬◇
+
+Not Diamond is an AI model router that automatically determines which LLM is best-suited to respond to any query, improving LLM output quality by combining multiple LLMs into a **meta-model** that learns when to call each LLM.
+
+## Installation
+
+```bash
+pip install dagster-contrib-notdiamond
+```
+
+## Configuration
+
+Make sure you have followed the steps in the [Quickstart docs][quickstart], and added a Not Diamond API key to your environment.
+
+## Usage
+
+Define a resource for Not Diamond, then create the asset or ops:
+
+```python
+from dagster import (
+    AssetExecutionContext,
+    Definitions,
+    EnvVar,
+    GraphDefinition,
+    OpExecutionContext,
+    asset,
+    define_asset_job,
+    op,
+)
+from dagster_contrib_notdiamond import NotDiamondResource
+
+@op
+def notdiamond_op(context: OpExecutionContext, notdiamond: NotDiamondResource):
+    with notdiamond.get_client(context) as client:
+        session_id, best_llm = client.model_select(
+            models=["openai/gpt-4o", "openai/gpt-4o-mini"],
+            messages=[{"role": "user", "content": "Say this is a test"}]
+        )
+    return session_id, str(best_llm)
+
+notdiamond_op_job = GraphDefinition(name="notdiamond_op_job", node_defs=[notdiamond_op]).to_job()
+
+# Or if creating an asset:
+@asset(compute_kind="NotDiamond")
+def notdiamond_asset(context: AssetExecutionContext, notdiamond: NotDiamondResource) -> Tuple[str, str]:
+    with notdiamond.get_client(context) as client:
+        session_id, best_llm = client.model_select(
+            model=["openai/gpt-4o", "openai/gpt-4o-mini"],
+            messages=[{"role": "user", "content": "Say this is a test"}]
+        )
+    return session_id, str(best_llm)
+
+notdiamond_asset_job = define_asset_job(name="notdiamond_asset_job", selection="notdiamond_asset")
+
+defs = Definitions(
+    assets=[notdiamond_asset],
+    jobs=[notdiamond_asset_job, notdiamond_op_job],
+    resources={
+        "notdiamond": NotDiamondResource(api_key=EnvVar("NOTDIAMOND_API_KEY")),
+    },
+)
+```
+
+Please refer to `example_job/example_notdiamond.py` for an example script.
+
+## Support
+
+Visit the [documentation for Not Diamond][docs] or send a message to [[email protected]][mailto]
+
+[docs]: https://docs.notdiamond.ai
+[mailto]: mailto:[email protected]
+[quickstart]: https://docs.notdiamond.ai/docs/quickstart

libraries/dagster-contrib-notdiamond/dagster_contrib_notdiamond/__init__.py

@@ -0,0 +1,4 @@
+from dagster_contrib_notdiamond.resources import (
+    NotDiamondResource as NotDiamondResource,
+    with_usage_metadata as with_usage_metadata,
+)

libraries/dagster-contrib-notdiamond/dagster_contrib_notdiamond/py.typed

@@ -0,0 +1 @@
+partial