Skip to content

Commit 10350c1

Browse files
feat(api): more pagination schemes
1 parent b764734 commit 10350c1

File tree

6 files changed

+175
-48
lines changed

6 files changed

+175
-48
lines changed

.stats.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
configured_endpoints: 55
22
openapi_spec_hash: f17890d85522687a4c68702da9ad2efb
3-
config_hash: 86582a50eb22b2866777cbd4d94f4e8d
3+
config_hash: 8f6e5c3b064cbb77569a6bf654954a56

api.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,7 @@ Methods:
214214
- <code title="get /api/projects/{project_id}/query_logs/{query_log_id}">client.projects.query_logs.<a href="./src/codex/resources/projects/query_logs.py">retrieve</a>(query_log_id, \*, project_id) -> <a href="./src/codex/types/projects/query_log_retrieve_response.py">QueryLogRetrieveResponse</a></code>
215215
- <code title="get /api/projects/{project_id}/query_logs/">client.projects.query_logs.<a href="./src/codex/resources/projects/query_logs.py">list</a>(project_id, \*\*<a href="src/codex/types/projects/query_log_list_params.py">params</a>) -> <a href="./src/codex/types/projects/query_log_list_response.py">SyncOffsetPageQueryLogs[QueryLogListResponse]</a></code>
216216
- <code title="get /api/projects/{project_id}/query_logs/logs_by_group">client.projects.query_logs.<a href="./src/codex/resources/projects/query_logs.py">list_by_group</a>(project_id, \*\*<a href="src/codex/types/projects/query_log_list_by_group_params.py">params</a>) -> <a href="./src/codex/types/projects/query_log_list_by_group_response.py">QueryLogListByGroupResponse</a></code>
217-
- <code title="get /api/projects/{project_id}/query_logs/groups">client.projects.query_logs.<a href="./src/codex/resources/projects/query_logs.py">list_groups</a>(project_id, \*\*<a href="src/codex/types/projects/query_log_list_groups_params.py">params</a>) -> <a href="./src/codex/types/projects/query_log_list_groups_response.py">QueryLogListGroupsResponse</a></code>
217+
- <code title="get /api/projects/{project_id}/query_logs/groups">client.projects.query_logs.<a href="./src/codex/resources/projects/query_logs.py">list_groups</a>(project_id, \*\*<a href="src/codex/types/projects/query_log_list_groups_params.py">params</a>) -> <a href="./src/codex/types/projects/query_log_list_groups_response.py">SyncOffsetPageQueryLogGroups[QueryLogListGroupsResponse]</a></code>
218218
- <code title="post /api/projects/{project_id}/query_logs/{query_log_id}/start_remediation">client.projects.query_logs.<a href="./src/codex/resources/projects/query_logs.py">start_remediation</a>(query_log_id, \*, project_id) -> <a href="./src/codex/types/projects/query_log_start_remediation_response.py">QueryLogStartRemediationResponse</a></code>
219219

220220
## Remediations

src/codex/pagination.py

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@
1818
"AsyncOffsetPageQueryLogs",
1919
"SyncOffsetPageRemediations",
2020
"AsyncOffsetPageRemediations",
21+
"SyncOffsetPageQueryLogGroups",
22+
"AsyncOffsetPageQueryLogGroups",
23+
"SyncOffsetPageQueryLogsByGroup",
24+
"AsyncOffsetPageQueryLogsByGroup",
2125
]
2226

2327
_BaseModelT = TypeVar("_BaseModelT", bound=BaseModel)
@@ -265,3 +269,123 @@ def next_page_info(self) -> Optional[PageInfo]:
265269
return PageInfo(params={"offset": current_count})
266270

