[Codex update] Improve how validation happens through the project (#95)

Author: ulya-tkch

pyproject.toml

@@ -25,8 +25,8 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dependencies = [
-  "cleanlab-tlm~=1.1",
-  "codex-sdk==0.1.0a22",
+  "cleanlab-tlm~=1.1,>=1.1.14",
+  "codex-sdk==0.1.0a23",
   "pydantic>=2.0.0, <3",
 ]
 
@@ -46,6 +46,7 @@ extra-dependencies = [
   "smolagents; python_version >= '3.10'",
   "thefuzz",
   "langchain-core",
+  "openai"
 ]
 
 [[tool.hatch.envs.types.matrix]]
@@ -63,6 +64,7 @@ extra-dependencies = [
   "smolagents; python_version >= '3.10'",
   "thefuzz",
   "langchain-core",
+  "openai",
 ]
 
 [tool.hatch.envs.hatch-test.env-vars]

src/cleanlab_codex/__init__.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: MIT
 from cleanlab_codex.client import Client
 from cleanlab_codex.project import Project
-from cleanlab_codex.validator import Validator
 
-__all__ = ["Client", "Project", "Validator"]
+__all__ = ["Client", "Project"]

src/cleanlab_codex/project.py

@@ -4,9 +4,10 @@
 
 from datetime import datetime
 from typing import TYPE_CHECKING as _TYPE_CHECKING
-from typing import Dict, Optional
+from typing import Dict, Optional, Union, cast
 
 from codex import AuthenticationError
+from codex.types.project_validate_params import Response
 
 from cleanlab_codex.internal.analytics import _AnalyticsMetadata
 from cleanlab_codex.internal.sdk_client import client_from_access_key
@@ -17,6 +18,7 @@
 
     from codex import Codex as _Codex
     from codex.types.project_validate_response import ProjectValidateResponse
+    from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
 
 
 _ERROR_CREATE_ACCESS_KEY = (
@@ -145,37 +147,60 @@ def create_access_key(
 
     def validate(
         self,
-        context: str,
-        prompt: str,
-        query: str,
-        response: str,
         *,
-        custom_metadata: Optional[object] = None,
+        messages: list[ChatCompletionMessageParam],
+        response: Union[ChatCompletion, str],
+        query: str,
+        context: str,
+        rewritten_query: Optional[str] = None,
+        metadata: Optional[object] = None,
         eval_scores: Optional[Dict[str, float]] = None,
-        eval_thresholds: Optional[Dict[str, float]] = None,
     ) -> ProjectValidateResponse:
-        """Run validation on a query to an AI system.
+        """Evaluate the quality of an AI-generated response using the structured message history, query, and retrieved context.
+
+        This method runs validation on an AI response using the full `messages` history (formatted as OpenAI-style chat messages),
+        which should include the latest user query and any preceding system or assistant messages.
+
+        **For single-turn Q&A apps, `messages` can be a minimal list with one user message. For multi-turn conversations, provide the full dialog
+        leading up to the final response.
+
+        The function assesses the trustworthiness and quality of the AI `response` in light of the provided `context` and
+        `query`, which should align with the most recent user message in `messages`.
+
+        If your AI response is flagged as problematic, then this method will:
+            - return an expert answer if one was previously provided for a similar query
+            - otherwise log this query for future SME review (to consider providing an expert answer) in the Web interface.
 
         Args:
-            context (str): The context used by the AI system to generate a response for the query.
-            prompt (str): The full prompt (including system instructions, context, and the original query) used by the AI system to generate a response for the query.
-            query (str): The original user input to the AI system.
-            response (str): The response generated by the AI system for the query.
-            custom_metadata (object, optional): Custom metadata to log in Codex for the query.
-            eval_scores (Dict[str, float], optional): Optional scores to use for the query. When provided, Codex will skip running TrustworthyRAG evaluations on the query and use the provided scores instead.
-            eval_thresholds (Dict[str, float], optional): Optional thresholds to use for evaluating the query. We recommend configuring thresholds on the Project instead and using the same thresholds for all queries.
+            messages (list[ChatCompletionMessageParam]): The full message history from the AI conversation, formatted for OpenAI-style chat completion.
+                This must include the final user message that triggered the AI response. All other arguments—`query`, `context`, and `response`—
+                must correspond specifically to this final user message.
+            response (ChatCompletion | str): Your AI-response that was generated based on the given `messages`. This is the response being evaluated, and should not appear in the `messages`.
+            query (str): The user query that the `response` is answering. This query should be the latest user message in `messages`.
+            context (str): The retrieved context (e.g., from your RAG system) that was supplied to the AI when generating the `response` to the final user query in `messages`.
+            rewritten_query (str, optional): An optional reformulation of `query` (e.g. made self-contained w.r.t multi-turn conversations) to improve retrieval quality.
+            metadata (object, optional): Arbitrary metadata to associate with this validation for logging or analysis inside the Codex project.
+            eval_scores (dict[str, float], optional): Precomputed evaluation scores to bypass automatic scoring. Providing `eval_scores` for specific evaluations bypasses automated scoring and uses the supplied scores instead. Consider providing these scores if you already have them precomputed to reduce runtime.
 
         Returns:
-            ProjectValidateResponse: The response from the validation.
+            ProjectValidateResponse: A structured object with the following fields:
+
+                - should_guardrail (bool): True if the AI system should suppress or modify the response before returning it to the user. When True, the response is considered problematic and may require further review or modification.
+                - escalated_to_sme (bool): True if the query should be escalated to SME for review. When True, the query is logged and may be answered by an expert.
+                - eval_scores (dict[str, ThresholdedEvalScore]): Evaluation scores for different response attributes (e.g., trustworthiness, helpfulness, ...). Each includes a numeric score and a `failed` flag indicating whether the score falls below threshold.
+                - expert_answer (str | None): If it was auto-determined that this query should be escalated to SME, and a prior SME answer for a similar query was found, then this will return that expert answer.  Otherwise, it is None.
+
+                When available, consider swapping your AI response with the expert answer before serving the response to your user.
         """
+
         return self._sdk_client.projects.validate(
             self._id,
+            messages=messages,
+            response=cast(Response, response),
             context=context,
-            prompt=prompt,
             query=query,
-            response=response,
-            custom_eval_thresholds=eval_thresholds,
-            custom_metadata=custom_metadata,
+            rewritten_question=rewritten_query,
+            custom_metadata=metadata,
             eval_scores=eval_scores,
         )
 

src/cleanlab_codex/validator.py

@@ -1,3 +1,17 @@
 from tests.fixtures.client import default_headers, mock_client_from_access_key, mock_client_from_api_key
+from tests.fixtures.validate import (
+    openai_chat_completion,
+    openai_messages_bad_no_user,
+    openai_messages_conversational,
+    openai_messages_single_turn,
+)
 
-__all__ = ["mock_client_from_access_key", "mock_client_from_api_key", "default_headers"]
+__all__ = [
+    "mock_client_from_access_key",
+    "mock_client_from_api_key",
+    "default_headers",
+    "openai_chat_completion",
+    "openai_messages_conversational",
+    "openai_messages_single_turn",
+    "openai_messages_bad_no_user",
+]

tests/conftest.py

@@ -0,0 +1,64 @@
+from typing import cast
+
+import pytest
+from openai.types.chat import (
+    ChatCompletion,
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionUserMessageParam,
+)
+
+
+@pytest.fixture
+def openai_chat_completion() -> ChatCompletion:
+    """Fixture that returns a static fake OpenAI ChatCompletion object."""
+    raw_response = {
+        "id": "chatcmpl-test123",
+        "object": "chat.completion",
+        "created": 1719876543,
+        "model": "gpt-4",
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "Paris",
+                },
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {
+            "prompt_tokens": 5,
+            "completion_tokens": 1,
+            "total_tokens": 6,
+        },
+    }
+
+    return ChatCompletion.model_validate(raw_response)
+
+
+@pytest.fixture
+def openai_messages_single_turn() -> list[ChatCompletionMessageParam]:
+    """Fixture that returns a single-turn message format."""
+    return [cast(ChatCompletionUserMessageParam, {"role": "user", "content": "What is the capital of France?"})]
+
+
+@pytest.fixture
+def openai_messages_bad_no_user() -> list[ChatCompletionMessageParam]:
+    """Fixture that returns invalid messages (missing required user message)."""
+    return [
+        cast(ChatCompletionAssistantMessageParam, {"role": "assistant", "content": "hi"}),
+        cast(ChatCompletionSystemMessageParam, {"role": "system", "content": "sys"}),
+    ]
+
+
+@pytest.fixture
+def openai_messages_conversational() -> list[ChatCompletionMessageParam]:
+    """Fixture that returns a conversational message format."""
+    return [
+        cast(ChatCompletionSystemMessageParam, {"role": "system", "content": "You are a helpful assistant."}),
+        cast(ChatCompletionUserMessageParam, {"role": "user", "content": "I love France!"}),
+        cast(ChatCompletionAssistantMessageParam, {"role": "assistant", "content": "That's great!"}),
+        cast(ChatCompletionUserMessageParam, {"role": "user", "content": "What is its capital?"}),
+    ]