generated from coulomb/repo-seed
143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
import json
|
|
|
|
from llm_connect.gemini import GeminiAdapter
|
|
from llm_connect.models import RunConfig
|
|
from llm_connect.openai import OpenAIAdapter
|
|
from llm_connect.openrouter import OpenRouterAdapter
|
|
|
|
|
|
STRUCTURED_SCHEMA = {
|
|
"type": "object",
|
|
"properties": {
|
|
"summary": {"type": "string"},
|
|
"recommendations": {"type": "array", "items": {"type": "string"}},
|
|
},
|
|
"required": ["summary", "recommendations"],
|
|
}
|
|
|
|
|
|
SMOKE_CONFIG = RunConfig(
|
|
model_name="gpt-4",
|
|
temperature=0.1,
|
|
max_tokens=300,
|
|
model_params={
|
|
"reasoning_effort": "medium",
|
|
"max_depth": 3,
|
|
"json_schema": STRUCTURED_SCHEMA,
|
|
},
|
|
)
|
|
|
|
|
|
def test_openrouter_structured_output_payload_and_model_routing(monkeypatch):
|
|
captured: dict[str, object] = {}
|
|
|
|
def fake_post_json(url, payload, headers=None, timeout=300): # noqa: ANN001
|
|
captured["url"] = url
|
|
captured["payload"] = payload
|
|
captured["headers"] = headers
|
|
captured["timeout"] = timeout
|
|
return {
|
|
"id": "or-response",
|
|
"model": payload["model"],
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": json.dumps(
|
|
{"summary": "ok", "recommendations": ["keep payload clean"]}
|
|
)
|
|
},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3},
|
|
}
|
|
|
|
monkeypatch.setattr("llm_connect.openrouter.post_json", fake_post_json)
|
|
adapter = OpenRouterAdapter(
|
|
model="anthropic/claude-sonnet-4",
|
|
api_key="or-test",
|
|
api_base="https://openrouter.example/api/v1",
|
|
)
|
|
|
|
response = adapter.execute_prompt("Return JSON.", SMOKE_CONFIG)
|
|
payload = captured["payload"]
|
|
|
|
assert response.model == "anthropic/claude-sonnet-4"
|
|
assert payload["model"] == "anthropic/claude-sonnet-4"
|
|
assert payload["response_format"]["json_schema"]["schema"] == STRUCTURED_SCHEMA
|
|
assert payload["response_format"]["json_schema"]["strict"] is False
|
|
assert "reasoning_effort" not in payload
|
|
assert "max_depth" not in payload
|
|
assert "json_schema" not in payload
|
|
|
|
|
|
def test_openai_structured_output_payload(monkeypatch):
|
|
captured: dict[str, object] = {}
|
|
|
|
def fake_post_json(url, payload, headers=None, timeout=300): # noqa: ANN001
|
|
captured["payload"] = payload
|
|
return {
|
|
"id": "oa-response",
|
|
"model": payload["model"],
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": json.dumps({"summary": "ok", "recommendations": []})
|
|
},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3},
|
|
}
|
|
|
|
monkeypatch.setattr("llm_connect.openai.post_json", fake_post_json)
|
|
adapter = OpenAIAdapter(model="gpt-4.1-mini", api_key="sk-test")
|
|
|
|
response = adapter.execute_prompt("Return JSON.", SMOKE_CONFIG)
|
|
payload = captured["payload"]
|
|
|
|
assert response.model == "gpt-4.1-mini"
|
|
assert payload["model"] == "gpt-4.1-mini"
|
|
assert payload["response_format"]["json_schema"]["schema"] == STRUCTURED_SCHEMA
|
|
assert "reasoning_effort" not in payload
|
|
assert "max_depth" not in payload
|
|
assert "json_schema" not in payload
|
|
|
|
|
|
def test_gemini_structured_output_payload(monkeypatch):
|
|
captured: dict[str, object] = {}
|
|
|
|
def fake_post_json(url, payload, headers=None, timeout=300): # noqa: ANN001
|
|
captured["url"] = url
|
|
captured["payload"] = payload
|
|
return {
|
|
"candidates": [
|
|
{
|
|
"content": {
|
|
"parts": [
|
|
{"text": json.dumps({"summary": "ok", "recommendations": []})}
|
|
]
|
|
},
|
|
"finishReason": "STOP",
|
|
}
|
|
],
|
|
"usageMetadata": {
|
|
"promptTokenCount": 1,
|
|
"candidatesTokenCount": 2,
|
|
"totalTokenCount": 3,
|
|
},
|
|
}
|
|
|
|
monkeypatch.setattr("llm_connect.gemini.post_json", fake_post_json)
|
|
adapter = GeminiAdapter(model="gemini-2.5-flash", api_key="gemini-test")
|
|
|
|
response = adapter.execute_prompt("Return JSON.", SMOKE_CONFIG)
|
|
payload = captured["payload"]
|
|
|
|
assert response.model == "gemini-2.5-flash"
|
|
assert payload["generationConfig"]["responseMimeType"] == "application/json"
|
|
assert payload["generationConfig"]["responseSchema"] == STRUCTURED_SCHEMA
|
|
assert "reasoning_effort" not in payload
|
|
assert "max_depth" not in payload
|
|
assert "json_schema" not in payload
|