z-ai-sdk-python/tests/integration_tests/test_batches.py

import logging
import logging.config
import os

import zai
from zai import ZaiClient


def test_batch_input_file(test_file_path, logging_conf) -> None:
	logging.config.dictConfig(logging_conf)  # type: ignore
	client = ZaiClient()  # Fill in your own API Key

	try:
		batch_input_file = client.files.create(
			file=open(os.path.join(test_file_path, 'batchinput.jsonl'), 'rb'),
			purpose='batch',
		)

		print(batch_input_file)

	#   FileObject(id='20240514_ea19d21b-d256-4586-b0df-e80a45e3c286', bytes=490, created_at=1715673494, filename=None, object='file', purpose='batch', status=None, status_details=None, fileName='batchinput.jsonl')

	except zai.core._errors.APIRequestFailedError as err:
		print(err)
	except zai.core._errors.APIInternalError as err:
		print(err)
	except zai.core._errors.APIStatusError as err:
		print(err)


def test_batch_create(logging_conf) -> None:
	logging.config.dictConfig(logging_conf)  # type: ignore
	client = ZaiClient()  # Fill in your own API Key
	try:
		create = client.batches.create(
			input_file_id='20240514_ea19d21b-d256-4586-b0df-e80a45e3c286',
			endpoint='/v4/chat/completions',
			completion_window='24h',
			metadata={'description': 'job test'},
			auto_delete_input_file=True,
		)
		print(create)
		# Batch(id='batch_1790292763050508288', completion_window='24h', created_at=1715674031399, endpoint='/v4/chat/completions', input_file_id='20240514_ea19d21b-d256-4586-b0df-e80a45e3c286', object='batch', status='validating', cancelled_at=None, cancelling_at=None, completed_at=None, error_file_id=None, errors=None, expired_at=None, expires_at=None, failed_at=None, finalizing_at=None, in_progress_at=None, metadata={'description': 'job test'}, output_file_id=None, request_counts=BatchRequestCounts(completed=None, failed=None, total=None))

	except zai.core._errors.APIRequestFailedError as err:
		print(err)
	except zai.core._errors.APIInternalError as err:
		print(err)
	except zai.core._errors.APIStatusError as err:
		print(err)


def test_batch_retrieve(logging_conf) -> None:
	logging.config.dictConfig(logging_conf)  # type: ignore
	client = ZaiClient()  # Fill in your own API Key
	try:
		retrieve = client.batches.retrieve('batch_1790291013237211136')
		print(retrieve)

	#   Batch(id='batch_1790291013237211136', completion_window='24h', created_at=1715673614000, endpoint='/v4/chat/completions', input_file_id='20240514_ea19d21b-d256-4586-b0df-e80a45e3c286', object='batch', status='validating', cancelled_at=None, cancelling_at=None, completed_at=None, error_file_id='', errors=None, expired_at=None, expires_at=None, failed_at=None, finalizing_at=None, in_progress_at=None, metadata={'description': 'job test'}, output_file_id='', request_counts=BatchRequestCounts(completed=None, failed=None, total=None))
	except zai.core._errors.APIRequestFailedError as err:
		print(err)
	except zai.core._errors.APIInternalError as err:
		print(err)
	except zai.core._errors.APIStatusError as err:
		print(err)


def test_batch_cancel(logging_conf) -> None:
	logging.config.dictConfig(logging_conf)  # type: ignore
	client = ZaiClient()  # Fill in your own API Key
	try:
		cancel = client.batches.cancel('batch_1790291013237211136')
		print(cancel)
		# Batch(id='batch_1790291013237211136', completion_window='24h', created_at=1715673614000, endpoint='/v4/chat/completions', input_file_id='20240514_ea19d21b-d256-4586-b0df-e80a45e3c286', object='batch', status='cancelling', cancelled_at=None, cancelling_at=1715673698775, completed_at=None, error_file_id='', errors=None, expired_at=None, expires_at=None, failed_at=None, finalizing_at=None, in_progress_at=None, metadata={'description': 'job test'}, output_file_id='', request_counts=BatchRequestCounts(completed=None, failed=None, total=None))
	except zai.core._errors.APIRequestFailedError as err:
		print(err)
	except zai.core._errors.APIInternalError as err:
		print(err)
	except zai.core._errors.APIStatusError as err:
		print(err)


def test_batch_list(logging_conf) -> None:
	logging.config.dictConfig(logging_conf)  # type: ignore
	client = ZaiClient()  # Fill in your own API Key
	try:
		list = client.batches.list(limit=10)
		print(list)
		if list.has_more:
			print('_________get_next_page___________')
			batch = list.get_next_page()
			print(batch)
			print('_________iter_pages___________')
			for batch in list.iter_pages():
				print(batch)
	except zai.core._errors.APIRequestFailedError as err:
		print(err)
	except zai.core._errors.APIInternalError as err:
		print(err)
	except zai.core._errors.APIStatusError as err:
		print(err)


def test_batch_result(test_file_path, logging_conf) -> None:
	logging.config.dictConfig(logging_conf)  # type: ignore
	client = ZaiClient()  # Fill in your own API Key
	try:
		content = client.files.content('file-QDpVyDIhxj8mcFiduUydNqQN')
		with open(os.path.join(test_file_path, 'content_batchoutput.jsonl'), 'wb') as f:
			f.write(content.content)
		content.write_to_file(os.path.join(test_file_path, 'write_to_file_batchoutput.jsonl'))

		assert content.content == open(os.path.join(test_file_path, 'batchoutput.jsonl'), 'rb').read()
		assert content.content == open(os.path.join(test_file_path, 'write_to_file_batchoutput.jsonl'), 'rb').read()
	except zai.core._errors.APIRequestFailedError as err:
		print(err)
	except zai.core._errors.APIInternalError as err:
		print(err)
	except zai.core._errors.APIStatusError as err:
		print(err)