Skip to content

Commit 557ec3e

Browse files
romartinromartin
committed
AAP-58794: Enable Lightspeed Service's chatbot API Key authentication.
Signed-off-by: romartin <roger600@gmail.com>
1 parent 288d464 commit 557ec3e

File tree

5 files changed

+114
-2
lines changed

5 files changed

+114
-2
lines changed

‎ansible_ai_connect/ai/api/model_pipelines/http/pipelines.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,11 @@ def self_test(self) -> HealthCheckSummary:
146146
}
147147
)
148148
try:
149-
headers = {"Content-Type": "application/json"}
149+
headers = {
150+
"Content-Type": "application/json",
151+
}
152+
if settings.CHATBOT_API_KEY is not None:
153+
headers["Authorization"] = f"Bearer {settings.CHATBOT_API_KEY}"
150154
r = self.session.get(
151155
self.config.inference_url + "/readiness",
152156
headers=headers,
@@ -213,6 +217,8 @@ def invoke(self, params: ChatBotParameters) -> ChatBotResponse:
213217
data["no_tools"] = bool(no_tools)
214218

215219
headers = self.headers or {}
220+
if params.auth_header:
221+
headers["Authorization"] = params.auth_header
216222
if params.mcp_headers:
217223
headers["MCP-HEADERS"] = json.dumps(params.mcp_headers)
218224

@@ -323,7 +329,8 @@ async def async_invoke(self, params: StreamingChatBotParameters) -> AsyncGenerat
323329
"Content-Type": "application/json",
324330
"Accept": "application/json,text/event-stream",
325331
}
326-
332+
if params.auth_header:
333+
headers["Authorization"] = params.auth_header
327334
if params.mcp_headers:
328335
headers["MCP-HEADERS"] = json.dumps(params.mcp_headers)
329336

‎ansible_ai_connect/ai/api/model_pipelines/http/tests/test_http_chatbot_pipeline.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -387,3 +387,94 @@ def test_invoke_with_no_tools_true(self):
387387
json_data = call_args[1]["json"]
388388
self.assertIn("no_tools", json_data)
389389
self.assertTrue(json_data["no_tools"])
390+
391+
392+
@override_settings(CHATBOT_DEFAULT_SYSTEM_PROMPT="You are a helpful assistant")
393+
class TestHttpChatBotPipelineAuthorizationHeader(WisdomServiceLogAwareTestCase):
394+
"""
395+
Test HTTP ChatBot Pipeline's Authorization header.
396+
"""
397+
398+
def setUp(self):
399+
super().setUp()
400+
# Use HTTPS configuration consistent with SSL/TLS enablement PR
401+
config = mock_pipeline_config(
402+
"http", inference_url="https://example.com:8443", verify_ssl=True, ca_cert_file=None
403+
)
404+
assert isinstance(config, HttpConfiguration)
405+
self.config = config
406+
self.pipeline = HttpChatBotPipeline(self.config)
407+
408+
# Mock the session to prevent actual HTTP calls
409+
self.pipeline.session = Mock()
410+
411+
@staticmethod
412+
def get_params(auth_header=None) -> ChatBotParameters:
413+
"""Helper to create test parameters"""
414+
return ChatBotParameters(
415+
query="Hello, how are you?",
416+
conversation_id="test-conversation-123",
417+
provider="test-provider",
418+
model_id="test-model",
419+
system_prompt="You are a helpful assistant",
420+
no_tools=False,
421+
mcp_headers=None,
422+
auth_header=auth_header,
423+
)
424+
425+
def test_invoke_with_auth_header(self):
426+
"""Test that invoke correctly includes Authorization header in the HTTP request"""
427+
auth_header = {"Authentication": "***"}
428+
429+
response_data = {
430+
"response": "Hello",
431+
"truncated": False,
432+
"referenced_documents": [],
433+
}
434+
435+
self.pipeline.session.post.return_value = MockResponse(response_data, 200)
436+
437+
params = self.get_params(auth_header=auth_header)
438+
result = self.pipeline.invoke(params)
439+
440+
# Verify the response
441+
self.assertEqual(result["response"], "Hello")
442+
443+
# Verify the HTTP call was made with correct parameters
444+
self.pipeline.session.post.assert_called_once()
445+
call_args = self.pipeline.session.post.call_args
446+
447+
# Check that Authorization header are included in the request headers
448+
headers = call_args[1]["headers"]
449+
self.assertEqual(headers["Content-Type"], "application/json")
450+
self.assertIn("Authorization", headers)
451+
452+
# Verify the Authorization header are correctly JSON encoded
453+
auth_header_in_request = headers["Authorization"]
454+
self.assertEqual(auth_header_in_request, auth_header)
455+
456+
457+
def test_invoke_with_no_auth_header(self):
458+
"""Test that invoke correctly handles no Authorization header in the HTTP request"""
459+
response_data = {
460+
"response": "Hello",
461+
"truncated": False,
462+
"referenced_documents": [],
463+
}
464+
465+
self.pipeline.session.post.return_value = MockResponse(response_data, 200)
466+
467+
params = self.get_params(auth_header=None)
468+
result = self.pipeline.invoke(params)
469+
470+
# Verify the response
471+
self.assertEqual(result["response"], "Hello")
472+
473+
# Verify the HTTP call was made with correct parameters
474+
self.pipeline.session.post.assert_called_once()
475+
call_args = self.pipeline.session.post.call_args
476+
477+
# Check that Authorization header are included in the request headers
478+
headers = call_args[1]["headers"]
479+
self.assertEqual(headers["Content-Type"], "application/json")
480+
self.assertNotIn("Authorization", headers)

