Feature - add custom metadata property for query (#70)

* Feature - add custom metadata property for query

* bump version and changelog

* add test for metadata

* hatch fmt

* fix type

* remove 'for analytics purposes'

Author: LukeMainwaring

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [1.0.8] - 2025-04-03
+
+- Update `Project.query()` method with optional `metadata` property to log and store arbitrary metadata.
 - Remove `response_validation.py` module.
 
 ## [1.0.7] - 2025-04-02
@@ -44,7 +47,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Initial release of the `cleanlab-codex` client library.
 
-[Unreleased]: https://github.com/cleanlab/cleanlab-codex/compare/v1.0.7...HEAD
+[Unreleased]: https://github.com/cleanlab/cleanlab-codex/compare/v1.0.8...HEAD
+[1.0.8]: https://github.com/cleanlab/cleanlab-codex/compare/v1.0.7...v1.0.8
 [1.0.7]: https://github.com/cleanlab/cleanlab-codex/compare/v1.0.6...v1.0.7
 [1.0.6]: https://github.com/cleanlab/cleanlab-codex/compare/v1.0.5...v1.0.6
 [1.0.5]: https://github.com/cleanlab/cleanlab-codex/compare/v1.0.4...v1.0.5

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 ]
 dependencies = [
   "cleanlab-tlm~=1.0.12",
-  "codex-sdk==0.1.0-alpha.13",
+  "codex-sdk==0.1.0-alpha.14",
   "pydantic>=2.0.0, <3",
 ]
 

src/cleanlab_codex/__about__.py

@@ -1,2 +1,2 @@
 # SPDX-License-Identifier: MIT
-__version__ = "1.0.7"
+__version__ = "1.0.8"

src/cleanlab_codex/codex_tool.py

@@ -114,7 +114,7 @@ def query(
         return self._project.query(
             question=question,
             fallback_answer=self._fallback_answer,
-            analytics_metadata=_AnalyticsMetadata(integration_type=IntegrationType.TOOL),
+            _analytics_metadata=_AnalyticsMetadata(integration_type=IntegrationType.TOOL),
         )[0]
 
     def to_openai_tool(self) -> dict[str, Any]:

src/cleanlab_codex/project.py

@@ -4,7 +4,7 @@
 
 from datetime import datetime
 from typing import TYPE_CHECKING as _TYPE_CHECKING
-from typing import Optional
+from typing import Any, Optional
 
 from codex import AuthenticationError
 
@@ -178,37 +178,43 @@ def query(
         question: str,
         *,
         fallback_answer: Optional[str] = None,
-        analytics_metadata: Optional[_AnalyticsMetadata] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        _analytics_metadata: Optional[_AnalyticsMetadata] = None,
     ) -> tuple[Optional[str], Entry]:
         """Query Codex to check if this project contains an answer to the question. If the question is not yet in the project, it will be added for SME review.
 
         Args:
             question (str): The question to ask the Codex API.
             fallback_answer (str, optional): Optional fallback answer to return if Codex is unable to answer the question.
+            metadata (dict, optional): Additional custom metadata to associate with the query.
 
         Returns:
             tuple[Optional[str], Entry]: A tuple representing the answer for the query and the existing or new entry in the Codex project.
                 If Codex is able to answer the question, the first element will be the answer returned by Codex and the second element will be the existing [`Entry`](/codex/api/python/types.entry#class-entry) in the Codex project.
                 If Codex is unable to answer the question, the first element will be `fallback_answer` if provided, otherwise None. The second element will be a new [`Entry`](/codex/api/python/types.entry#class-entry) in the Codex project.
         """
-        if not analytics_metadata:
-            analytics_metadata = _AnalyticsMetadata(integration_type=IntegrationType.BACKUP)
+        if not _analytics_metadata:
+            _analytics_metadata = _AnalyticsMetadata(integration_type=IntegrationType.BACKUP)
 
         return self._query_project(
             question=question,
             fallback_answer=fallback_answer,
-            analytics_metadata=analytics_metadata,
+            client_metadata=metadata,
+            analytics_metadata=_analytics_metadata,
         )
 
     def _query_project(
         self,
         question: str,
         *,
         fallback_answer: Optional[str] = None,
+        client_metadata: Optional[dict[str, Any]] = None,
         analytics_metadata: Optional[_AnalyticsMetadata] = None,
     ) -> tuple[Optional[str], Entry]:
         extra_headers = analytics_metadata.to_headers() if analytics_metadata else None
-        query_res = self._sdk_client.projects.entries.query(self._id, question=question, extra_headers=extra_headers)
+        query_res = self._sdk_client.projects.entries.query(
+            self._id, question=question, client_metadata=client_metadata, extra_headers=extra_headers
+        )
 
         entry = Entry.model_validate(query_res.entry.model_dump())
         if query_res.answer is not None:

tests/test_codex_tool.py

@@ -26,7 +26,7 @@ def test_tool_query_passes_metadata(mock_client_from_access_key: MagicMock) -> N
         assert kwargs["question"] == "what is the capital of France?"
         assert kwargs["fallback_answer"] == CodexTool.DEFAULT_FALLBACK_ANSWER
         assert (
-            kwargs["analytics_metadata"].to_headers()
+            kwargs["_analytics_metadata"].to_headers()
             == _AnalyticsMetadata(integration_type=IntegrationType.TOOL).to_headers()
         )
 

tests/test_project.py

@@ -184,6 +184,23 @@ def test_query_answer_found(mock_client_from_access_key: MagicMock) -> None:
     assert res[1].model_dump() == answered_entry.model_dump()
 
 
+def test_query_answer_found_with_metadata(mock_client_from_access_key: MagicMock) -> None:
+    answered_entry = SDKEntry(
+        id=str(uuid.uuid4()),
+        question="What is the capital of France?",
+        answer="Paris",
+        client_query_metadata=[{"trustworthiness_score": 0.95}],
+    )
+    mock_client_from_access_key.projects.entries.query.return_value = EntryQueryResponse(
+        answer="Paris", entry=answered_entry
+    )
+    project = Project(mock_client_from_access_key, FAKE_PROJECT_ID)
+    res = project.query("What is the capital of France?", metadata={"trustworthiness_score": 0.95})
+    assert res[0] == answered_entry.answer
+    assert res[1] is not None
+    assert res[1].model_dump() == answered_entry.model_dump()  # metadata should be included in the entry
+
+
 def test_add_entries_empty_list(mock_client_from_access_key: MagicMock) -> None:
     """Test adding an empty list of entries"""
     project = Project(mock_client_from_access_key, FAKE_PROJECT_ID)