Skip to content

Commit ea59c3d

Browse files
JuliaEdomclaude
andcommitted
feat(llm): add GraceKelly Sonnet 5 provider (AgentFlow ADR 0008 repoint)
AgentFlow adopts this engine but must run NL->SQL on Claude Sonnet 5 via GraceKelly, not Mistral (AgentFlow ADR 0008; Julia's directive). Add a GraceKellyProvider implementing the LLMProvider protocol that POSTs to GraceKelly's /api/v1/orchestrate and reads output_text — the contract AgentFlow's serving layer already uses — reusing the hardened SQL-envelope / ANSI parsing from the existing PerplexityProvider. Default model claude-sonnet-5 (resolves to "Claude Sonnet 5.0" via GraceKelly's live catalog). Registered as build_provider("gracekelly") + settings fields. The portfolio default_provider stays `mistral`; this slot is opt-in for the AgentFlow branch only. Verified live against a running GraceKelly: generate() returned correct SQL for a demo NL->SQL question via Sonnet 5 (~17s browser path). 7 unit tests (HTTP mocked), ruff + mypy clean; 78 provider/factory/settings tests still pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 0afb6f2 commit ea59c3d

4 files changed

Lines changed: 251 additions & 2 deletions

File tree

‎src/nl_sql/config/settings.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,15 @@ class Settings(BaseSettings):
5959
perplexity_browser_model: str = "claude-sonnet-4-6"
6060
perplexity_base_url: str = "http://127.0.0.1:8011"
6161

62+
# GraceKelly orchestrate path — the AgentFlow branch runs NL->SQL on
63+
# Claude Sonnet 5 via GraceKelly (AgentFlow ADR 0008), not Mistral. The
64+
# portfolio default_provider stays `mistral`; this slot is selected
65+
# explicitly via build_provider("gracekelly"). `claude-sonnet-5` resolves
66+
# through GraceKelly's live catalog to "Claude Sonnet 5.0".
67+
gracekelly_model: str = "claude-sonnet-5"
68+
gracekelly_base_url: str = "http://127.0.0.1:8011"
69+
gracekelly_timeout_seconds: float = 180.0
70+
6271
mistral_api_key: str = Field(default="", validation_alias="MISTRAL_API_KEY")
6372
github_token: str = Field(default="", validation_alias="GITHUB_TOKEN")
6473
groq_api_key: str = Field(default="", validation_alias="GROQ_API_KEY")