‎ansible_ai_connect/ai/api/model_pipelines/pipelines.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,7 @@ class ChatBotParameters:
231231
model_id: str
232232
conversation_id: Optional[str]
233233
system_prompt: str
234+
auth_header: Optional[str] = field(kw_only=True, default=None)
234235
mcp_headers: Optional[dict[str, dict[str, str]]] = field(kw_only=True, default=None)
235236
no_tools: bool
236237

@@ -242,6 +243,7 @@ def init(
242243
model_id: Optional[str] = None,
243244
conversation_id: Optional[str] = None,
244245
system_prompt: Optional[str] = None,
246+
auth_header: Optional[str] = None,
245247
mcp_headers: Optional[dict[str, dict[str, str]]] = None,
246248
no_tools: Optional[bool] = False,
247249
):
@@ -251,6 +253,7 @@ def init(
251253
model_id=model_id,
252254
conversation_id=conversation_id,
253255
system_prompt=system_prompt,
256+
auth_header=auth_header,
254257
mcp_headers=mcp_headers,
255258
no_tools=no_tools,
256259
)
@@ -274,6 +277,7 @@ def init(
274277
system_prompt: Optional[str] = None,
275278
media_type: Optional[str] = None,
276279
event: Optional[Any] = None,
280+
auth_header: Optional[str] = None,
277281
mcp_headers: Optional[dict[str, dict[str, str]]] = None,
278282
no_tools: Optional[bool] = False,
279283
):
@@ -285,6 +289,7 @@ def init(
285289
system_prompt=system_prompt,
286290
media_type=media_type,
287291
event=event,
292+
auth_header=auth_header,
288293
mcp_headers=mcp_headers,
289294
no_tools=no_tools,
290295
)

‎ansible_ai_connect/ai/api/views.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -261,6 +261,10 @@ def finalize_response(self, request, response, *args, **kwargs):
261261
send_schema1_event(self.event)
262262
return response
263263

264+
@staticmethod
265+
def get_auth_header(request: Request, config: HttpConfiguration) -> str:
266+
return f"Bearer {settings.CHATBOT_API_KEY}"
267+
264268
@staticmethod
265269
def get_mcp_headers(request: Request, config: HttpConfiguration) -> dict:
266270
mcp_headers = {}
@@ -1122,6 +1126,7 @@ def post(self, request) -> Response:
11221126
self.event.modelName = self.req_model_id or self.llm.config.model_id
11231127
self.event.no_tools = no_tools
11241128

1129+
auth_header = self.get_auth_header(request, self.llm.config)
11251130
mcp_headers = self.get_mcp_headers(request, self.llm.config)
11261131

11271132
data = self.llm.invoke(
@@ -1131,6 +1136,7 @@ def post(self, request) -> Response:
11311136
model_id=self.req_model_id or self.llm.config.model_id,
11321137
provider=req_provider,
11331138
conversation_id=conversation_id,
1139+
auth_header=auth_header,
11341140
mcp_headers=mcp_headers,
11351141
no_tools=no_tools,
11361142
)
@@ -1221,6 +1227,7 @@ def post(self, request) -> StreamingHttpResponse:
12211227
self.event.modelName = self.req_model_id or self.llm.config.model_id
12221228
self.event.no_tools = no_tools
12231229

1230+
auth_header = self.get_auth_header(request, self.llm.config)
12241231
mcp_headers = self.get_mcp_headers(request, self.llm.config)
12251232

12261233
return self.llm.invoke(
@@ -1232,6 +1239,7 @@ def post(self, request) -> StreamingHttpResponse:
12321239
conversation_id=conversation_id,
12331240
media_type=media_type,
12341241
event=copy.copy(self.event),
1242+
auth_header=auth_header,
12351243
mcp_headers=mcp_headers,
12361244
no_tools=no_tools,
12371245
)

‎ansible_ai_connect/main/settings/base.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -595,6 +595,7 @@ def is_ssl_enabled(value: str) -> bool:
595595
CHATBOT_DEFAULT_PROVIDER = os.getenv("CHATBOT_DEFAULT_PROVIDER")
596596
CHATBOT_DEBUG_UI = os.getenv("CHATBOT_DEBUG_UI", "False").lower() == "true"
597597
CHATBOT_DEFAULT_SYSTEM_PROMPT = os.getenv("CHATBOT_DEFAULT_SYSTEM_PROMPT")
598+
CHATBOT_API_KEY = os.getenv("CHATBOT_API_KEY")
598599
# ==========================================
599600

600601
# ==========================================

0 commit comments

Comments
 (0)