OpenAI Completions

from openai import OpenAI
from athina_logger.log_stream_inference.openai_completion_stream import LogOpenAiCompletionStreamInference
from athina_logger.api_key import AthinaApiKey
from athina_logger.exception.custom_exception import CustomException

client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
AthinaApiKey.set_api_key('ATHINA_API_KEY')

# Using python openai SDK
def stream_openai_completion_response(user_query, prompt):
	response = client.completions.create(
		model='text-davinci-003',
		prompt=prompt,
		stream=True,
	)
	logger = LogOpenAiCompletionStreamInference(
		prompt=prompt,
		language_model_id="text-davinci-003",
		prompt_slug="customer_query",
		context=context, # OPTIONAL
		response_time=response_time_ms, # OPTIONAL
		customer_id="stripe", # OPTIONAL
		customer_user_id="shiv@athina.ai", # OPTIONAL
		session_id="c45g-1234-s6g4-43d3", # OPTIONAL
		user_query=user_query, # OPTIONAL
		environment="production", # OPTIONAL
		external_reference_id="5e838eaf-7dd0-4b6f-a32c-26110dd54e58", # OPTIONAL; If passed, should be unique across all inference calls
		custom_attributes={
            "name": "John Doe"
        } # OPTIONAL;
	)

	# Here are 2 ways to log openai chat streams
	# OPTION 1
	try:
		for chunk in response:
			logger.collect_stream_inference_by_chunk(chunk)
		logger.log_stream_inference()
	except Exception as e:
		if isinstance(e, CustomException):
			print(e.status_code)
			print(e.message)
		else:
			print(e)


	# OPTION 2
	try:
		logger.collect_stream_inference(response)
		logger.log_stream_inference()
	except Exception as e:
		if isinstance(e, CustomException):
			print(e.status_code)
			print(e.message)
		else:
			print(e)

# Using OpenAI API call and SSE
def stream_openai_completion_response_with_openai_api(user_query, prompt):
	reqUrl = 'https://api.openai.com/v1/completions'
	reqHeaders = {
		'Accept': 'text/event-stream',
		'Authorization': 'Bearer ' + 'OPENAI_API_KEY'
	}
	reqBody = {
		"model": "text-davinci-003",
		"prompt": prompt,
		"stream": True,
	}
	request = requests.post(reqUrl, stream=True,
							headers=reqHeaders, json=reqBody)
	logger = LogOpenAiCompletionStreamInference(
		prompt=prompt,
		language_model_id="text-davinci-003",
		prompt_slug="customer_query",
		context=context, # OPTIONAL
		response_time=response_time_ms, # OPTIONAL
		customer_id="stripe", # OPTIONAL
		customer_user_id="shiv@athina.ai", # OPTIONAL
		session_id="c45g-1234-s6g4-43d3", # OPTIONAL
		user_query=user_query, # OPTIONAL
		environment="production", # OPTIONAL
		external_reference_id="5e838eaf-7dd0-4b6f-a32c-26110dd54e58", # OPTIONAL; If passed, should be unique across all inference calls
		custom_attributes={
            "name": "John Doe"
        } # OPTIONAL;
	)
	client = sseclient.SSEClient(request)
	try:
		for event in client.events():
			if event.data != '[DONE]':
				logger.collect_stream_inference_by_chunk(json.loads(event.data))
		logger.log_stream_inference()
	except Exception as e:
		if isinstance(e, CustomException):
			print(e.status_code)
			print(e.message)
		else:
			print(e)