267271
return None
272+
273+
274+
class SyncOffsetPageQueryLogGroups(BaseSyncPage[_T], BasePage[_T], Generic[_T]):
275+
query_log_groups: List[_T]
276+
total_count: Optional[int] = None
277+
278+
@override
279+
def _get_page_items(self) -> List[_T]:
280+
query_log_groups = self.query_log_groups
281+
if not query_log_groups:
282+
return []
283+
return query_log_groups
284+
285+
@override
286+
def next_page_info(self) -> Optional[PageInfo]:
287+
offset = self._options.params.get("offset") or 0
288+
if not isinstance(offset, int):
289+
raise ValueError(f'Expected "offset" param to be an integer but got {offset}')
290+
291+
length = len(self._get_page_items())
292+
current_count = offset + length
293+
294+
total_count = self.total_count
295+
if total_count is None:
296+
return None
297+
298+
if current_count < total_count:
299+
return PageInfo(params={"offset": current_count})
300+
301+
return None
302+
303+
304+
class AsyncOffsetPageQueryLogGroups(BaseAsyncPage[_T], BasePage[_T], Generic[_T]):
305+
query_log_groups: List[_T]
306+
total_count: Optional[int] = None
307+
308+
@override
309+
def _get_page_items(self) -> List[_T]:
310+
query_log_groups = self.query_log_groups
311+
if not query_log_groups:
312+
return []
313+
return query_log_groups
314+
315+
@override
316+
def next_page_info(self) -> Optional[PageInfo]:
317+
offset = self._options.params.get("offset") or 0
318+
if not isinstance(offset, int):
319+
raise ValueError(f'Expected "offset" param to be an integer but got {offset}')
320+
321+
length = len(self._get_page_items())
322+
current_count = offset + length
323+
324+
total_count = self.total_count
325+
if total_count is None:
326+
return None
327+
328+
if current_count < total_count:
329+
return PageInfo(params={"offset": current_count})
330+
331+
return None
332+
333+
334+
class SyncOffsetPageQueryLogsByGroup(BaseSyncPage[_T], BasePage[_T], Generic[_T]):
335+
query_logs_by_group: List[_T]
336+
total_count: Optional[int] = None
337+
338+
@override
339+
def _get_page_items(self) -> List[_T]:
340+
query_logs_by_group = self.query_logs_by_group
341+
if not query_logs_by_group:
342+
return []
343+
return query_logs_by_group
344+
345+
@override
346+
def next_page_info(self) -> Optional[PageInfo]:
347+
offset = self._options.params.get("offset") or 0
348+
if not isinstance(offset, int):
349+
raise ValueError(f'Expected "offset" param to be an integer but got {offset}')
350+
351+
length = len(self._get_page_items())
352+
current_count = offset + length
353+
354+
total_count = self.total_count
355+
if total_count is None:
356+
return None
357+
358+
if current_count < total_count:
359+
return PageInfo(params={"offset": current_count})
360+
361+
return None
362+
363+
364+
class AsyncOffsetPageQueryLogsByGroup(BaseAsyncPage[_T], BasePage[_T], Generic[_T]):
365+
query_logs_by_group: List[_T]
366+
total_count: Optional[int] = None
367+
368+
@override
369+
def _get_page_items(self) -> List[_T]:
370+
query_logs_by_group = self.query_logs_by_group
371+
if not query_logs_by_group:
372+
return []
373+
return query_logs_by_group
374+
375+
@override
376+
def next_page_info(self) -> Optional[PageInfo]:
377+
offset = self._options.params.get("offset") or 0
378+
if not isinstance(offset, int):
379+
raise ValueError(f'Expected "offset" param to be an integer but got {offset}')
380+
381+
length = len(self._get_page_items())
382+
current_count = offset + length
383+
384+
total_count = self.total_count
385+
if total_count is None:
386+
return None
387+
388+
if current_count < total_count:
389+
return PageInfo(params={"offset": current_count})
390+
391+
return None