‎src/nl_sql/llm/providers/factory.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from nl_sql.config import Settings, get_settings
66
from nl_sql.llm.providers.base import LLMProvider, ProviderError
77
from nl_sql.llm.providers.github_models import GitHubModelsProvider
8+
from nl_sql.llm.providers.gracekelly import GraceKellyProvider
89
from nl_sql.llm.providers.groq import GroqProvider
910
from nl_sql.llm.providers.mistral import MistralProvider
1011
from nl_sql.llm.providers.ollama import OllamaProvider
@@ -16,11 +17,17 @@ def build_provider(name: str, settings: Settings | None = None) -> LLMProvider:
1617
"""Build an LLMProvider by short name.
1718
1819
Recognized names: ``mistral``, ``github_models``, ``groq``, ``ollama``,
19-
``perplexity``, ``openrouter``. Raises ProviderError for unknown names
20-
or missing credentials.
20+
``perplexity``, ``openrouter``, ``gracekelly``. Raises ProviderError for
21+
unknown names or missing credentials.
2122
"""
2223
s = settings or get_settings()
2324
match name:
25+
case "gracekelly":
26+
return GraceKellyProvider(
27+
model=s.gracekelly_model,
28+
base_url=s.gracekelly_base_url,
29+
timeout_seconds=s.gracekelly_timeout_seconds,
30+
)
2431
case "mistral":
2532
return MistralProvider(
2633
api_key=s.mistral_api_key,
Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
"""GraceKelly orchestrate provider — routes generation through a local
2+
GraceKelly instance that drives a frontier model via its browser adapter.
3+
4+
This is the AgentFlow-branch model slot (see AgentFlow ADR 0008): AgentFlow
5+
adopts this engine but runs NL->SQL on **Claude Sonnet 5 via GraceKelly**, not
6+
Mistral. The portfolio NL_SQL default provider stays `mistral`; this provider is
7+
selected explicitly (`build_provider("gracekelly")`).
8+
9+
Difference from `PerplexityProvider`: that one predates the V2 API and posts to
10+
`/api/v1/pipeline` reading `answer`. This provider posts to `/api/v1/orchestrate`
11+
reading `output_text`, which is the contract AgentFlow's serving layer already
12+
standardised on (`serving/semantic_layer/nl_engine._llm_translate`). GraceKelly
13+
owns model selection/execution — this provider never talks to a model API
14+
directly. The default model `claude-sonnet-5` resolves through GraceKelly's live
15+
catalog to "Claude Sonnet 5.0".
16+
17+
Latency is ~20-40s per call (browser path), so this provider is for evaluation
18+
runs and one-off probes, not an interactive chat surface.
19+
"""
20+
21+
from __future__ import annotations
22+
23+
import json
24+
import time
25+
from urllib import error as urlerror
26+
from urllib import request as urlrequest
27+
28+
from nl_sql.llm.providers.base import (
29+
GenerateRequest,
30+
GenerateResponse,
31+
ProviderError,
32+
)
33+
34+
# Reuse the hardened Perplexity/GraceKelly output parsing: strip terminal ANSI
35+
# codes the Markdown pipeline can leave behind, and unwrap the
36+
# `{"sql": "...", "rationale": "..."}` envelope Sonnet sometimes emits instead
37+
# of a bare statement. Same browser backend => same output quirks.
38+
from nl_sql.llm.providers.perplexity import _ANSI_RE, _unwrap_sql_json
39+
40+
41+
class GraceKellyProvider:
42+
"""LLMProvider that proxies generate() to a local GraceKelly orchestrate API."""
43+
44+
name: str = "gracekelly"
45+
46+
def __init__(
47+
self,
48+
*,
49+
model: str = "claude-sonnet-5",
50+
base_url: str = "http://127.0.0.1:8011",
51+
timeout_seconds: float = 180.0,
52+
) -> None:
53+
if not model.strip():
54+
raise ProviderError("GraceKellyProvider requires non-empty model")
55+
self.model = model
56+
self._base_url = base_url.rstrip("/")
57+
self._timeout = timeout_seconds
58+
59+
def generate(self, req: GenerateRequest) -> GenerateResponse:
60+
prompt = req.prompt
61+
if req.system:
62+
prompt = f"{req.system}\n\n{prompt}"
63+
payload = json.dumps({"prompt": prompt, "model": self.model}).encode("utf-8")
64+
65+
http_request = urlrequest.Request(
66+
f"{self._base_url}/api/v1/orchestrate",
67+
data=payload,
68+
headers={"Content-Type": "application/json", "Accept": "application/json"},
69+
method="POST",
70+
)
71+
t0 = time.perf_counter()
72+
try:
73+
with urlrequest.urlopen(http_request, timeout=self._timeout) as response:
74+
body = response.read().decode("utf-8")
75+
except urlerror.HTTPError as exc:
76+
raise ProviderError(
77+
f"GraceKelly /api/v1/orchestrate returned {exc.code}: "
78+
f"{exc.read().decode('utf-8', errors='replace')[:400]}"
79+
) from exc
80+
except (urlerror.URLError, TimeoutError) as exc:
81+
raise ProviderError(
82+
f"GraceKelly unreachable at {self._base_url}: {exc!r}. "
83+
"Start it with `python -m uvicorn gracekelly.main:create_app "
84+
"--factory --host 127.0.0.1 --port 8011` "
85+
"and set GRACEKELLY_EXECUTION_PROFILE=hybrid."
86+
) from exc
87+
88+
elapsed_ms = (time.perf_counter() - t0) * 1000.0
89+
parsed = json.loads(body)
90+
91+
failure = parsed.get("failure_message") or parsed.get("failure_code")
92+
if failure and not parsed.get("output_text"):
93+
raise ProviderError(f"GraceKelly orchestrate failed: {failure}")
94+
95+
answer = _ANSI_RE.sub("", str(parsed.get("output_text") or ""))
96+
answer = _unwrap_sql_json(answer)
97+
98+
model_field = parsed.get("model")
99+
model_id = model_field.get("id") if isinstance(model_field, dict) else model_field
100+
# Browser path does not surface token counts; use a word-count proxy so
101+
# eval reports show something plausible without faking billing units.
102+
approx_in = max(1, len(prompt.split()))
103+
approx_out = max(1, len(answer.split()))
104+
return GenerateResponse(
105+
text=answer,
106+
model=str(model_id or self.model),
107+
input_tokens=approx_in,
108+
output_tokens=approx_out,
109+
latency_ms=elapsed_ms,
110+
)

‎tests/test_gracekelly_provider.py‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
"""Unit tests for the GraceKelly orchestrate provider (AgentFlow ADR 0008 slot).
2+
3+
Mock the HTTP layer so these run without a live GraceKelly server; a live smoke
4+
is done separately against a running instance.
5+
"""
6+
7+
from __future__ import annotations
8+
9+
import io
10+
import json
11+
from typing import Any
12+
from unittest import mock
13+
from urllib import error as urlerror
14+
15+
import pytest
16+
17+
from nl_sql.llm.providers import GenerateRequest, ProviderError, build_provider
18+
from nl_sql.llm.providers.gracekelly import GraceKellyProvider
19+
20+
21+
class _FakeResponse(io.BytesIO):
22+
def __enter__(self) -> _FakeResponse:
23+
return self
24+
25+
def __exit__(self, *exc: object) -> None:
26+
self.close()
27+
28+
29+
def _orchestrate_body(**overrides: Any) -> bytes:
30+
payload = {
31+
"status": "completed",
32+
"output_text": "SELECT COUNT(*) FROM orders_v2 WHERE status = 'cancelled';",
33+
"model": {"id": "claude-sonnet-5-0", "display_name": "Claude Sonnet 5.0"},
34+
"failure_message": None,
35+
}
36+
payload.update(overrides)
37+
return json.dumps(payload).encode("utf-8")
38+
39+
40+
def test_factory_builds_gracekelly_with_sonnet_5() -> None:
41+
provider = build_provider("gracekelly")
42+
assert isinstance(provider, GraceKellyProvider)
43+
assert provider.name == "gracekelly"
44+
assert provider.model == "claude-sonnet-5"
45+
46+
47+
def test_generate_posts_orchestrate_and_parses_output_text() -> None:
48+
captured: dict[str, Any] = {}
49+
50+
def fake_urlopen(req: Any, timeout: float | None = None) -> _FakeResponse:
51+
captured["url"] = req.full_url
52+
captured["body"] = json.loads(req.data.decode("utf-8"))
53+
return _FakeResponse(_orchestrate_body())
54+
55+
provider = GraceKellyProvider(model="claude-sonnet-5", base_url="http://127.0.0.1:8011")
56+
with mock.patch("nl_sql.llm.providers.gracekelly.urlrequest.urlopen", fake_urlopen):
57+
resp = provider.generate(GenerateRequest(prompt="how many orders were cancelled?"))
58+
59+
assert captured["url"] == "http://127.0.0.1:8011/api/v1/orchestrate"
60+
assert captured["body"] == {
61+
"prompt": "how many orders were cancelled?",
62+
"model": "claude-sonnet-5",
63+
}
64+
assert resp.text == "SELECT COUNT(*) FROM orders_v2 WHERE status = 'cancelled';"
65+
assert resp.model == "claude-sonnet-5-0"
66+
67+
68+
def test_generate_prepends_system_prompt() -> None:
69+
captured: dict[str, Any] = {}
70+
71+
def fake_urlopen(req: Any, timeout: float | None = None) -> _FakeResponse:
72+
captured["body"] = json.loads(req.data.decode("utf-8"))
73+
return _FakeResponse(_orchestrate_body())
74+
75+
provider = GraceKellyProvider()
76+
with mock.patch("nl_sql.llm.providers.gracekelly.urlrequest.urlopen", fake_urlopen):
77+
provider.generate(
78+
GenerateRequest(system="You are a DuckDB SQL generator.", prompt="revenue?")
79+
)
80+
81+
assert captured["body"]["prompt"] == "You are a DuckDB SQL generator.\n\nrevenue?"
82+
83+
84+
def test_generate_unwraps_sql_json_envelope() -> None:
85+
envelope = json.dumps({"sql": "SELECT 1", "rationale": "trivial"})
86+
87+
def fake_urlopen(req: Any, timeout: float | None = None) -> _FakeResponse:
88+
return _FakeResponse(_orchestrate_body(output_text=envelope))
89+
90+
provider = GraceKellyProvider()
91+
with mock.patch("nl_sql.llm.providers.gracekelly.urlrequest.urlopen", fake_urlopen):
92+
resp = provider.generate(GenerateRequest(prompt="q"))
93+
94+
assert resp.text == "SELECT 1"
95+
96+
97+
def test_generate_raises_on_orchestrate_failure() -> None:
98+
def fake_urlopen(req: Any, timeout: float | None = None) -> _FakeResponse:
99+
return _FakeResponse(_orchestrate_body(output_text="", failure_message="browser timeout"))
100+
101+
provider = GraceKellyProvider()
102+
with (
103+
mock.patch("nl_sql.llm.providers.gracekelly.urlrequest.urlopen", fake_urlopen),
104+
pytest.raises(ProviderError, match="browser timeout"),
105+
):
106+
provider.generate(GenerateRequest(prompt="q"))
107+
108+
109+
def test_generate_wraps_http_error() -> None:
110+
def fake_urlopen(req: Any, timeout: float | None = None) -> _FakeResponse:
111+
raise urlerror.HTTPError(req.full_url, 422, "Unsupported model", {}, io.BytesIO(b"nope")) # type: ignore[arg-type]
112+
113+
provider = GraceKellyProvider()
114+
with (
115+
mock.patch("nl_sql.llm.providers.gracekelly.urlrequest.urlopen", fake_urlopen),
116+
pytest.raises(ProviderError, match="422"),
117+
):
118+
provider.generate(GenerateRequest(prompt="q"))
119+
120+
121+
def test_empty_model_rejected() -> None:
122+
with pytest.raises(ProviderError):
123+
GraceKellyProvider(model=" ")

0 commit comments

Comments
 (0)