Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions backend/services/llm_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from core.exceptions import LLMServiceError
from services.circuit_breaker import provider_circuit_breaker
from services.retry import retry_transient
from pydantic import BaseModel, Field
from pydantic import BaseModel, ConfigDict, Field
from services.llm_provider_urls import build_llm_provider_http_client

logger = logging.getLogger(__name__)
Expand All @@ -26,6 +26,8 @@
class ExtractionResult(BaseModel):
"""Result model for email extraction."""

model_config = ConfigDict(extra="forbid", strict=True)

summary: str
action_items: list[str]
provenance: str | None = None
Expand Down Expand Up @@ -76,7 +78,7 @@ async def extract_action_items_and_summary(
{"role": "user", "content": email_body},
],
response_format=ExtractionResult,
),
),
operation_name="summary extraction",
),
)
Expand Down
8 changes: 7 additions & 1 deletion backend/services/project_graph/llm_extractor.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
from typing import Iterable

from openai import AsyncOpenAI
from pydantic import BaseModel
from pydantic import BaseModel, ConfigDict

from services.llm_provider_urls import build_llm_provider_http_client

Expand Down Expand Up @@ -79,6 +79,8 @@


class ExtractedObjectPayload(BaseModel):
model_config = ConfigDict(extra="forbid", strict=True)

object_type: str
title: str
summary: str
Expand All @@ -91,13 +93,17 @@ class ExtractedObjectPayload(BaseModel):


class ExtractedRelationPayload(BaseModel):
model_config = ConfigDict(extra="forbid", strict=True)

source_local_key: str
target_local_key: str
relation_type: str
confidence: float


class ExtractionPayload(BaseModel):
model_config = ConfigDict(extra="forbid", strict=True)

objects: list[ExtractedObjectPayload]
relations: list[ExtractedRelationPayload] = []

Expand Down
4 changes: 3 additions & 1 deletion backend/services/rag_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
import logging

from openai import AsyncOpenAI
from pydantic import BaseModel
from pydantic import BaseModel, ConfigDict

from core.config import settings
from services.llm_provider_urls import build_llm_provider_http_client
Expand All @@ -25,6 +25,8 @@


class GroundedAnswerPayload(BaseModel):
model_config = ConfigDict(extra="forbid", strict=True)

answer: str
cited_email_ids: list[int]

Expand Down
177 changes: 177 additions & 0 deletions backend/tests/test_llm_structured_output_contract.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,177 @@
"""OpenAI-compatible structured-output contracts used by Naruon consumers."""

import json

import httpx
import pytest
from openai import AsyncOpenAI
from openai.lib._parsing import type_to_response_format_param
from pydantic import ValidationError

from core.exceptions import LLMServiceError
from services.llm_service import ExtractionResult, extract_action_items_and_summary
from services.project_graph import llm_extractor
from services.project_graph.llm_extractor import ExtractionPayload
from services import rag_service
from services.rag_service import GroundedAnswerPayload


class _CaptureTransport(httpx.AsyncBaseTransport):
def __init__(self, content: str) -> None:
self.body = None
self.content = content

async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
self.body = json.loads(request.content)
return httpx.Response(
200,
request=request,
json={
"id": "chatcmpl-contract",
"object": "chat.completion",
"created": 0,
"model": "orchestrator/free",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": self.content},
"finish_reason": "stop",
}
],
},
)


@pytest.mark.parametrize(
("payload_model", "valid_payload"),
[
(ExtractionResult, {"summary": "s", "action_items": []}),
(ExtractionPayload, {"objects": [], "relations": []}),
(GroundedAnswerPayload, {"answer": "a", "cited_email_ids": []}),
],
)
def test_structured_payloads_use_strict_json_schema_and_reject_extra_fields(
payload_model,
valid_payload,
):
envelope = type_to_response_format_param(payload_model)