src/codex/resources/projects/query_logs.py

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,12 @@
1818
async_to_raw_response_wrapper,
1919
async_to_streamed_response_wrapper,
2020
)
21-
from ...pagination import SyncOffsetPageQueryLogs, AsyncOffsetPageQueryLogs
21+
from ...pagination import (
22+
SyncOffsetPageQueryLogs,
23+
AsyncOffsetPageQueryLogs,
24+
SyncOffsetPageQueryLogGroups,
25+
AsyncOffsetPageQueryLogGroups,
26+
)
2227
from ..._base_client import AsyncPaginator, make_request_options
2328
from ...types.projects import query_log_list_params, query_log_list_groups_params, query_log_list_by_group_params
2429
from ...types.projects.query_log_list_response import QueryLogListResponse
@@ -290,7 +295,7 @@ def list_groups(
290295
extra_query: Query | None = None,
291296
extra_body: Body | None = None,
292297
timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
293-
) -> QueryLogListGroupsResponse:
298+
) -> SyncOffsetPageQueryLogGroups[QueryLogListGroupsResponse]:
294299
"""
295300
List query log groups by project ID.
296301
@@ -323,8 +328,9 @@ def list_groups(
323328
"""
324329
if not project_id:
325330
raise ValueError(f"Expected a non-empty value for `project_id` but received {project_id!r}")
326-
return self._get(
331+
return self._get_api_list(
327332
f"/api/projects/{project_id}/query_logs/groups",
333+
page=SyncOffsetPageQueryLogGroups[QueryLogListGroupsResponse],
328334
options=make_request_options(
329335
extra_headers=extra_headers,
330336
extra_query=extra_query,
@@ -349,7 +355,7 @@ def list_groups(
349355
query_log_list_groups_params.QueryLogListGroupsParams,
350356
),
351357
),
352-
cast_to=QueryLogListGroupsResponse,
358+
model=QueryLogListGroupsResponse,
353359
)
354360

355361
def start_remediation(
@@ -622,7 +628,7 @@ async def list_by_group(
622628
cast_to=QueryLogListByGroupResponse,
623629
)
624630

625-
async def list_groups(
631+
def list_groups(
626632
self,
627633
project_id: str,
628634
*,
@@ -649,7 +655,7 @@ async def list_groups(
649655
extra_query: Query | None = None,
650656
extra_body: Body | None = None,
651657
timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
652-
) -> QueryLogListGroupsResponse:
658+
) -> AsyncPaginator[QueryLogListGroupsResponse, AsyncOffsetPageQueryLogGroups[QueryLogListGroupsResponse]]:
653659
"""
654660
List query log groups by project ID.
655661
@@ -682,14 +688,15 @@ async def list_groups(
682688
"""
683689
if not project_id:
684690
raise ValueError(f"Expected a non-empty value for `project_id` but received {project_id!r}")
685-
return await self._get(
691+
return self._get_api_list(
686692
f"/api/projects/{project_id}/query_logs/groups",
693+
page=AsyncOffsetPageQueryLogGroups[QueryLogListGroupsResponse],
687694
options=make_request_options(
688695
extra_headers=extra_headers,
689696
extra_query=extra_query,
690697
extra_body=extra_body,
691698
timeout=timeout,
692-
query=await async_maybe_transform(
699+
query=maybe_transform(
693700
{
694701
"created_at_end": created_at_end,
695702
"created_at_start": created_at_start,
@@ -708,7 +715,7 @@ async def list_groups(
708715
query_log_list_groups_params.QueryLogListGroupsParams,
709716
),
710717
),
711-
cast_to=QueryLogListGroupsResponse,
718+
model=QueryLogListGroupsResponse,
712719
)
713720

714721
async def start_remediation(

src/codex/types/projects/query_log_list_groups_response.py

Lines changed: 19 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -8,41 +8,40 @@
88

99
__all__ = [
1010
"QueryLogListGroupsResponse",
11-
"QueryLogGroup",
12-
"QueryLogGroupFormattedEscalationEvalScores",
13-
"QueryLogGroupFormattedEvalScores",
14-
"QueryLogGroupFormattedGuardrailEvalScores",
15-
"QueryLogGroupFormattedNonGuardrailEvalScores",
16-
"QueryLogGroupContext",
17-
"QueryLogGroupDeterministicGuardrailsResults",
11+
"FormattedEscalationEvalScores",
12+
"FormattedEvalScores",
13+
"FormattedGuardrailEvalScores",
14+
"FormattedNonGuardrailEvalScores",
15+
"Context",
16+
"DeterministicGuardrailsResults",
1817
]
1918

2019

21-
class QueryLogGroupFormattedEscalationEvalScores(BaseModel):
20+
class FormattedEscalationEvalScores(BaseModel):
2221
score: float
2322

2423
status: Literal["pass", "fail"]
2524

2625

27-
class QueryLogGroupFormattedEvalScores(BaseModel):
26+
class FormattedEvalScores(BaseModel):
2827
score: float
2928

3029
status: Literal["pass", "fail"]
3130

3231

33-
class QueryLogGroupFormattedGuardrailEvalScores(BaseModel):
32+
class FormattedGuardrailEvalScores(BaseModel):
3433
score: float
3534

3635
status: Literal["pass", "fail"]
3736

3837

39-
class QueryLogGroupFormattedNonGuardrailEvalScores(BaseModel):
38+
class FormattedNonGuardrailEvalScores(BaseModel):
4039
score: float
4140

4241
status: Literal["pass", "fail"]
4342

4443

45-
class QueryLogGroupContext(BaseModel):
44+
class Context(BaseModel):
4645
content: str
4746
"""The actual content/text of the document."""
4847

@@ -59,32 +58,32 @@ class QueryLogGroupContext(BaseModel):
5958
"""Title or heading of the document. Useful for display and context."""
6059

6160

62-
class QueryLogGroupDeterministicGuardrailsResults(BaseModel):
61+
class DeterministicGuardrailsResults(BaseModel):
6362
guardrail_name: str
6463

6564
should_guardrail: bool
6665

6766
matches: Optional[List[str]] = None
6867

6968

70-
class QueryLogGroup(BaseModel):
69+
class QueryLogListGroupsResponse(BaseModel):
7170
id: str
7271

7372
created_at: datetime
7473

75-
formatted_escalation_eval_scores: Optional[Dict[str, QueryLogGroupFormattedEscalationEvalScores]] = None
74+
formatted_escalation_eval_scores: Optional[Dict[str, FormattedEscalationEvalScores]] = None
7675

77-
formatted_eval_scores: Optional[Dict[str, QueryLogGroupFormattedEvalScores]] = None
76+
formatted_eval_scores: Optional[Dict[str, FormattedEvalScores]] = None
7877
"""Format evaluation scores for frontend display with pass/fail status.
7978
8079
Returns: Dictionary mapping eval keys to their formatted representation: {
8180
"eval_key": { "score": float, "status": "pass" | "fail" } } Returns None if
8281
eval_scores is None.
8382
"""
8483

85-
formatted_guardrail_eval_scores: Optional[Dict[str, QueryLogGroupFormattedGuardrailEvalScores]] = None
84+
formatted_guardrail_eval_scores: Optional[Dict[str, FormattedGuardrailEvalScores]] = None
8685

87-
formatted_non_guardrail_eval_scores: Optional[Dict[str, QueryLogGroupFormattedNonGuardrailEvalScores]] = None
86+
formatted_non_guardrail_eval_scores: Optional[Dict[str, FormattedNonGuardrailEvalScores]] = None
8887

8988
is_bad_response: bool
9089

@@ -103,7 +102,7 @@ class QueryLogGroup(BaseModel):
103102
was_cache_hit: Optional[bool] = None
104103
"""If similar query already answered, or None if cache was not checked"""
105104

106-
context: Optional[List[QueryLogGroupContext]] = None
105+
context: Optional[List[Context]] = None
107106
"""RAG context used for the query"""
108107

109108
custom_metadata: Optional[object] = None
@@ -112,7 +111,7 @@ class QueryLogGroup(BaseModel):
112111
custom_metadata_keys: Optional[List[str]] = None
113112
"""Keys of the custom metadata"""
114113

115-
deterministic_guardrails_results: Optional[Dict[str, QueryLogGroupDeterministicGuardrailsResults]] = None
114+
deterministic_guardrails_results: Optional[Dict[str, DeterministicGuardrailsResults]] = None
116115
"""Results of deterministic guardrails applied to the query"""
117116

118117
escalated: Optional[bool] = None
@@ -144,11 +143,3 @@ class QueryLogGroup(BaseModel):
144143

145144
primary_eval_issue_score: Optional[float] = None
146145
"""Score of the primary eval issue"""
147-
148-
149-
class QueryLogListGroupsResponse(BaseModel):
150-
custom_metadata_columns: List[str]
151-
152-
query_log_groups: List[QueryLogGroup]
153-
154-
total_count: int

0 commit comments

Comments
 (0)