Skip to content

Commit b863224

Browse files
committed
delete vector store, create/delete/list vector store files endpoints
1 parent 181ecf1 commit b863224

File tree

16 files changed

+543
-12
lines changed

16 files changed

+543
-12
lines changed

openai-client/src/main/scala/io/cequence/openaiscala/JsonFormats.scala

Lines changed: 65 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import io.cequence.openaiscala.domain.AssistantToolResource.{
55
FileSearchResources
66
}
77
import io.cequence.openaiscala.domain.Batch._
8+
import io.cequence.openaiscala.domain.ChunkingStrategy.StaticChunkingStrategy
89
import io.cequence.openaiscala.domain.FineTune.WeightsAndBiases
910
import io.cequence.openaiscala.domain.response.AssistantToolResourceResponse.{
1011
CodeInterpreterResourcesResponse,
@@ -18,9 +19,11 @@ import io.cequence.openaiscala.domain.response.ResponseFormat.{
1819
import io.cequence.openaiscala.domain.response._
1920
import io.cequence.openaiscala.domain.{ThreadMessageFile, _}
2021
import io.cequence.wsclient.JsonUtil
21-
import io.cequence.wsclient.JsonUtil.enumFormat
22+
import io.cequence.wsclient.JsonUtil.{enumFormat, snakeEnumFormat}
23+
import io.cequence.wsclient.domain.EnumValue
2224
import play.api.libs.functional.syntax._
2325
import play.api.libs.json.Json.toJson
26+
import play.api.libs.json.JsonNaming.SnakeCase
2427
import play.api.libs.json.{Format, JsValue, Json, _}
2528

2629
import java.{util => ju}
@@ -714,9 +717,69 @@ object JsonFormats {
714717
implicit lazy val createBatchResponsesFormat: Format[CreateBatchResponses] =
715718
Json.format[CreateBatchResponses]
716719

717-
implicit lazy val fileCountsFormat: Format[FileCounts] =
720+
implicit lazy val fileCountsFormat: Format[FileCounts] = {
721+
implicit val config: JsonConfiguration = JsonConfiguration(SnakeCase)
718722
Json.format[FileCounts]
723+
}
719724

720725
implicit lazy val vectorStoreFormat: Format[VectorStore] =
721726
Json.format[VectorStore]
727+
728+
implicit lazy val vectorStoreFileFormat: Format[VectorStoreFile] = {
729+
implicit val config: JsonConfiguration = JsonConfiguration(SnakeCase)
730+
Json.format[VectorStoreFile]
731+
}
732+
733+
implicit lazy val vectorStoreFileStatusFormat: Format[VectorStoreFileStatus] = {
734+
import VectorStoreFileStatus._
735+
enumFormat(Cancelled, Completed, InProgress, Failed)
736+
}
737+
738+
implicit lazy val lastErrorFormat: Format[LastError] = Json.format[LastError]
739+
implicit lazy val lastErrorCodeFormat: Format[LastErrorCode] = {
740+
import LastErrorCode._
741+
snakeEnumFormat(ServerError, RateLimitExceeded)
742+
}
743+
implicit lazy val chunkingStrategyAutoFormat
744+
: Format[ChunkingStrategy.AutoChunkingStrategy.type] =
745+
Json.format[ChunkingStrategy.AutoChunkingStrategy.type]
746+
implicit lazy val chunkingStrategyStaticFormat
747+
: Format[ChunkingStrategy.StaticChunkingStrategy.type] =
748+
Json.format[ChunkingStrategy.StaticChunkingStrategy.type]
749+
750+
val chunkingStrategyFormatReads: Reads[ChunkingStrategy] =
751+
(
752+
(__ \ "max_chunk_size_tokens").readNullable[Int] and
753+
(__ \ "chunk_overlap_tokens").readNullable[Int]
754+
)(
755+
(
756+
maxChunkSizeTokens: Option[Int],
757+
chunkOverlapTokens: Option[Int]
758+
) => StaticChunkingStrategy(maxChunkSizeTokens, chunkOverlapTokens)
759+
)
760+
761+
implicit lazy val chunkingStrategyFormat: Format[ChunkingStrategy] = {
762+
val reads: Reads[ChunkingStrategy] = Reads { json =>
763+
import ChunkingStrategy._
764+
(json \ "type").validate[String].flatMap {
765+
case "auto" => JsSuccess(AutoChunkingStrategy)
766+
case "static" =>
767+
(json.validate[ChunkingStrategy](chunkingStrategyFormatReads))
768+
case "" => JsSuccess(AutoChunkingStrategy)
769+
case _ => JsError("Unknown chunking strategy type")
770+
}
771+
}
772+
773+
val writes: Writes[ChunkingStrategy] = Writes {
774+
case ChunkingStrategy.AutoChunkingStrategy => Json.obj("type" -> "auto")
775+
case ChunkingStrategy.StaticChunkingStrategy(maxChunkSizeTokens, chunkOverlapTokens) =>
776+
Json.obj(
777+
"type" -> "static",
778+
"max_chunk_size_tokens" -> maxChunkSizeTokens,
779+
"chunk_overlap_tokens" -> chunkOverlapTokens
780+
)
781+
}
782+
783+
Format(reads, writes)
784+
}
722785
}

openai-client/src/main/scala/io/cequence/openaiscala/service/impl/EndPoint.scala

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ object EndPoint {
2323
case object batches extends EndPoint
2424
case object assistants extends EndPoint
2525
case object vector_stores extends EndPoint
26+
case object vector_store_files extends EndPoint("vector_stores/files")
2627
}
2728

2829
sealed trait Param extends EnumValue
@@ -100,4 +101,6 @@ object Param {
100101
case object input_file_id extends Param
101102
case object endpoint extends Param
102103
case object completion_window extends Param
104+
case object chunking_strategy extends Param
105+
case object filter extends Param
103106
}

openai-client/src/main/scala/io/cequence/openaiscala/service/impl/OpenAIServiceImpl.scala

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,55 @@ private[service] trait OpenAIServiceImpl
379379
readAttribute(response, "data").asSafeArray[VectorStore]
380380
}
381381

382+
override def deleteVectorStore(
383+
vectorStoreId: String
384+
): Future[DeleteResponse] =
385+
execDELETEWithStatus(
386+
EndPoint.vector_stores,
387+
endPointParam = Some(vectorStoreId)
388+
).map(handleDeleteEndpointResponse)
389+
390+
override def createVectorStoreFile(
391+
vectorStoreId: String,
392+
fileId: String,
393+
chunkingStrategy: ChunkingStrategy = ChunkingStrategy.AutoChunkingStrategy
394+
): Future[VectorStoreFile] =
395+
execPOST(
396+
EndPoint.vector_stores,
397+
endPointParam = Some(s"$vectorStoreId/files"),
398+
bodyParams = jsonBodyParams(
399+
Param.file_id -> Some(fileId),
400+
Param.chunking_strategy -> Some(Json.toJson(chunkingStrategy))
401+
)
402+
).map(
403+
_.asSafe[VectorStoreFile]
404+
)
405+
406+
override def listVectorStoreFiles(
407+
vectorStoreId: String,
408+
pagination: Pagination = Pagination.default,
409+
order: Option[SortOrder] = None,
410+
filter: Option[VectorStoreFileStatus] = None
411+
): Future[Seq[VectorStoreFile]] =
412+
execGET(
413+
EndPoint.vector_stores,
414+
endPointParam = Some(s"$vectorStoreId/files"),
415+
params = paginationParams(pagination) :+
416+
Param.order -> order :+
417+
Param.filter -> filter
418+
).map { response =>
419+
readAttribute(response, "data").asSafeArray[VectorStoreFile]
420+
}
421+
422+
override def deleteVectorStoreFile(
423+
vectorStoreId: String,
424+
fileId: String
425+
): Future[DeleteResponse] =
426+
execDELETEWithStatus(
427+
EndPoint.vector_stores,
428+
endPointParam = Some(s"$vectorStoreId/files/$fileId")
429+
).map(handleDeleteEndpointResponse)
430+
382431
override def createFineTune(
383432
training_file: String,
384433
validation_file: Option[String] = None,

openai-client/src/test/scala/io/cequence/openaiscala/JsonFormatsSpec.scala

Lines changed: 135 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@ object JsonFormatsSpec {
3232
}
3333
}
3434

35-
@Ignore
3635
class JsonFormatsSpec extends AnyWordSpecLike with Matchers {
3736

3837
private val textResponseJson =
@@ -102,6 +101,15 @@ class JsonFormatsSpec extends AnyWordSpecLike with Matchers {
102101
| }
103102
|}""".stripMargin
104103

104+
private val fileCountsJson =
105+
"""{
106+
| "in_progress" : 1,
107+
| "completed" : 2,
108+
| "cancelled" : 3,
109+
| "failed" : 4,
110+
| "total" : 10
111+
|}""".stripMargin
112+
105113
private val attachmentJson =
106114
"""{
107115
| "file_id" : {
@@ -136,7 +144,7 @@ class JsonFormatsSpec extends AnyWordSpecLike with Matchers {
136144
)
137145
}
138146

139-
"serialize and deserialize file search's resources" in {
147+
"serialize and deserialize file search's resources" ignore {
140148
testCodec[AssistantToolResource](
141149
FileSearchResources(
142150
Seq("vs_xxx"),
@@ -188,6 +196,131 @@ class JsonFormatsSpec extends AnyWordSpecLike with Matchers {
188196
)
189197
}
190198

199+
"serialize and deserialize FileCounts" in {
200+
val integration = FileCounts(
201+
inProgress = 1,
202+
completed = 2,
203+
cancelled = 3,
204+
failed = 4,
205+
total = 10
206+
)
207+
testCodec[FileCounts](
208+
integration,
209+
fileCountsJson,
210+
Pretty
211+
)
212+
}
213+
214+
"serialize and deserialize VectorStore" in {
215+
val vectorStore = VectorStore(
216+
fileIds = Seq(FileId("file-123")),
217+
metadata = Map("key" -> "value")
218+
)
219+
testCodec[VectorStore](
220+
vectorStore,
221+
"""{
222+
| "file_ids" : [ {
223+
| "file_id" : "file-123"
224+
| } ],
225+
| "metadata" : {
226+
| "key" : "value"
227+
| }
228+
|}""".stripMargin,
229+
Pretty
230+
)
231+
}
232+
233+
"serialize and deserialize VectorStoreFileStatus" in {
234+
import VectorStoreFileStatus._
235+
testCodec[VectorStoreFileStatus](Cancelled, "\"cancelled\"".stripMargin, Pretty)
236+
testCodec[VectorStoreFileStatus](Completed, "\"completed\"".stripMargin, Pretty)
237+
testCodec[VectorStoreFileStatus](InProgress, "\"in_progress\"".stripMargin, Pretty)
238+
testCodec[VectorStoreFileStatus](Failed, "\"failed\"".stripMargin, Pretty)
239+
}
240+
241+
"serialize and deserialize LastErrorCode" in {
242+
import LastErrorCode._
243+
testCodec[LastErrorCode](ServerError, "\"server_error\"".stripMargin, Pretty)
244+
testCodec[LastErrorCode](
245+
RateLimitExceeded,
246+
"\"rate_limit_exceeded\"".stripMargin,
247+
Pretty
248+
)
249+
}
250+
251+
"serialize and deserialize ChunkingStrategy" in {
252+
import ChunkingStrategy._
253+
testCodec[ChunkingStrategy](
254+
AutoChunkingStrategy,
255+
"""{
256+
| "type" : "auto"
257+
|}""".stripMargin.stripMargin,
258+
Pretty
259+
)
260+
testCodec[ChunkingStrategy](
261+
StaticChunkingStrategy(None, None),
262+
"""{
263+
| "type" : "static",
264+
| "max_chunk_size_tokens" : 800,
265+
| "chunk_overlap_tokens" : 400
266+
|}""".stripMargin,
267+
Pretty
268+
)
269+
testDeserialization[ChunkingStrategy](
270+
StaticChunkingStrategy(None, None),
271+
"""{
272+
| "type" : "static"
273+
|}""".stripMargin
274+
)
275+
testDeserialization[ChunkingStrategy](
276+
StaticChunkingStrategy(Some(1000), None),
277+
"""{
278+
| "type" : "static",
279+
| "max_chunk_size_tokens" : 1000
280+
|}""".stripMargin
281+
)
282+
testDeserialization[ChunkingStrategy](
283+
StaticChunkingStrategy(None, Some(150)),
284+
"""{
285+
| "type" : "static",
286+
| "chunk_overlap_tokens" : 150
287+
|}""".stripMargin
288+
)
289+
}
290+
291+
"serialize and deserialize VectorStoreFile" in {
292+
val vectorStoreFile = VectorStoreFile(
293+
id = "file-id-1",
294+
`object` = "file",
295+
usageBytes = 100,
296+
createdAt = 1000,
297+
vectorStoreId = "vs-123",
298+
status = VectorStoreFileStatus.Completed,
299+
lastError = Some(LastError(LastErrorCode.ServerError, "error message")),
300+
chunkingStrategy = ChunkingStrategy.StaticChunkingStrategy(Some(1000), Some(500))
301+
)
302+
testCodec[VectorStoreFile](
303+
vectorStoreFile,
304+
"""{
305+
| "id" : "file-id-1",
306+
| "object" : "file",
307+
| "usage_bytes" : 100,
308+
| "created_at" : 1000,
309+
| "vector_store_id" : "vs-123",
310+
| "status" : "completed",
311+
| "last_error" : {
312+
| "code" : "server_error",
313+
| "message" : "error message"
314+
| },
315+
| "chunking_strategy" : {
316+
| "type" : "static",
317+
| "max_chunk_size_tokens" : 1000,
318+
| "chunk_overlap_tokens" : 500
319+
| }
320+
|}""".stripMargin,
321+
Pretty
322+
)
323+
}
191324
}
192325

193326
private def testCodec[A](
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package io.cequence.openaiscala.domain
2+
3+
sealed trait ChunkingStrategy
4+
5+
object ChunkingStrategy {
6+
case object AutoChunkingStrategy extends ChunkingStrategy
7+
case class StaticChunkingStrategy private (
8+
maxChunkSizeTokens: Int,
9+
chunkOverlapTokens: Int
10+
) extends ChunkingStrategy
11+
12+
object StaticChunkingStrategy {
13+
def apply(
14+
maybeMaxChunkSizeTokens: Option[Int],
15+
maybeChunkOverlapTokens: Option[Int]
16+
): StaticChunkingStrategy = {
17+
val maxChunkSizeTokens = maybeMaxChunkSizeTokens.getOrElse(800)
18+
val chunkOverlapTokens = maybeChunkOverlapTokens.getOrElse(400)
19+
20+
if (maxChunkSizeTokens < 100 || maxChunkSizeTokens > 4096)
21+
throw new IllegalArgumentException("maxChunkSizeTokens must be between 100 and 4096")
22+
if (chunkOverlapTokens < 0 || chunkOverlapTokens > maxChunkSizeTokens / 2)
23+
throw new IllegalArgumentException(
24+
"chunkOverlapTokens must be between 0 and maxChunkSizeTokens/2"
25+
)
26+
new StaticChunkingStrategy(maxChunkSizeTokens, chunkOverlapTokens)
27+
}
28+
}
29+
}

openai-core/src/main/scala/io/cequence/openaiscala/domain/VectorStore.scala

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ case class VectorStore(
1717
)
1818

1919
case class FileCounts(
20-
in_progress: Int,
20+
inProgress: Int,
2121
completed: Int,
2222
cancelled: Int,
2323
failed: Int,
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
package io.cequence.openaiscala.domain
2+
3+
import io.cequence.wsclient.domain.EnumValue
4+
5+
sealed trait LastErrorCode extends EnumValue
6+
7+
object LastErrorCode {
8+
case object ServerError extends LastErrorCode
9+
case object RateLimitExceeded extends LastErrorCode
10+
}
11+
12+
case class LastError(
13+
code: LastErrorCode,
14+
message: String
15+
)
16+
17+
case class VectorStoreFile(
18+
id: String,
19+
`object`: String,
20+
usageBytes: Int,
21+
createdAt: Int,
22+
vectorStoreId: String,
23+
status: VectorStoreFileStatus,
24+
lastError: Option[LastError],
25+
chunkingStrategy: ChunkingStrategy
26+
)

0 commit comments

Comments
 (0)