client

type AIBrixExtraBody ¶

type AIBrixExtraBody struct {
	JobID              string                       `json:"job_id,omitempty"`
	Runtime            *plannerapi.RuntimeRef       `json:"runtime,omitempty"`
	ResourceAllocation ResourceAllocation           `json:"resource_allocation,omitempty"`
	ModelTemplate      *plannerapi.ModelTemplateRef `json:"model_template,omitempty"`
	Model              string                       `json:"model,omitempty"`
}

AIBrixExtraBody is the AIBrix-specific extension the BatchClient serializes onto POST /v1/batches via the openai-go SDK's extra_body channel. Everything else on the submission rides on openai.BatchNewParams directly.

type BatchClient ¶

type BatchClient interface {
	CreateBatch(ctx context.Context, params openai.BatchNewParams, aibrix AIBrixExtraBody) (*openai.Batch, error)
	GetBatch(ctx context.Context, batchID string) (*openai.Batch, error)
	CancelBatch(ctx context.Context, batchID string) (*openai.Batch, error)
	ListBatches(ctx context.Context, req *ListBatchesRequest) (*ListBatchesResponse, error)
}

BatchClient is the planner -> MDS adapter for all batch operations. It is the single mocking seam for MDS in tests: a fake BatchClient covers both the worker's submit path and the planner's read overlay.

All methods return *openai.Batch directly (no planner-side wrapper): the wire shape is OpenAI-compatible and the planner has no overlay fields to add today. Re-introduce a thin wrapper here when planner state needs to ride alongside the MDS batch on read.

type DefaultResourceAllocation ¶

type DefaultResourceAllocation struct {
	ProvisionID string `json:"provision_id,omitempty"`
}

DefaultResourceAllocation is the shape used by backends that only need to carry the provision ID (kubernetes / aws / lambdaCloud today).

type ListBatchesRequest ¶

type ListBatchesRequest struct {
	// Limit caps the page size. Zero means "use the upstream default"
	// (typically 20).
	Limit int `json:"limit,omitempty"`
	// After is the cursor returned (implicitly, as the trailing batch
	// ID) by a prior page. Empty means "first page".
	After string `json:"after,omitempty"`
}

ListBatchesRequest is the planner -> MDS paginated read for the batch list endpoint. The cursor semantics match MDS (and OpenAI): pass the last batch ID from the previous page as After.

type ListBatchesResponse ¶

type ListBatchesResponse struct {
	Data    []*openai.Batch `json:"data"`
	HasMore bool            `json:"has_more"`
}

ListBatchesResponse is the normalized MDS list payload. The slice is in upstream order (newest-first by MDS convention); cursor advancement is the caller's responsibility - re-issue ListBatches with After = Data[len(Data)-1].ID until HasMore is false.

type OpenAIBatchClient ¶

type OpenAIBatchClient struct {
	// contains filtered or unexported fields
}

OpenAIBatchClient implements BatchClient against an OpenAI-compatible /v1/batches endpoint (the metadata service). All AIBrix-namespaced fields ride along under the SDK's extra_body channel via option.WithJSONSet.

func NewOpenAIBatchClient(metadataServiceURL string, injector error_injection.Injector) *OpenAIBatchClient

func (c *OpenAIBatchClient) CancelBatch(ctx context.Context, batchID string) (*openai.Batch, error)

func (c *OpenAIBatchClient) CreateBatch(ctx context.Context, params openai.BatchNewParams, aibrix AIBrixExtraBody) (*openai.Batch, error)

func (c *OpenAIBatchClient) GetBatch(ctx context.Context, batchID string) (*openai.Batch, error)

func (c *OpenAIBatchClient) ListBatches(ctx context.Context, req *ListBatchesRequest) (*ListBatchesResponse, error)

type ResourceAllocation ¶

type ResourceAllocation interface {
	// contains filtered or unexported methods
}

ResourceAllocation is the backend-agnostic planner allocation metadata carried on AIBrixExtraBody.ResourceAllocation. Each backend defines its own concrete type implementing this interface.

isResourceAllocation is an unexported, no-op marker method (empty body, never called at runtime) that exists purely to seal the interface: because it is unexported, only types declared in this package can satisfy ResourceAllocation, preventing arbitrary structs from accidentally satisfying it.