Spaces:
Configuration error
Configuration error
File size: 7,342 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 |
import json
import os
import sys
from unittest.mock import AsyncMock, patch
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from enterprise.litellm_enterprise.enterprise_callbacks.session_handler import (
_ENTERPRISE_ResponsesSessionHandler,
)
@pytest.mark.asyncio
async def test_get_chat_completion_message_history_for_previous_response_id():
"""
Test get_chat_completion_message_history_for_previous_response_id with mock data
"""
# Mock data based on the provided spend logs (simplified version)
mock_spend_logs = [
{
"request_id": "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb",
"call_type": "aresponses",
"api_key": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
"spend": 0.004803,
"total_tokens": 329,
"prompt_tokens": 11,
"completion_tokens": 318,
"startTime": "2025-05-30T03:17:06.703+00:00",
"endTime": "2025-05-30T03:17:11.894+00:00",
"model": "claude-3-5-sonnet-latest",
"session_id": "a96757c4-c6dc-4c76-b37e-e7dfa526b701",
"proxy_server_request": {
"input": "who is Michael Jordan",
"model": "anthropic/claude-3-5-sonnet-latest",
},
"response": {
"id": "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb",
"model": "claude-3-5-sonnet-20241022",
"object": "chat.completion",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "Michael Jordan (born February 17, 1963) is widely considered the greatest basketball player of all time. Here are some key points about him...",
"tool_calls": None,
"function_call": None,
},
"finish_reason": "stop",
}
],
"created": 1748575031,
"usage": {
"total_tokens": 329,
"prompt_tokens": 11,
"completion_tokens": 318,
},
},
"status": "success",
},
{
"request_id": "chatcmpl-370760c9-39fa-4db7-b034-d1f8d933c935",
"call_type": "aresponses",
"api_key": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
"spend": 0.010437,
"total_tokens": 967,
"prompt_tokens": 339,
"completion_tokens": 628,
"startTime": "2025-05-30T03:17:28.600+00:00",
"endTime": "2025-05-30T03:17:39.921+00:00",
"model": "claude-3-5-sonnet-latest",
"session_id": "a96757c4-c6dc-4c76-b37e-e7dfa526b701",
"proxy_server_request": {
"input": "can you tell me more about him",
"model": "anthropic/claude-3-5-sonnet-latest",
"previous_response_id": "resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmFudGhyb3BpYzttb2RlbF9pZDplMGYzMDJhMTQxMmU3ODQ3MGViYjI4Y2JlZDAxZmZmNWY4OGMwZDMzMWM2NjdlOWYyYmE0YjQxM2M2ZmJkMjgyO3Jlc3BvbnNlX2lkOmNoYXRjbXBsLTkzNWI4ZGFkLWZkYzItNDY2ZS1hOGNhLWUyNmU1YThhMjFiYg==",
},
"response": {
"id": "chatcmpl-370760c9-39fa-4db7-b034-d1f8d933c935",
"model": "claude-3-5-sonnet-20241022",
"object": "chat.completion",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "Here's more detailed information about Michael Jordan...",
"tool_calls": None,
"function_call": None,
},
"finish_reason": "stop",
}
],
"created": 1748575059,
"usage": {
"total_tokens": 967,
"prompt_tokens": 339,
"completion_tokens": 628,
},
},
"status": "success",
},
]
# Mock the get_all_spend_logs_for_previous_response_id method
with patch.object(
_ENTERPRISE_ResponsesSessionHandler,
"get_all_spend_logs_for_previous_response_id",
new_callable=AsyncMock,
) as mock_get_spend_logs:
mock_get_spend_logs.return_value = mock_spend_logs
# Test the function
previous_response_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb"
result = await _ENTERPRISE_ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
previous_response_id
)
# Verify the mock was called with correct parameters
mock_get_spend_logs.assert_called_once_with(previous_response_id)
# Verify the returned ChatCompletionSession structure
assert "messages" in result
assert "litellm_session_id" in result
# Verify session_id is extracted correctly
assert result["litellm_session_id"] == "a96757c4-c6dc-4c76-b37e-e7dfa526b701"
# Verify messages structure
messages = result["messages"]
assert len(messages) == 4 # 2 user messages + 2 assistant messages
# Check the message sequence
# First user message
assert messages[0].get("role") == "user"
assert messages[0].get("content") == "who is Michael Jordan"
# First assistant response
assert messages[1].get("role") == "assistant"
content_1 = messages[1].get("content", "")
if isinstance(content_1, str):
assert "Michael Jordan" in content_1
assert content_1.startswith("Michael Jordan (born February 17, 1963)")
# Second user message
assert messages[2].get("role") == "user"
assert messages[2].get("content") == "can you tell me more about him"
# Second assistant response
assert messages[3].get("role") == "assistant"
content_3 = messages[3].get("content", "")
if isinstance(content_3, str):
assert "Here's more detailed information about Michael Jordan" in content_3
@pytest.mark.asyncio
async def test_get_chat_completion_message_history_empty_spend_logs():
"""
Test get_chat_completion_message_history_for_previous_response_id with empty spend logs
"""
with patch.object(
_ENTERPRISE_ResponsesSessionHandler,
"get_all_spend_logs_for_previous_response_id",
new_callable=AsyncMock,
) as mock_get_spend_logs:
mock_get_spend_logs.return_value = []
previous_response_id = "non-existent-id"
result = await _ENTERPRISE_ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
previous_response_id
)
# Verify empty result structure
assert result.get("messages") == []
assert result.get("litellm_session_id") is None
|