Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 2 additions & 21 deletions tests/ai_test_model.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,4 @@
import collections.abc
from typing import override

import httpx
from httpx import Auth, Request, Response
from pydantic import BaseModel

from splunklib.ai import OpenAIModel
Expand Down Expand Up @@ -37,20 +33,6 @@ async def create_model(s: TestLLMSettings) -> PredefinedModel:
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

Expand Down Expand Up @@ -79,14 +61,13 @@ async def _buildInternalAIModel(

token = _TokenResponse.model_validate_json(response.text).access_token

auth_handler = _InternalAIAuth(token)
model = "gpt-5-nano"

return OpenAIModel(
model=model,
base_url=f"{base_url}/{model}",
api_key="", # unused
api_key="test-api-key", # unused
extra_body={"user": f'{{"appkey":"{app_key}"}}'},
httpx_client=httpx.AsyncClient(auth=auth_handler),
httpx_client=httpx.AsyncClient(headers={"api-key": token}),
temperature=0.0,
)