Spaces:
Build error
Build error
File size: 5,908 Bytes
9c7d451 04a98a0 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 e4a41fa 9c7d451 | 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 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | """Shared pytest fixtures for Headroom tests."""
# CRITICAL: Must be set before ANY imports that could trigger sentence_transformers
# The Rust tokenizers use parallelism that deadlocks with pytest-asyncio
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import json
import tempfile
from datetime import datetime
from pathlib import Path
from unittest.mock import Mock
import pytest
# Sample messages fixtures
@pytest.fixture
def sample_messages():
"""Basic conversation messages."""
return [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello, how are you?"},
{"role": "assistant", "content": "I'm doing well, thank you!"},
]
@pytest.fixture
def sample_messages_with_tools():
"""Conversation with tool calls and responses."""
return [
{"role": "system", "content": "You are a helpful assistant with tools."},
{"role": "user", "content": "Search for user 12345"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {"name": "search_user", "arguments": '{"user_id": "12345"}'},
}
],
},
{
"role": "tool",
"tool_call_id": "call_123",
"content": '{"id": "12345", "name": "Alice", "email": "alice@example.com"}',
},
{"role": "assistant", "content": "I found user Alice with ID 12345."},
]
@pytest.fixture
def sample_tool_output_large():
"""Large tool output for compression testing (100 items)."""
return json.dumps(
[
{
"id": i,
"name": f"Item {i}",
"score": i * 0.1,
"status": "active" if i % 2 == 0 else "inactive",
}
for i in range(100)
]
)
@pytest.fixture
def sample_tool_output_with_errors():
"""Tool output containing error items."""
items = [{"id": i, "status": "success"} for i in range(20)]
items[5] = {"id": 5, "status": "error", "message": "Connection refused"}
items[15] = {"id": 15, "status": "failed", "exception": "TimeoutError"}
return json.dumps(items)
@pytest.fixture
def sample_system_prompt_with_date():
"""System prompt containing dynamic date."""
return "You are a helpful assistant. Current date: 2025-01-06. Help the user with their tasks."
@pytest.fixture
def sample_anthropic_messages():
"""Anthropic-style messages with content blocks."""
return [
{
"role": "user",
"content": [
{"type": "text", "text": "Analyze this image"},
{
"type": "image",
"source": {"type": "base64", "media_type": "image/png", "data": "..."},
},
],
}
]
# Mock client fixtures
@pytest.fixture
def mock_openai_response():
"""Mock OpenAI API response."""
mock = Mock()
mock.id = "chatcmpl-123"
mock.model = "gpt-4o"
mock.usage = Mock()
mock.usage.prompt_tokens = 100
mock.usage.completion_tokens = 50
mock.usage.total_tokens = 150
mock.choices = [Mock()]
mock.choices[0].message = Mock()
mock.choices[0].message.content = "This is a response."
mock.choices[0].message.role = "assistant"
mock.choices[0].finish_reason = "stop"
return mock
@pytest.fixture
def mock_openai_client(mock_openai_response):
"""Mock OpenAI client."""
client = Mock()
client.chat = Mock()
client.chat.completions = Mock()
client.chat.completions.create = Mock(return_value=mock_openai_response)
return client
# Storage fixtures
@pytest.fixture
def temp_sqlite_db():
"""Temporary SQLite database path."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
yield f.name
Path(f.name).unlink(missing_ok=True)
@pytest.fixture
def temp_jsonl_file():
"""Temporary JSONL file path."""
with tempfile.NamedTemporaryFile(suffix=".jsonl", delete=False) as f:
yield f.name
Path(f.name).unlink(missing_ok=True)
# Provider fixtures
@pytest.fixture
def openai_provider():
"""OpenAI provider instance."""
from headroom.providers.openai import OpenAIProvider
return OpenAIProvider()
@pytest.fixture
def openai_tokenizer():
"""OpenAI token counter for gpt-4o."""
from headroom.providers.openai import OpenAITokenCounter
return OpenAITokenCounter("gpt-4o")
# Config fixtures
@pytest.fixture
def default_config():
"""Default HeadroomConfig."""
from headroom.config import HeadroomConfig
return HeadroomConfig()
@pytest.fixture
def smart_crusher_config():
"""SmartCrusher config for testing."""
from headroom.config import SmartCrusherConfig
return SmartCrusherConfig(
enabled=True,
min_items_to_analyze=3,
min_tokens_to_crush=0, # Always crush for tests
max_items_after_crush=10,
)
# Helper for creating RequestMetrics
@pytest.fixture
def sample_request_metrics():
"""Sample RequestMetrics for storage tests."""
from headroom.config import RequestMetrics
return RequestMetrics(
request_id="test-123",
timestamp=datetime(2025, 1, 6, 12, 0, 0),
model="gpt-4o",
stream=False,
mode="audit",
tokens_input_before=1000,
tokens_input_after=800,
tokens_output=200,
block_breakdown={"system": 100, "user": 200, "assistant": 500},
waste_signals={"json_bloat": 50},
stable_prefix_hash="abc123",
cache_alignment_score=85.0,
cached_tokens=100,
transforms_applied=["CacheAligner", "SmartCrusher"],
tool_units_dropped=1,
turns_dropped=0,
messages_hash="def456",
)
|