address greptile issues - add response obj and test file

litellm/integrations/akto/akto_logger.py

@@ -234,7 +234,9 @@ class AktoLogger(CustomLogger):
         try:
             data = self.extract_logging_data(kwargs)
             status = self.get_failure_status_code(kwargs)
-            payload = self.build_akto_payload(data, status_code=status)
+            payload = self.build_akto_payload(
+                data, status_code=status, response_obj=response_obj
+            )
             self.sync_http_handler.post(**self.request_kwargs(payload))
         except Exception as e:
             verbose_logger.error("Akto logging error (failure): %s", e)
@@ -243,7 +245,9 @@ class AktoLogger(CustomLogger):
         try:
             data = self.extract_logging_data(kwargs)
             status = self.get_failure_status_code(kwargs)
-            payload = self.build_akto_payload(data, status_code=status)
+            payload = self.build_akto_payload(
+                data, status_code=status, response_obj=response_obj
+            )
             await self.async_http_handler.post(**self.request_kwargs(payload))
         except Exception as e:
             verbose_logger.error("Akto logging error (failure): %s", e)

tests/guardrails_tests/test_akto_logger.py

@@ -0,0 +1,235 @@
+import json
+import os
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+import pytest
+
+from litellm.integrations.akto.akto_logger import AktoLogger
+
+
+# ── Fixtures ──
+
+
+@pytest.fixture
+def akto_env():
+    with patch.dict(
+        os.environ,
+        {
+            "AKTO_DATA_INGESTION_API_BASE": "http://localhost:9090",
+            "AKTO_API_KEY": "test-token",
+        },
+    ):
+        yield
+
+
+@pytest.fixture
+def logger(akto_env):
+    return AktoLogger()
+
+
+@pytest.fixture
+def sample_kwargs():
+    return {
+        "messages": [{"role": "user", "content": "Hello"}],
+        "model": "gpt-4",
+        "litellm_params": {
+            "metadata": {
+                "user_api_key_request_route": "/v1/chat/completions",
+                "user_api_key_user_id": "user-1",
+                "user_api_key_team_id": "team-1",
+            },
+            "proxy_server_request": {
+                "headers": {
+                    "host": "my-litellm.example.com",
+                    "content-type": "application/json",
+                    "x-forwarded-for": "10.0.0.1",
+                },
+            },
+        },
+    }
+
+
+# ── Init ──
+
+
+def test_init_requires_env_vars():
+    with patch.dict(os.environ, {}, clear=True):
+        with pytest.raises(Exception, match="AKTO_DATA_INGESTION_API_BASE"):
+            AktoLogger()
+
+
+def test_init_requires_api_key():
+    with patch.dict(
+        os.environ,
+        {"AKTO_DATA_INGESTION_API_BASE": "http://x"},
+        clear=True,
+    ):
+        with pytest.raises(Exception, match="AKTO_API_KEY"):
+            AktoLogger()
+
+
+def test_init_success(logger):
+    assert logger.akto_base_url == "http://localhost:9090"
+    assert logger.akto_api_key == "test-token"
+    assert logger.akto_account_id == "1000000"
+    assert logger.akto_vxlan_id == "0"
+
+
+# ── extract_logging_data ──
+
+
+def test_extract_logging_data_promotes_from_litellm_params():
+    kwargs = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "litellm_params": {
+            "metadata": {"user_api_key_request_route": "/v1/chat/completions"},
+            "proxy_server_request": {"headers": {"host": "example.com"}},
+        },
+    }
+    data = AktoLogger.extract_logging_data(kwargs)
+    assert data["metadata"]["user_api_key_request_route"] == "/v1/chat/completions"
+    assert data["proxy_server_request"]["headers"]["host"] == "example.com"
+
+
+def test_extract_logging_data_does_not_overwrite():
+    kwargs = {
+        "metadata": {"user_api_key_request_route": "/chat/completions"},
+        "litellm_params": {
+            "metadata": {"user_api_key_request_route": "/v1/chat/completions"},
+        },
+    }
+    data = AktoLogger.extract_logging_data(kwargs)
+    assert data["metadata"]["user_api_key_request_route"] == "/chat/completions"
+
+
+# ── Payload ──
+
+
+def test_build_akto_payload(logger, sample_kwargs):
+    data = AktoLogger.extract_logging_data(sample_kwargs)
+    payload = logger.build_akto_payload(data)
+
+    assert payload["path"] == "/v1/chat/completions"
+    assert payload["method"] == "POST"
+    assert payload["statusCode"] == "200"
+    assert payload["source"] == "MIRRORING"
+    assert payload["contextSource"] == "AGENTIC"
+    assert payload["ip"] == "10.0.0.1"
+    assert payload["akto_account_id"] == "1000000"
+
+    req_body = json.loads(payload["requestPayload"])
+    assert req_body["model"] == "gpt-4"
+    assert req_body["messages"][0]["content"] == "Hello"
+
+    req_headers = json.loads(payload["requestHeaders"])
+    assert req_headers["host"] == "my-litellm.example.com"
+
+    tag = json.loads(payload["tag"])
+    assert tag == {"gen-ai": "Gen AI", "user_id": "user-1", "team_id": "team-1"}
+
+
+def test_build_akto_payload_failure_status(logger):
+    payload = logger.build_akto_payload({}, status_code=500)
+    assert payload["statusCode"] == "500"
+    assert payload["status"] == "500"
+
+
+# ── Sensitive headers ──
+
+
+def test_build_request_headers_strips_sensitive():
+    headers = AktoLogger.build_request_headers(
+        {
+            "proxy_server_request": {
+                "headers": {
+                    "host": "myhost.com",
+                    "Authorization": "Bearer sk-secret",
+                    "X-Api-Key": "key-123",
+                    "Cookie": "session=abc",
+                    "user-agent": "test",
+                }
+            }
+        }
+    )
+    assert "authorization" not in headers
+    assert "x-api-key" not in headers
+    assert "cookie" not in headers
+    assert headers["host"] == "myhost.com"
+    assert headers["user-agent"] == "test"
+
+
+# ── get_failure_status_code ──
+
+
+def test_get_failure_status_code_with_status():
+    exc = MagicMock()
+    exc.status_code = 403
+    assert AktoLogger.get_failure_status_code({"exception": exc}) == 403
+
+
+def test_get_failure_status_code_no_status():
+    assert AktoLogger.get_failure_status_code({"exception": Exception("err")}) == 500
+
+
+def test_get_failure_status_code_no_exception():
+    assert AktoLogger.get_failure_status_code({}) == 500
+
+
+# ── Async success ──
+
+
+@pytest.mark.asyncio
+async def test_async_log_success_event(logger, sample_kwargs):
+    logger.async_http_handler.post = AsyncMock(return_value=MagicMock(status_code=200))
+    mock_resp = MagicMock()
+    mock_resp.model_dump.return_value = {
+        "choices": [{"message": {"content": "Hi!", "role": "assistant"}}]
+    }
+
+    await logger.async_log_success_event(
+        kwargs=sample_kwargs, response_obj=mock_resp, start_time=None, end_time=None
+    )
+
+    logger.async_http_handler.post.assert_called_once()
+    call = logger.async_http_handler.post.call_args.kwargs
+    assert call["params"]["ingest_data"] == "true"
+    payload = json.loads(call["data"])
+    assert payload["statusCode"] == "200"
+    assert (
+        json.loads(payload["responsePayload"])["choices"][0]["message"]["content"]
+        == "Hi!"
+    )
+
+
+# ── Async failure ──
+
+
+@pytest.mark.asyncio
+async def test_async_log_failure_event(logger, sample_kwargs):
+    logger.async_http_handler.post = AsyncMock(return_value=MagicMock(status_code=200))
+
+    await logger.async_log_failure_event(
+        kwargs=sample_kwargs, response_obj=None, start_time=None, end_time=None
+    )
+
+    logger.async_http_handler.post.assert_called_once()
+    payload = json.loads(logger.async_http_handler.post.call_args.kwargs["data"])
+    assert payload["statusCode"] == "500"
+
+
+# ── Error handling ──
+
+
+@pytest.mark.asyncio
+async def test_async_log_swallows_errors(logger):
+    logger.async_http_handler.post = AsyncMock(
+        side_effect=httpx.ConnectError("refused")
+    )
+    # Should not raise
+    await logger.async_log_success_event(
+        kwargs={"messages": [], "model": "m", "litellm_params": {}},
+        response_obj=None,
+        start_time=None,
+        end_time=None,
+    )