Run Agentic SDK tests in CI (#35)

This change enables Agentic SDK tests to be run in CI, additionally it
fixes system tests, these were broken, since the message refactor.

Our CI now uses our internal OpenAI-compatible model,
before GA we might also add support for OpenAI/Ollama for CI testing,
such that external developers can also run our test suite.

Author: mateusz834

.env.template

@@ -14,3 +14,9 @@ version=9.0
 #splunkToken="<Bearer-token>"
 # Session key for authentication
 #token="<Session-Key>"
+
+#internal_ai_client_id=""
+#internal_ai_client_secret=""
+#internal_ai_app_key=""
+#internal_ai_token_url=""
+#internal_ai_base_url=""

.github/workflows/test.yml

@@ -19,9 +19,22 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
-        run: python -m pip install . --group test
+        run: python -m pip install '.[openai]' --group test
       - name: Set up .env
         run: cp .env.template .env
+      - name: Write internal AI secrets to .env
+        run: |
+          echo "internal_ai_app_key=$INTERNAL_AI_APP_KEY" >> .env
+          echo "internal_ai_client_id=$INTERNAL_AI_CLIENT_ID" >> .env
+          echo "internal_ai_client_secret=$INTERNAL_AI_CLIENT_SECRET" >> .env
+          echo "internal_ai_token_url=$INTERNAL_AI_TOKEN_URL" >> .env
+          echo "internal_ai_base_url=$INTERNAL_AI_BASE_URL" >> .env
+        env:
+          INTERNAL_AI_APP_KEY: ${{ secrets.INTERNAL_AI_APP_KEY }}
+          INTERNAL_AI_CLIENT_ID: ${{ secrets.INTERNAL_AI_CLIENT_ID }}
+          INTERNAL_AI_CLIENT_SECRET: ${{ secrets.INTERNAL_AI_CLIENT_SECRET }}
+          INTERNAL_AI_TOKEN_URL: ${{ secrets.INTERNAL_AI_TOKEN_URL }}
+          INTERNAL_AI_BASE_URL: ${{ secrets.INTERNAL_AI_BASE_URL }}
       - name: Run unit tests
         run: make test-unit
       - name: Run entire test suite

tests/ai_test_model.py

@@ -0,0 +1,89 @@
+import collections.abc
+from typing import override
+
+import httpx
+from httpx import Auth, Request, Response
+from pydantic import BaseModel
+
+from splunklib.ai import OpenAIModel
+from splunklib.ai.model import PredefinedModel
+
+
+class InternalAIModel(BaseModel):
+    client_id: str
+    client_secret: str
+    app_key: str
+
+    token_url: str
+    base_url: str
+
+
+class TestLLMSettings(BaseModel):
+    # TODO: Currently we only support our internal OpenAI-compatible model,
+    # once we are close to GA we should also support OpenAI and probably Ollama, such
+    # that external developers can also run our test suite suite locally.
+    internal_ai: InternalAIModel | None = None
+
+
+async def create_model(s: TestLLMSettings) -> PredefinedModel:
+    if s.internal_ai is not None:
+        return await _buildInternalAIModel(
+            token_url=s.internal_ai.token_url,
+            base_url=s.internal_ai.base_url,
+            client_id=s.internal_ai.client_id,
+            client_secret=s.internal_ai.client_secret,
+            app_key=s.internal_ai.app_key,
+        )
+    raise Exception("unreachable")
+
+
+class _InternalAIAuth(Auth):
+    token: str
+
+    def __init__(self, token: str) -> None:
+        self.token = token
+
+    @override
+    def auth_flow(
+        self, request: Request
+    ) -> collections.abc.Generator[Request, Response, None]:
+        request.headers["api-key"] = self.token
+        yield request
+
+
+class _TokenResponse(BaseModel):
+    access_token: str
+
+
+async def _buildInternalAIModel(
+    token_url: str,
+    base_url: str,
+    client_id: str,
+    client_secret: str,
+    app_key: str,
+) -> OpenAIModel:
+    headers = {
+        "Accept": "*/*",
+        "Content-Type": "application/x-www-form-urlencoded",
+    }
+
+    http = httpx.AsyncClient()
+    response = await http.post(
+        url=token_url,
+        headers=headers,
+        data={"grant_type": "client_credentials"},
+        auth=(client_id, client_secret),
+    )
+
+    token = _TokenResponse.model_validate_json(response.text).access_token
+
+    auth_handler = _InternalAIAuth(token)
+    model = "gpt-4.1"
+
+    return OpenAIModel(
+        model=model,
+        base_url=f"{base_url}/{model}",
+        api_key="",  # unused
+        extra_body={"user": f'{{"appkey":"{app_key}"}}'},
+        httpx_client=httpx.AsyncClient(auth=auth_handler),
+    )

tests/ai_testlib.py

@@ -0,0 +1,32 @@
+from splunklib.ai.model import PredefinedModel
+from tests.ai_test_model import InternalAIModel, TestLLMSettings, create_model
+from tests.testlib import SDKTestCase
+
+
+class AITestCase(SDKTestCase):
+    _model: PredefinedModel | None = None
+
+    @property
+    def test_llm_settings(self) -> TestLLMSettings:
+        client_id: str = self.opts.kwargs["internal_ai_client_id"]
+        client_secret: str = self.opts.kwargs["internal_ai_client_secret"]
+        app_key: str = self.opts.kwargs["internal_ai_app_key"]
+        token_url: str = self.opts.kwargs["internal_ai_token_url"]
+        base_url: str = self.opts.kwargs["internal_ai_base_url"]
+        return TestLLMSettings(
+            internal_ai=InternalAIModel(
+                client_id=client_id,
+                client_secret=client_secret,
+                app_key=app_key,
+                token_url=token_url,
+                base_url=base_url,
+            )
+        )
+
+    async def model(self) -> PredefinedModel:
+        if self._model is not None:
+            return self._model
+
+        model = await create_model(self.test_llm_settings)
+        self._model = model
+        return model

tests/cretestlib.py tests/cre_testlib.pytests/cretestlib.py renamed to tests/cre_testlib.py

@@ -19,14 +19,17 @@
 from abc import abstractmethod
 from http.cookies import SimpleCookie
 
+from splunklib.ai.model import PredefinedModel
 from splunklib.binding import _spliturl
 from splunklib.client import Service, connect
+from tests.ai_test_model import TestLLMSettings, create_model
 
 try:
     import splunk
 
     class CRETestHandler(splunk.rest.BaseRestHandler):
         _service: Service | None = None
+        _model: PredefinedModel | None = None
 
         def handle_POST(self) -> None:
             async def run() -> None:
@@ -45,6 +48,16 @@ async def run() -> None:
         @abstractmethod
         async def run(self) -> None: ...
 
+        async def model(self) -> PredefinedModel:
+            if self._model is not None:
+                return self._model
+
+            raw_body = str(self.request["payload"])
+            s = TestLLMSettings.model_validate_json(raw_body)
+            model = await create_model(s)
+            self._model = model
+            return model
+
         @property
         def service(self) -> Service:
             if self._service is not None: