Spaces:
Sleeping
Sleeping
File size: 4,376 Bytes
2e818da | 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 | from unittest.mock import MagicMock, patch
from pydantic import BaseModel
from app.agents.cerebras_client import CerebrasClient
from app.agents.cerebras_errors import CerebrasErrorKind, CerebrasError
class FakeOutput(BaseModel):
answer: str
confidence: int
def _make_mock_response(content: str):
msg = MagicMock()
msg.content = content
choice = MagicMock()
choice.message = msg
resp = MagicMock()
resp.choices = [choice]
# Real SDK responses always carry a numeric usage -> the speed-log line divides
# by it, so a bare MagicMock here would crash. Give it a realistic int.
resp.usage.completion_tokens = 12
return resp
def test_structured_complete_parses_json():
client = CerebrasClient.__new__(CerebrasClient)
client._health = {"status": "ok"}
client._rate_limit_until = 0.0
mock_sdk = MagicMock()
client._client = mock_sdk
mock_sdk.chat.completions.create.return_value = _make_mock_response(
'{"answer": "entropy is disorder", "confidence": 90}'
)
result = client.structured_complete(
[{"role": "user", "content": "define entropy"}], FakeOutput
)
assert isinstance(result, FakeOutput)
assert result.confidence == 90
def test_schema_build_rejects_open_ended_dict_field():
client = CerebrasClient.__new__(CerebrasClient)
client._health = {"status": "ok"}
client._rate_limit_until = 0.0
client._client = MagicMock()
class HasOpenDict(BaseModel):
panel_ids_by_requirement: dict[str, list[str]]
try:
client._build_schema(HasOpenDict)
assert False, "Should have raised ValueError for an open-ended dict field"
except ValueError as err:
assert "open-ended dict field" in str(err)
def test_schema_strips_defs_and_sets_additional_properties():
client = CerebrasClient.__new__(CerebrasClient)
client._health = {"status": "ok"}
client._rate_limit_until = 0.0
client._client = MagicMock()
class Nested(BaseModel):
value: int
class Outer(BaseModel):
nested: Nested
schema = client._build_schema(Outer)
assert "$defs" not in schema
assert schema["additionalProperties"] is False
def test_rate_limit_short_circuit():
import time
client = CerebrasClient.__new__(CerebrasClient)
client._health = {"status": "ok"}
client._rate_limit_until = time.time() + 60
client._client = MagicMock()
try:
client.structured_complete([{"role": "user", "content": "test"}], FakeOutput)
assert False, "Should have raised CerebrasError"
except CerebrasError as err:
assert err.kind == CerebrasErrorKind.RATE_LIMITED
def test_complete_with_tools_returns_message_with_tool_calls():
client = CerebrasClient.__new__(CerebrasClient)
client._health = {"status": "ok"}
client._rate_limit_until = 0.0
mock_sdk = MagicMock()
client._client = mock_sdk
tool_call = MagicMock()
tool_call.id = "call_1"
tool_call.function.name = "web_search"
tool_call.function.arguments = '{"query": "VAE"}'
message = MagicMock()
message.tool_calls = [tool_call]
message.content = ""
resp = MagicMock()
resp.choices = [MagicMock(message=message)]
mock_sdk.chat.completions.create.return_value = resp
tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}]
result = client.complete_with_tools([{"role": "user", "content": "what is a VAE"}], tools)
assert result.tool_calls[0].function.name == "web_search"
# tools + tool_choice forwarded to the SDK
_, kwargs = mock_sdk.chat.completions.create.call_args
assert kwargs["tools"] == tools
assert kwargs["tool_choice"] == "auto"
def test_complete_with_tools_passes_through_no_tool_answer():
client = CerebrasClient.__new__(CerebrasClient)
client._health = {"status": "ok"}
client._rate_limit_until = 0.0
mock_sdk = MagicMock()
client._client = mock_sdk
message = MagicMock()
message.tool_calls = None
message.content = "A VAE is a generative model."
resp = MagicMock()
resp.choices = [MagicMock(message=message)]
mock_sdk.chat.completions.create.return_value = resp
result = client.complete_with_tools([{"role": "user", "content": "vae?"}], [])
assert result.tool_calls is None
assert "generative model" in result.content
|