diff --git a/posthog/egress/test/test_typesafe.py b/posthog/egress/test/test_typesafe.py index 8d8156fd1738..81a3f247029c 100644 --- a/posthog/egress/test/test_typesafe.py +++ b/posthog/egress/test/test_typesafe.py @@ -10,17 +10,9 @@ from prometheus_client import REGISTRY from posthog.egress.limiter.policies import Priority, resolve_policy -from posthog.egress.typesafe.client import ( - ChoiceAnswer, - ChoiceQuestion, - NoulAnswer, - NoulQuestion, - Question, - TypeSafeNotConfigured, - TypeSafeRequestFailed, - system_one, -) +from posthog.egress.typesafe.client import TypeSafeNotConfigured, TypeSafeRequestFailed, system_one from posthog.egress.typesafe.limiter import typesafe_account_key +from posthog.llm.system_one import ChoiceAnswer, ChoiceQuestion, NoulAnswer, NoulQuestion, Question _FAKE_API_KEY = "fake-key-for-tests" diff --git a/posthog/egress/typesafe/README.md b/posthog/egress/typesafe/README.md index b4d7490d54fd..3d0970abccd0 100644 --- a/posthog/egress/typesafe/README.md +++ b/posthog/egress/typesafe/README.md @@ -3,6 +3,11 @@ TypeSafe serves Jev, a System One model that answers typed questions (`noul`, `choice`, `score`) about a `state`. It returns probabilities, not free text. +The request and answer types live in `posthog/llm/system_one.py`, because the Go ai-gateway serves the same API with models PostHog hosts. +A caller builds its client with `build_system_one_client` from `posthog/llm/system_one_client.py`. +It reaches the ai-gateway where `AI_GATEWAY_URL` is set. +It falls back to this domain only when the caller passes a `TypeSafeFallback`, and every caller that does so meets the usage policy below. + ## Usage policy TypeSafe is approved for experiments only. diff --git a/posthog/egress/typesafe/__init__.py b/posthog/egress/typesafe/__init__.py index e4375fc81b42..4aae693df6d9 100644 --- a/posthog/egress/typesafe/__init__.py +++ b/posthog/egress/typesafe/__init__.py @@ -1,29 +1,8 @@ -from posthog.egress.typesafe.client import ( - JEV_LATEST, - Answer, - ChoiceAnswer, - ChoiceQuestion, - JsonValue, - NoulAnswer, - NoulQuestion, - Question, - SystemOneResult, - TypeSafeNotConfigured, - TypeSafeRequestFailed, - system_one, -) +from posthog.egress.typesafe.client import JEV_LATEST, TypeSafeNotConfigured, TypeSafeRequestFailed, system_one from posthog.egress.typesafe.transport import TypeSafeEgressBudgetExhausted, typesafe_request __all__ = [ "JEV_LATEST", - "Answer", - "ChoiceAnswer", - "ChoiceQuestion", - "JsonValue", - "NoulAnswer", - "NoulQuestion", - "Question", - "SystemOneResult", "TypeSafeEgressBudgetExhausted", "TypeSafeNotConfigured", "TypeSafeRequestFailed", diff --git a/posthog/egress/typesafe/client.py b/posthog/egress/typesafe/client.py index 5cb358213b7a..e16624de70cf 100644 --- a/posthog/egress/typesafe/client.py +++ b/posthog/egress/typesafe/client.py @@ -1,13 +1,7 @@ """Typed client for TypeSafe's System One endpoint, ``POST /v1/systemone``. -A System One model such as Jev answers typed questions about a ``state`` and returns judgments -rather than text. A noul question returns the probability that a yes/no statement holds. A choice -question picks one option from the criteria the caller supplies, with a probability for every -option. Score questions exist too and stay unwired until a caller needs one, so the request body -remains a checked shape. - -Keep user text in ``state`` and refer to it from the instructions by name. Never interpolate it -into ``instructions``, so user text cannot become an instruction. +The request and answer types live in ``posthog.llm.system_one``, because the Go ai-gateway serves +the same API. This module adds TypeSafe's host, key, and egress budget. USAGE POLICY. TypeSafe is approved for experiments only. Read "Usage policy" in this package's README.md before you add a caller. In short: gate every caller behind a feature flag that reaches @@ -15,150 +9,37 @@ opt-in from each customer and sign-off from leadership first. """ -import math -from collections.abc import Mapping, Sequence -from typing import cast +from collections.abc import Mapping from django.conf import settings -from posthog.dataclasses import frozen from posthog.egress.limiter.policies import Priority from posthog.egress.typesafe.transport import DEFAULT_TIMEOUT, typesafe_request +from posthog.llm.system_one import ( + SYSTEM_ONE_PATH, + JsonValue, + Question, + SystemOneNotConfigured, + SystemOneRequestFailed, + SystemOneResult, + build_system_one_body, + parse_system_one_response, +) TYPESAFE_API_BASE = "https://api.typesafe.ai" -SYSTEM_ONE_ENDPOINT = "/v1/systemone" +SYSTEM_ONE_ENDPOINT = SYSTEM_ONE_PATH # The alias moves to each new release. A caller that tunes thresholds against one version pins that # version's id instead, so a release cannot shift its answers without a code change. JEV_LATEST = "jev-latest" -# TypeSafe rejects a choice question with more options than this. -MAX_CHOICE_OPTIONS = 255 - -type JsonValue = str | int | float | bool | None | Mapping[str, JsonValue] | Sequence[JsonValue] - - -class TypeSafeNotConfigured(Exception): - """No TypeSafe API key is configured on this instance, so no call was made. Self-hosted - deployments run without one, so callers must treat this as a normal degraded path.""" - - -class TypeSafeRequestFailed(Exception): - """TypeSafe was reached but did not return a usable answer (HTTP error, or a body that does not - match the documented shape). ``status_code`` is set for an HTTP error, so a caller can defer a - 429 or 529 and drop the rest.""" - - def __init__(self, message: str, *, status_code: int | None = None) -> None: - super().__init__(message) - self.status_code = status_code - - -@frozen -class NoulQuestion: - """A yes/no question. ``criteria_true`` and ``criteria_false`` optionally describe what each answer means.""" - - instructions: JsonValue - criteria_true: JsonValue = None - criteria_false: JsonValue = None - - def to_json(self) -> dict[str, JsonValue]: - body: dict[str, JsonValue] = {"type": "noul", "instructions": self.instructions} - criteria = { - key: value - for key, value in (("true", self.criteria_true), ("false", self.criteria_false)) - if value is not None - } - if criteria: - body["criteria"] = criteria - return body - - -@frozen -class ChoiceQuestion: - """A pick of one option. ``criteria`` maps each option to its description, or to ``None`` when - the option name is clear on its own.""" - - instructions: JsonValue - criteria: Mapping[str, JsonValue] - - def __post_init__(self) -> None: - if not 1 <= len(self.criteria) <= MAX_CHOICE_OPTIONS: - raise ValueError(f"A choice question needs between 1 and {MAX_CHOICE_OPTIONS} options") - - def to_json(self) -> dict[str, JsonValue]: - return {"type": "choice", "instructions": self.instructions, "criteria": dict(self.criteria)} +class TypeSafeNotConfigured(SystemOneNotConfigured): + """No TypeSafe API key is configured on this instance, so no call was made.""" -type Question = NoulQuestion | ChoiceQuestion - -@frozen -class NoulAnswer: - probability: float - - -@frozen -class ChoiceAnswer: - choice: str - confidence: float - probabilities: Mapping[str, float] - - -type Answer = NoulAnswer | ChoiceAnswer - - -@frozen -class SystemOneResult: - """The answers to one request, keyed by the question ids the caller chose. ``model`` is the - versioned id that answered, also when the request named an alias.""" - - model: str - answers: Mapping[str, Answer] - input_tokens: int | None - - -def _as_mapping(value: object) -> Mapping[str, object] | None: - # An isinstance check cannot narrow the key type, but these mappings come from a parsed JSON - # body, where every key is a string. - return cast(Mapping[str, object], value) if isinstance(value, Mapping) else None - - -def _as_probability(value: object) -> float | None: - if isinstance(value, bool) or not isinstance(value, int | float): - return None - try: - probability = float(value) - except OverflowError: - # JSON integers have no size limit, and one past the float range cannot be a probability. - return None - if not math.isfinite(probability) or not 0 <= probability <= 1: - return None - return probability - - -def _parse_answer(question_id: str, question: Question, raw: object) -> Answer: - answer = _as_mapping(raw) - if answer is None: - raise TypeSafeRequestFailed(f"TypeSafe returned no answer for {question_id!r}") - - if isinstance(question, NoulQuestion): - probability = _as_probability(answer.get("noul")) - if answer.get("type") != "noul" or probability is None: - raise TypeSafeRequestFailed(f"TypeSafe returned a malformed noul for {question_id!r}") - return NoulAnswer(probability=probability) - - choice = answer.get("choice") - confidence = _as_probability(answer.get("confidence")) - if answer.get("type") != "choice" or not isinstance(choice, str) or choice not in question.criteria: - raise TypeSafeRequestFailed(f"TypeSafe returned a malformed choice for {question_id!r}") - if confidence is None: - raise TypeSafeRequestFailed(f"TypeSafe returned a malformed confidence for {question_id!r}") - raw_probabilities = _as_mapping(answer.get("probabilities")) or {} - probabilities = {option: _as_probability(raw_probabilities.get(option)) for option in question.criteria} - complete = {option: probability for option, probability in probabilities.items() if probability is not None} - if len(complete) != len(question.criteria): - raise TypeSafeRequestFailed(f"TypeSafe returned incomplete probabilities for {question_id!r}") - return ChoiceAnswer(choice=choice, confidence=confidence, probabilities=complete) +class TypeSafeRequestFailed(SystemOneRequestFailed): + """TypeSafe was reached but did not return a usable answer.""" def system_one( @@ -193,11 +74,7 @@ def system_one( priority=priority, timeout=timeout, allow_redirects=False, - json={ - "model": model, - "state": state, - "questions": {question_id: question.to_json() for question_id, question in questions.items()}, - }, + json=build_system_one_body(state=state, questions=questions, model=model), ) if not response.ok: @@ -208,22 +85,7 @@ def system_one( payload: object = response.json() except ValueError as exc: raise TypeSafeRequestFailed("TypeSafe returned a non-JSON body") from exc - - body = _as_mapping(payload) - raw_answers = _as_mapping(body.get("answers")) if body is not None else None - if body is None or raw_answers is None: - raise TypeSafeRequestFailed("TypeSafe returned no answers") - answered_model = body.get("model") - if not isinstance(answered_model, str) or not answered_model: - raise TypeSafeRequestFailed("TypeSafe returned no model") - - usage = _as_mapping(body.get("usage")) or {} - input_tokens = usage.get("input_tokens") - return SystemOneResult( - model=answered_model, - answers={ - question_id: _parse_answer(question_id, question, raw_answers.get(question_id)) - for question_id, question in questions.items() - }, - input_tokens=input_tokens if isinstance(input_tokens, int) and not isinstance(input_tokens, bool) else None, - ) + try: + return parse_system_one_response(payload, questions) + except SystemOneRequestFailed as exc: + raise TypeSafeRequestFailed(str(exc)) from exc diff --git a/posthog/llm/system_one.py b/posthog/llm/system_one.py new file mode 100644 index 000000000000..999df9043dc1 --- /dev/null +++ b/posthog/llm/system_one.py @@ -0,0 +1,187 @@ +"""Wire types for the System One API, ``POST /v1/systemone``, shared by every server that speaks it. + +A System One model such as Jev answers typed questions about a ``state`` and returns judgments +rather than text. A noul question returns the probability that a yes/no statement holds. A choice +question picks one option from the criteria the caller supplies, with a probability for every +option. Score questions exist too and stay unwired until a caller needs one, so the request body +remains a checked shape. + +Keep user text in ``state`` and refer to it from the instructions by name. Never interpolate it +into ``instructions``, so user text cannot become an instruction. + +TypeSafe (``posthog.egress.typesafe``) and the Go ai-gateway both serve this API. Callers pick one +through ``posthog.llm.system_one_client.build_system_one_client``. +""" + +import math +from collections.abc import Mapping, Sequence +from typing import cast + +from posthog.dataclasses import frozen + +SYSTEM_ONE_PATH = "/v1/systemone" + +# The most options a System One server takes in one choice question. A model can take fewer. +MAX_CHOICE_OPTIONS = 255 + +type JsonValue = str | int | float | bool | None | Mapping[str, JsonValue] | Sequence[JsonValue] + + +class SystemOneNotConfigured(Exception): + """No System One server is configured on this instance, so no call was made. Self-hosted + deployments run without one, so callers must treat this as a normal degraded path.""" + + +class SystemOneRequestFailed(Exception): + """A System One server was reached but did not return a usable answer (HTTP error, or a body that + does not match the documented shape). ``status_code`` is set for an HTTP error, so a caller can + defer a 429 or 529 and drop the rest.""" + + def __init__(self, message: str, *, status_code: int | None = None) -> None: + super().__init__(message) + self.status_code = status_code + + +@frozen +class NoulQuestion: + """A yes/no question. ``criteria_true`` and ``criteria_false`` optionally describe what each answer means.""" + + instructions: JsonValue + criteria_true: JsonValue = None + criteria_false: JsonValue = None + + def to_json(self) -> dict[str, JsonValue]: + body: dict[str, JsonValue] = {"type": "noul", "instructions": self.instructions} + criteria = { + key: value + for key, value in (("true", self.criteria_true), ("false", self.criteria_false)) + if value is not None + } + if criteria: + body["criteria"] = criteria + return body + + +@frozen +class ChoiceQuestion: + """A pick of one option. ``criteria`` maps each option to its description, or to ``None`` when + the option name is clear on its own.""" + + instructions: JsonValue + criteria: Mapping[str, JsonValue] + + def __post_init__(self) -> None: + if not 1 <= len(self.criteria) <= MAX_CHOICE_OPTIONS: + raise ValueError(f"A choice question needs between 1 and {MAX_CHOICE_OPTIONS} options") + + def to_json(self) -> dict[str, JsonValue]: + return {"type": "choice", "instructions": self.instructions, "criteria": dict(self.criteria)} + + +type Question = NoulQuestion | ChoiceQuestion + + +@frozen +class NoulAnswer: + probability: float + + +@frozen +class ChoiceAnswer: + choice: str + confidence: float + probabilities: Mapping[str, float] + + +type Answer = NoulAnswer | ChoiceAnswer + + +@frozen +class SystemOneResult: + """The answers to one request, keyed by the question ids the caller chose. ``model`` is the + versioned id that answered, also when the request named an alias.""" + + model: str + answers: Mapping[str, Answer] + input_tokens: int | None + + +def _as_mapping(value: object) -> Mapping[str, object] | None: + # An isinstance check cannot narrow the key type, but these mappings come from a parsed JSON + # body, where every key is a string. + return cast(Mapping[str, object], value) if isinstance(value, Mapping) else None + + +def _as_probability(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + try: + probability = float(value) + except OverflowError: + # JSON integers have no size limit, and one past the float range cannot be a probability. + return None + if not math.isfinite(probability) or not 0 <= probability <= 1: + return None + return probability + + +def _parse_answer(question_id: str, question: Question, raw: object) -> Answer: + answer = _as_mapping(raw) + if answer is None: + raise SystemOneRequestFailed(f"The System One server returned no answer for {question_id!r}") + + if isinstance(question, NoulQuestion): + probability = _as_probability(answer.get("noul")) + # A server can leave out `type`, because the question already says which answer it expects. + if answer.get("type", "noul") != "noul" or probability is None: + raise SystemOneRequestFailed(f"The System One server returned a malformed noul for {question_id!r}") + return NoulAnswer(probability=probability) + + choice = answer.get("choice") + confidence = _as_probability(answer.get("confidence")) + if answer.get("type", "choice") != "choice" or not isinstance(choice, str) or choice not in question.criteria: + raise SystemOneRequestFailed(f"The System One server returned a malformed choice for {question_id!r}") + if confidence is None: + raise SystemOneRequestFailed(f"The System One server returned a malformed confidence for {question_id!r}") + raw_probabilities = _as_mapping(answer.get("probabilities")) or {} + probabilities = {option: _as_probability(raw_probabilities.get(option)) for option in question.criteria} + complete = {option: probability for option, probability in probabilities.items() if probability is not None} + if len(complete) != len(question.criteria): + raise SystemOneRequestFailed(f"The System One server returned incomplete probabilities for {question_id!r}") + return ChoiceAnswer(choice=choice, confidence=confidence, probabilities=complete) + + +def build_system_one_body( + *, state: JsonValue, questions: Mapping[str, Question], model: str | None +) -> dict[str, JsonValue]: + """The System One request body. ``model`` is left out when ``None``, so the server picks its default.""" + body: dict[str, JsonValue] = {} if model is None else {"model": model} + body["state"] = state + body["questions"] = {question_id: question.to_json() for question_id, question in questions.items()} + return body + + +def parse_system_one_response(payload: object, questions: Mapping[str, Question]) -> SystemOneResult: + """Validate a parsed System One JSON body against the questions that were asked. + + Raises :class:`SystemOneRequestFailed` unless the body holds a model id and a complete, + well-formed answer to every question. + """ + body = _as_mapping(payload) + raw_answers = _as_mapping(body.get("answers")) if body is not None else None + if body is None or raw_answers is None: + raise SystemOneRequestFailed("The System One server returned no answers") + answered_model = body.get("model") + if not isinstance(answered_model, str) or not answered_model: + raise SystemOneRequestFailed("The System One server returned no model") + + usage = _as_mapping(body.get("usage")) or {} + input_tokens = usage.get("input_tokens") + return SystemOneResult( + model=answered_model, + answers={ + question_id: _parse_answer(question_id, question, raw_answers.get(question_id)) + for question_id, question in questions.items() + }, + input_tokens=input_tokens if isinstance(input_tokens, int) and not isinstance(input_tokens, bool) else None, + ) diff --git a/posthog/llm/system_one_client.py b/posthog/llm/system_one_client.py new file mode 100644 index 000000000000..8b4452195eb2 --- /dev/null +++ b/posthog/llm/system_one_client.py @@ -0,0 +1,171 @@ +"""Picks the System One server a caller reaches, the way ``build_openai_client`` picks a chat gateway. + +The Go ai-gateway serves System One models that PostHog hosts, and bills the wallet of the team that +owns ``AI_GATEWAY_API_KEY``. TypeSafe serves Jev, but a caller reaches it only by passing a +``TypeSafeFallback``: TypeSafe is a third party, approved for experiments that send no customer data +(see ``posthog/egress/typesafe/README.md``). The two serve different models, so the fallback names +its own, and the result says which model answered. +""" + +from collections.abc import Mapping +from dataclasses import field +from urllib.parse import urlparse, urlunparse + +from django.conf import settings + +import httpx +import structlog + +from posthog.dataclasses import frozen +from posthog.egress.limiter.policies import Priority +from posthog.egress.typesafe.client import system_one +from posthog.llm.gateway_client import AIGatewayConfig, ai_gateway_headers, resolve_ai_gateway_config +from posthog.llm.system_one import ( + SYSTEM_ONE_PATH, + ChoiceQuestion, + JsonValue, + Question, + SystemOneNotConfigured, + SystemOneRequestFailed, + SystemOneResult, + build_system_one_body, + parse_system_one_response, +) + +logger = structlog.get_logger(__name__) + +DEFAULT_TIMEOUT_SECONDS = 30.0 + +# The decision models the gateway serves (JevK5) answer with one letter per option, A to P. +GATEWAY_MAX_CHOICE_OPTIONS = 16 +GATEWAY_MAX_QUESTIONS = 32 + + +@frozen +class TypeSafeFallback: + """Where no gateway is configured, ask TypeSafe for ``model`` from the ``source`` egress budget.""" + + model: str + source: str + priority: Priority = Priority.NORMAL + + +@frozen +class GatewaySystemOneClient: + url: str + api_key: str = field(repr=False) + headers: Mapping[str, str] + model: str + timeout: float + + def decide(self, *, state: JsonValue, questions: Mapping[str, Question]) -> SystemOneResult: + if not 1 <= len(questions) <= GATEWAY_MAX_QUESTIONS: + raise ValueError(f"A System One request needs between 1 and {GATEWAY_MAX_QUESTIONS} questions") + for question_id, question in questions.items(): + if isinstance(question, ChoiceQuestion) and len(question.criteria) > GATEWAY_MAX_CHOICE_OPTIONS: + raise ValueError(f"{question_id!r} has more than {GATEWAY_MAX_CHOICE_OPTIONS} options") + try: + with httpx.Client(trust_env=False, timeout=self.timeout) as client: + response = client.post( + self.url, + json=build_system_one_body(state=state, questions=questions, model=self.model), + headers={**self.headers, "Authorization": f"Bearer {self.api_key}"}, + ) + except httpx.HTTPError as exc: + raise SystemOneRequestFailed(f"The ai-gateway was not reached: {exc.__class__.__name__}") from exc + if response.status_code != 200: + # The error body can echo the state, so it stays out of the exception that gets logged. + raise SystemOneRequestFailed( + f"The ai-gateway returned HTTP {response.status_code}", status_code=response.status_code + ) + try: + payload: object = response.json() + except ValueError as exc: + raise SystemOneRequestFailed("The ai-gateway returned a non-JSON body") from exc + return parse_system_one_response(payload, questions) + + +@frozen +class TypeSafeSystemOneClient: + model: str + source: str + priority: Priority + timeout: float + + def decide(self, *, state: JsonValue, questions: Mapping[str, Question]) -> SystemOneResult: + return system_one( + state=state, + questions=questions, + source=self.source, + model=self.model, + priority=self.priority, + timeout=self.timeout, + ) + + +type SystemOneClient = GatewaySystemOneClient | TypeSafeSystemOneClient + + +def _system_one_url(gateway_url: str) -> str: + """The gateway URL setting carries the OpenAI ``/v1`` base path, and System One hangs off the origin.""" + parsed = urlparse(gateway_url) + path = parsed.path.rstrip("/").removesuffix("/v1") + return urlunparse(parsed._replace(path=path + SYSTEM_ONE_PATH, params="", query="", fragment="")) + + +def _usable_gateway() -> AIGatewayConfig | None: + """The gateway config, unless its key would travel in clear to a host off this machine.""" + gateway = resolve_ai_gateway_config() + if gateway is None: + return None + parsed = urlparse(gateway.url) + if parsed.scheme != "https" and parsed.hostname not in {"localhost", "127.0.0.1", "::1"}: + logger.warning("system_one_gateway_url_not_https") + return None + return gateway + + +def system_one_configured(typesafe_fallback: TypeSafeFallback | None = None) -> bool: + if _usable_gateway() is not None: + return True + return typesafe_fallback is not None and bool(settings.TYPESAFE_API_KEY) + + +def build_system_one_client( + *, + model: str, + ai_product: str, + typesafe_fallback: TypeSafeFallback | None = None, + distinct_id: str | None = None, + trace_id: str | None = None, + properties: Mapping[str, str] | None = None, + timeout: float = DEFAULT_TIMEOUT_SECONDS, +) -> SystemOneClient: + """A client for ``model`` on the Go ai-gateway when it is configured, else for TypeSafe when the + caller passes ``typesafe_fallback``. + + ``ai_product``, ``distinct_id``, ``trace_id`` and ``properties`` label the gateway's event. Raises + :class:`SystemOneNotConfigured` when no server the caller allows is configured. + """ + gateway = _usable_gateway() + if gateway is not None: + return GatewaySystemOneClient( + url=_system_one_url(gateway.url), + api_key=gateway.api_key, + headers=ai_gateway_headers( + ai_product=ai_product, trace_id=trace_id, properties=properties, distinct_id=distinct_id + ) + or {}, + model=model, + timeout=timeout, + ) + if typesafe_fallback is None: + raise SystemOneNotConfigured("Configure AI_GATEWAY_URL (https) and AI_GATEWAY_API_KEY") + if not settings.TYPESAFE_API_KEY: + raise SystemOneNotConfigured("Configure AI_GATEWAY_URL and AI_GATEWAY_API_KEY, or TYPESAFE_API_KEY") + return TypeSafeSystemOneClient( + model=typesafe_fallback.model, + source=typesafe_fallback.source, + priority=typesafe_fallback.priority, + timeout=timeout, + ) diff --git a/posthog/llm/test_system_one_client.py b/posthog/llm/test_system_one_client.py new file mode 100644 index 000000000000..b3a03d251f2a --- /dev/null +++ b/posthog/llm/test_system_one_client.py @@ -0,0 +1,170 @@ +import json + +from unittest.mock import patch + +from django.test import SimpleTestCase, override_settings + +import httpx +from parameterized import parameterized + +from posthog.egress.limiter.policies import Priority +from posthog.llm.system_one import ( + ChoiceAnswer, + ChoiceQuestion, + NoulAnswer, + NoulQuestion, + Question, + SystemOneNotConfigured, + SystemOneRequestFailed, +) +from posthog.llm.system_one_client import ( + GatewaySystemOneClient, + SystemOneClient, + TypeSafeFallback, + TypeSafeSystemOneClient, + build_system_one_client, +) + +GATEWAY_MODEL = "posthog/hogference/jevk5-fp8-0.2" +FALLBACK = TypeSafeFallback(model="jev-1.13.0", source="test", priority=Priority.BATCH) +GATEWAY = {"AI_GATEWAY_URL": "https://ai-gateway.example.com/v1", "AI_GATEWAY_API_KEY": "phs_test"} +NOTHING = {"AI_GATEWAY_URL": "", "AI_GATEWAY_API_KEY": "", "TYPESAFE_API_KEY": ""} +QUESTIONS: dict[str, Question] = { + "urgent": NoulQuestion(instructions="Is this urgent?"), + "team": ChoiceQuestion(instructions="Which team handles this?", criteria={"billing": None, "support": None}), +} +ANSWERS = { + "model": GATEWAY_MODEL, + "answers": { + # The hogference server leaves out `type`, and TypeSafe sends it. + "urgent": {"noul": 0.8}, + "team": { + "type": "choice", + "choice": "billing", + "confidence": 0.7, + "probabilities": {"billing": 0.7, "support": 0.3}, + }, + }, + "usage": {"input_tokens": 42}, +} + + +def _build(typesafe_fallback: TypeSafeFallback | None = FALLBACK) -> SystemOneClient: + return build_system_one_client( + model=GATEWAY_MODEL, ai_product="test_product", typesafe_fallback=typesafe_fallback, distinct_id="team-7" + ) + + +class TestBuildSystemOneClient(SimpleTestCase): + @parameterized.expand( + [ + ("gateway_wins", {**GATEWAY, "TYPESAFE_API_KEY": "ts-key"}, GatewaySystemOneClient, GATEWAY_MODEL), + ("typesafe_fallback", {"TYPESAFE_API_KEY": "ts-key"}, TypeSafeSystemOneClient, FALLBACK.model), + ( + "gateway_over_plain_http_falls_back", + { + "AI_GATEWAY_URL": "http://ai-gateway.example.com/v1", + "AI_GATEWAY_API_KEY": "phs_test", + "TYPESAFE_API_KEY": "ts-key", + }, + TypeSafeSystemOneClient, + FALLBACK.model, + ), + ] + ) + def test_picks_the_server_and_its_model( + self, _name: str, configured: dict, kind: type[SystemOneClient], model: str + ) -> None: + with override_settings(**{**NOTHING, **configured}): + client = _build() + + assert isinstance(client, kind) + assert client.model == model + + @parameterized.expand( + [ + ("nothing_configured", NOTHING, FALLBACK), + # A caller that passes no fallback must never reach TypeSafe, even with its key set. + ("typesafe_not_allowed", {**NOTHING, "TYPESAFE_API_KEY": "ts-key"}, None), + ( + "gateway_key_over_plain_http", + {**NOTHING, "AI_GATEWAY_URL": "http://ai-gateway.example.com/v1", "AI_GATEWAY_API_KEY": "phs_test"}, + FALLBACK, + ), + ] + ) + def test_refuses_without_a_usable_server( + self, _name: str, configured: dict, typesafe_fallback: TypeSafeFallback | None + ) -> None: + with override_settings(**configured), self.assertRaises(SystemOneNotConfigured): + _build(typesafe_fallback) + + def test_gateway_request_reaches_the_system_one_route_with_its_labels(self) -> None: + with override_settings(**{**NOTHING, **GATEWAY}): + client = _build() + with patch.object(httpx.Client, "send", return_value=httpx.Response(200, json=ANSWERS)) as send: + result = client.decide(state={"ticket": "Payouts fail"}, questions=QUESTIONS) + + request: httpx.Request = send.call_args.args[0] + assert str(request.url) == "https://ai-gateway.example.com/v1/systemone" + assert request.headers["Authorization"] == "Bearer phs_test" + assert request.headers["X-PostHog-Product"] == "test_product" + assert request.headers["X-PostHog-Distinct-Id"] == "team-7" + assert json.loads(request.content)["model"] == GATEWAY_MODEL + assert result.model == GATEWAY_MODEL + assert result.answers == { + "urgent": NoulAnswer(probability=0.8), + "team": ChoiceAnswer(choice="billing", confidence=0.7, probabilities={"billing": 0.7, "support": 0.3}), + } + + @parameterized.expand( + [ + ("http_error", httpx.Response(404, json={"error": "not found"}), 404), + ("unparseable_answer", httpx.Response(200, json={"model": GATEWAY_MODEL, "answers": {}}), None), + ("unreachable", httpx.ConnectError("refused"), None), + ] + ) + def test_gateway_failures_raise_request_failed(self, _name: str, outcome, status_code: int | None) -> None: + with override_settings(**{**NOTHING, **GATEWAY}): + client = _build() + mock = {"side_effect": outcome} if isinstance(outcome, Exception) else {"return_value": outcome} + with patch.object(httpx.Client, "send", **mock), self.assertRaises(SystemOneRequestFailed) as raised: + client.decide(state="Payouts fail", questions=QUESTIONS) + + assert raised.exception.status_code == status_code + + @parameterized.expand( + [ + ( + "choice_past_the_option_limit", + { + "pick": ChoiceQuestion( + instructions="Which one?", criteria={f"option_{index}": None for index in range(17)} + ) + }, + ), + ( + "past_the_question_limit", + {f"question_{index}": NoulQuestion(instructions="Is it?") for index in range(33)}, + ), + ] + ) + def test_gateway_rejects_an_oversized_request_before_sending( + self, _name: str, questions: dict[str, Question] + ) -> None: + with override_settings(**{**NOTHING, **GATEWAY}): + client = _build() + + with patch.object(httpx.Client, "send") as send, self.assertRaises(ValueError): + client.decide(state="x", questions=questions) + + send.assert_not_called() + + def test_typesafe_fallback_sends_its_model_and_lane(self) -> None: + with override_settings(**{**NOTHING, "TYPESAFE_API_KEY": "ts-key"}): + client = _build() + with patch("posthog.llm.system_one_client.system_one") as system_one: + client.decide(state="x", questions=QUESTIONS) + + kwargs = system_one.call_args.kwargs + assert (kwargs["source"], kwargs["model"], kwargs["priority"]) == ("test", FALLBACK.model, Priority.BATCH)