Skip to content

Commit df5d394

Browse files
author
michal.zyga
committed
add list batches method
1 parent 8da3dbd commit df5d394

File tree

7 files changed

+122
-31
lines changed

7 files changed

+122
-31
lines changed

core/src/main/scala/sttp/openai/OpenAI.scala

+20-2
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,15 @@ import sttp.openai.requests.audio.AudioResponseData.AudioResponse
1717
import sttp.openai.requests.audio.RecognitionModel
1818
import sttp.openai.requests.audio.transcriptions.TranscriptionConfig
1919
import sttp.openai.requests.audio.translations.TranslationConfig
20-
import sttp.openai.requests.batch.{BatchRequestBody, BatchResponse}
20+
import sttp.openai.requests.batch.{QueryParameters => _, _}
2121
import sttp.openai.requests.completions.CompletionsRequestBody.CompletionsBody
2222
import sttp.openai.requests.completions.CompletionsResponseData.CompletionsResponse
2323
import sttp.openai.requests.completions.chat.ChatRequestBody.ChatBody
2424
import sttp.openai.requests.completions.chat.ChatRequestResponseData.ChatResponse
2525
import sttp.openai.requests.embeddings.EmbeddingsRequestBody.EmbeddingsBody
2626
import sttp.openai.requests.embeddings.EmbeddingsResponseBody.EmbeddingResponse
2727
import sttp.openai.requests.files.FilesResponseData._
28-
import sttp.openai.requests.finetuning
28+
import sttp.openai.requests.{batch, finetuning}
2929
import sttp.openai.requests.finetuning._
3030
import sttp.openai.requests.images.ImageResponseData.ImageResponse
3131
import sttp.openai.requests.images.creation.ImageCreationRequestBody.ImageCreationBody
@@ -1159,6 +1159,24 @@ class OpenAI(authToken: String, baseUri: Uri = OpenAIUris.OpenAIBaseUri) {
11591159
.post(openAIUris.cancelBatch(batchId))
11601160
.response(asJson_parseErrors[BatchResponse])
11611161

1162+
/** List your organization's batches.
1163+
*
1164+
* [[https://platform.openai.com/docs/api-reference/batch/list]]
1165+
*
1166+
* @return
1167+
* A list of paginated Batch objects.
1168+
*/
1169+
def listBatches(
1170+
queryParameters: batch.QueryParameters = batch.QueryParameters.empty
1171+
): Request[Either[OpenAIException, ListBatchResponse]] = {
1172+
val uri = openAIUris.Batches
1173+
.withParams(queryParameters.toMap)
1174+
1175+
openAIAuthRequest
1176+
.get(uri)
1177+
.response(asJson_parseErrors[ListBatchResponse])
1178+
}
1179+
11621180
protected val openAIAuthRequest: PartialRequest[Either[String, String]] = basicRequest.auth
11631181
.bearer(authToken)
11641182

core/src/main/scala/sttp/openai/OpenAISyncClient.scala

+12-2
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,14 @@ import sttp.openai.requests.audio.AudioResponseData.AudioResponse
99
import sttp.openai.requests.audio.RecognitionModel
1010
import sttp.openai.requests.audio.transcriptions.TranscriptionConfig
1111
import sttp.openai.requests.audio.translations.TranslationConfig
12-
import sttp.openai.requests.batch.{BatchRequestBody, BatchResponse}
12+
import sttp.openai.requests.batch.{BatchRequestBody, BatchResponse, ListBatchResponse}
1313
import sttp.openai.requests.completions.CompletionsRequestBody.CompletionsBody
1414
import sttp.openai.requests.completions.CompletionsResponseData.CompletionsResponse
1515
import sttp.openai.requests.completions.chat.ChatRequestBody.ChatBody
1616
import sttp.openai.requests.completions.chat.ChatRequestResponseData.ChatResponse
1717
import sttp.openai.requests.embeddings.EmbeddingsRequestBody.EmbeddingsBody
1818
import sttp.openai.requests.embeddings.EmbeddingsResponseBody.EmbeddingResponse
1919
import sttp.openai.requests.files.FilesResponseData.{DeletedFileData, FileData, FilesResponse}
20-
import sttp.openai.requests.finetuning
2120
import sttp.openai.requests.finetuning._
2221
import sttp.openai.requests.images.ImageResponseData.ImageResponse
2322
import sttp.openai.requests.images.creation.ImageCreationRequestBody.ImageCreationBody
@@ -41,6 +40,7 @@ import sttp.openai.requests.vectorstore.file.VectorStoreFileResponseData.{
4140
ListVectorStoreFilesResponse,
4241
VectorStoreFile
4342
}
43+
import sttp.openai.requests.{batch, finetuning}
4444

4545
import java.io.File
4646

@@ -836,6 +836,16 @@ class OpenAISyncClient private (
836836
def cancelBatch(batchId: String): BatchResponse =
837837
sendOrThrow(openAI.cancelBatch(batchId))
838838

839+
/** List your organization's batches.
840+
*
841+
* [[https://platform.openai.com/docs/api-reference/batch/list]]
842+
*
843+
* @return
844+
* A list of paginated Batch objects.
845+
*/
846+
def listBatches(queryParameters: batch.QueryParameters = batch.QueryParameters.empty): ListBatchResponse =
847+
sendOrThrow(openAI.listBatches(queryParameters))
848+
839849
/** Closes and releases resources of http client if was not provided explicitly, otherwise works no-op. */
840850
def close(): Unit = if (closeClient) backend.close() else ()
841851

core/src/main/scala/sttp/openai/requests/batch/BatchResponse.scala

+12
Original file line numberDiff line numberDiff line change
@@ -117,3 +117,15 @@ case class RequestCounts(
117117
object RequestCounts {
118118
implicit val requestCountsR: SnakePickle.Reader[RequestCounts] = SnakePickle.macroR[RequestCounts]
119119
}
120+
121+
case class ListBatchResponse(
122+
`object`: String = "list",
123+
data: Seq[BatchResponse],
124+
firstId: String,
125+
lastId: String,
126+
hasMore: Boolean
127+
)
128+
129+
object ListBatchResponse {
130+
implicit val listBatchResponseR: SnakePickle.Reader[ListBatchResponse] = SnakePickle.macroR[ListBatchResponse]
131+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
package sttp.openai.requests.batch
2+
3+
/** @param after
4+
* A cursor for use in pagination. after is an object ID that defines your place in the list. For instance, if you make a list request
5+
* and receive 100 objects, ending with obj_foo, your subsequent call can include after=obj_foo in order to fetch the next page of the
6+
* list.
7+
* @param limit
8+
* A limit on the number of objects to be returned. Limit can range between 1 and 100, and the default is 20.
9+
*/
10+
case class QueryParameters(
11+
after: Option[String] = None,
12+
limit: Option[Int] = None
13+
) {
14+
15+
def toMap: Map[String, String] = {
16+
val queryParams = after.map("after" -> _) ++
17+
limit.map(_.toString).map("limit" -> _)
18+
queryParams.toMap
19+
}
20+
}
21+
22+
object QueryParameters {
23+
val empty: QueryParameters = QueryParameters(None, None)
24+
}

core/src/main/scala/sttp/openai/requests/finetuning/QueryParameters.scala

+3-5
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
11
package sttp.openai.requests.finetuning
22

33
/** @param after
4-
* A cursor for use in pagination. after is an object ID that defines your place in the list. For instance, if you make a list request
5-
* and receive 100 objects, ending with obj_foo, your subsequent call can include after=obj_foo in order to fetch the next page of the
6-
* list
4+
* Identifier for the last job from the previous pagination request.
75
* @param limit
8-
* Defaults to 20 A limit on the number of objects to be returned. Limit can range between 1 and 100, and the default is 20.
6+
* Number of fine-tuning jobs to retrieve.
97
*/
108
case class QueryParameters(
119
after: Option[String] = None,
@@ -14,7 +12,7 @@ case class QueryParameters(
1412

1513
def toMap: Map[String, String] = {
1614
val queryParams = after.map("after" -> _) ++
17-
limit.map(_.toString).map("order" -> _)
15+
limit.map(_.toString).map("limit" -> _)
1816
queryParams.toMap
1917
}
2018
}

core/src/test/scala/sttp/openai/fixtures/BatchFixture.scala

+32
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package sttp.openai.fixtures
22

3+
import sttp.openai.requests.batch.{BatchResponse, RequestCounts}
4+
35
object BatchFixture {
46

57
val jsonCreateBatchRequest: String = """{
@@ -42,4 +44,34 @@ object BatchFixture {
4244
| }
4345
|}""".stripMargin
4446

47+
val jsonListBatchResponse: String = s"""{
48+
| "object": "list",
49+
| "data": [$jsonCreateBatchResponse],
50+
| "first_id": "ftckpt_zc4Q7MP6XxulcVzj4MZdwsAB",
51+
| "last_id": "ftckpt_enQCFmOTGj3syEpYVhBRLTSy",
52+
| "has_more": true
53+
|}""".stripMargin
54+
55+
val batchResponse: BatchResponse = BatchResponse(
56+
id = "batch_abc123",
57+
endpoint = "/v1/completions",
58+
errors = None,
59+
inputFileId = "file-abc123",
60+
completionWindow = "24h",
61+
status = "completed",
62+
outputFileId = Some("file-cvaTdG"),
63+
errorFileId = Some("file-HOWS94"),
64+
createdAt = 1711471533,
65+
inProgressAt = Some(1711471538),
66+
expiresAt = Some(1711557933),
67+
finalizingAt = Some(1711493133),
68+
completedAt = Some(1711493163),
69+
failedAt = None,
70+
expiredAt = None,
71+
cancellingAt = None,
72+
cancelledAt = None,
73+
requestCounts = Some(RequestCounts(total = 100, completed = 95, failed = 5)),
74+
metadata = Some(Map("customer_id" -> "user_123456789", "batch_description" -> "Nightly eval job"))
75+
)
76+
4577
}

core/src/test/scala/sttp/openai/requests/BatchDataSpec.scala

+19-22
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import org.scalatest.flatspec.AnyFlatSpec
55
import org.scalatest.matchers.should.Matchers
66
import sttp.openai.fixtures.BatchFixture
77
import sttp.openai.json.{SnakePickle, SttpUpickleApiExtension}
8-
import sttp.openai.requests.batch.{BatchRequestBody, BatchResponse, RequestCounts}
8+
import sttp.openai.requests.batch.{BatchRequestBody, BatchResponse, ListBatchResponse}
99

1010
class BatchDataSpec extends AnyFlatSpec with Matchers with EitherValues {
1111

@@ -27,32 +27,29 @@ class BatchDataSpec extends AnyFlatSpec with Matchers with EitherValues {
2727
"Given create batch response as Json" should "be properly deserialized to case class" in {
2828
// given
2929
val jsonResponse = BatchFixture.jsonCreateBatchResponse
30-
val expectedResponse: BatchResponse = BatchResponse(
31-
id = "batch_abc123",
32-
endpoint = "/v1/completions",
33-
errors = None,
34-
inputFileId = "file-abc123",
35-
completionWindow = "24h",
36-
status = "completed",
37-
outputFileId = Some("file-cvaTdG"),
38-
errorFileId = Some("file-HOWS94"),
39-
createdAt = 1711471533,
40-
inProgressAt = Some(1711471538),
41-
expiresAt = Some(1711557933),
42-
finalizingAt = Some(1711493133),
43-
completedAt = Some(1711493163),
44-
failedAt = None,
45-
expiredAt = None,
46-
cancellingAt = None,
47-
cancelledAt = None,
48-
requestCounts = Some(RequestCounts(total = 100, completed = 95, failed = 5)),
49-
metadata = Some(Map("customer_id" -> "user_123456789", "batch_description" -> "Nightly eval job"))
50-
)
30+
val expectedResponse: BatchResponse = BatchFixture.batchResponse
5131
// when
5232
val deserializedJsonResponse: Either[Exception, BatchResponse] =
5333
SttpUpickleApiExtension.deserializeJsonSnake[BatchResponse].apply(jsonResponse)
5434
// then
5535
deserializedJsonResponse.value shouldBe expectedResponse
5636
}
5737

38+
"Given list batch response as Json" should "be properly deserialized to case class" in {
39+
// given
40+
val jsonResponse = BatchFixture.jsonListBatchResponse
41+
val expectedResponse: ListBatchResponse = ListBatchResponse(
42+
data = Seq(BatchFixture.batchResponse),
43+
hasMore = true,
44+
firstId = "ftckpt_zc4Q7MP6XxulcVzj4MZdwsAB",
45+
lastId = "ftckpt_enQCFmOTGj3syEpYVhBRLTSy"
46+
)
47+
// when
48+
val deserializedJsonResponse: Either[Exception, ListBatchResponse] =
49+
SttpUpickleApiExtension.deserializeJsonSnake[ListBatchResponse].apply(jsonResponse)
50+
51+
// then
52+
deserializedJsonResponse.value shouldBe expectedResponse
53+
}
54+
5855
}

0 commit comments

Comments
 (0)