assert envelope["type"] == "json_schema"
assert envelope["json_schema"]["strict"] is True
assert envelope["json_schema"]["schema"]["additionalProperties"] is False
with pytest.raises(ValidationError):
payload_model.model_validate({**valid_payload, "unexpected_field": True})


@pytest.mark.asyncio
async def test_summary_uses_real_sdk_json_schema_on_orchestrator_route(monkeypatch):
transport = _CaptureTransport(
json.dumps({"summary": "s", "action_items": [], "confidence": 90})
)
client = AsyncOpenAI(
api_key="test-token",
base_url="https://orchestrator.example/v1",
http_client=httpx.AsyncClient(transport=transport),
)
monkeypatch.setattr("services.llm_service.AsyncOpenAI", lambda **_: client)
monkeypatch.setattr(
"services.llm_service.build_llm_provider_http_client",
_validated_client,
)

result = await extract_action_items_and_summary(
"email",
"tenant-scoped-token",
base_url="https://orchestrator.example/v1",
model="orchestrator/free",
)

assert result.summary == "s"
assert transport.body["model"] == "orchestrator/free"
assert transport.body["response_format"]["type"] == "json_schema"
assert transport.body["response_format"]["json_schema"]["strict"] is True


@pytest.mark.asyncio
async def test_project_graph_uses_real_sdk_json_schema_on_orchestrator_route(
monkeypatch,
):
transport = _CaptureTransport(json.dumps({"objects": [], "relations": []}))
client = AsyncOpenAI(
api_key="test-token",
base_url="https://orchestrator.example/v1",
http_client=httpx.AsyncClient(transport=transport),
)
monkeypatch.setattr(llm_extractor, "AsyncOpenAI", lambda **_: client)
monkeypatch.setattr(
llm_extractor, "build_llm_provider_http_client", _validated_client
)

result = await llm_extractor._call_llm(
api_key="tenant-scoped-token",
base_url="https://orchestrator.example/v1",
model="orchestrator/free",
segments_json='{"segments": []}',
)

assert result.objects == []
assert transport.body["model"] == "orchestrator/free"
assert transport.body["response_format"]["type"] == "json_schema"


@pytest.mark.asyncio
async def test_grounded_answer_uses_real_sdk_json_schema_on_orchestrator_route(
monkeypatch,
):
transport = _CaptureTransport(
json.dumps({"answer": "grounded", "cited_email_ids": []})
)
client = AsyncOpenAI(
api_key="test-token",
base_url="https://orchestrator.example/v1",
http_client=httpx.AsyncClient(transport=transport),
)
monkeypatch.setattr(rag_service, "AsyncOpenAI", lambda **_: client)
monkeypatch.setattr(
rag_service, "build_llm_provider_http_client", _validated_client
)

result = await rag_service._call_llm(
api_key="tenant-scoped-token",
base_url="https://orchestrator.example/v1",
model="orchestrator/free",
question="q",
emails_json='{"emails": []}',
)

assert result.answer == "grounded"
assert transport.body["model"] == "orchestrator/free"
assert transport.body["response_format"]["type"] == "json_schema"


@pytest.mark.asyncio
async def test_schema_extra_field_fails_closed(monkeypatch):
transport = _CaptureTransport(
json.dumps({"summary": "s", "action_items": [], "unexpected": True})
)
client = AsyncOpenAI(
api_key="test-token",
base_url="https://orchestrator.example/v1",
http_client=httpx.AsyncClient(transport=transport),
)
monkeypatch.setattr("services.llm_service.AsyncOpenAI", lambda **_: client)
monkeypatch.setattr(
"services.llm_service.build_llm_provider_http_client", _validated_client
)

with pytest.raises(LLMServiceError, match="LLM API error during extraction"):
await extract_action_items_and_summary(
"email",
"tenant-scoped-token",
base_url="https://orchestrator.example/v1",
model="orchestrator/free",
)


async def _validated_client(*_):
return "https://orchestrator.example/v1", None
Loading