| 24 | |
| 25 | |
| 26 | class CompletionsWrapper(Completions): |
| 27 | openpipe_reporting_client: OpenPipe |
| 28 | openpipe_completions_client: OriginalOpenAI |
| 29 | |
| 30 | def __init__( |
| 31 | self, |
| 32 | client: OriginalOpenAI, |
| 33 | openpipe_reporting_client: OpenPipe, |
| 34 | openpipe_completions_client: OriginalOpenAI, |
| 35 | ) -> None: |
| 36 | super().__init__(client) |
| 37 | self.openpipe_reporting_client = openpipe_reporting_client |
| 38 | self.openpipe_completions_client = openpipe_completions_client |
| 39 | |
| 40 | def create( |
| 41 | self, *args, **kwargs |
| 42 | ) -> Union[ChatCompletion, Stream[ChatCompletionChunk]]: |
| 43 | openpipe_options = kwargs.pop("openpipe", {}) |
| 44 | |
| 45 | requested_at = int(time.time() * 1000) |
| 46 | model = kwargs.get("model", "") |
| 47 | |
| 48 | if model.startswith("openpipe:"): |
| 49 | extra_headers = get_extra_headers(kwargs, openpipe_options) |
| 50 | |
| 51 | return self.openpipe_completions_client.chat.completions.create( |
| 52 | **kwargs, extra_headers=extra_headers |
| 53 | ) |
| 54 | |
| 55 | try: |
| 56 | chat_completion = super().create(*args, **kwargs) |
| 57 | |
| 58 | if isinstance(chat_completion, Stream): |
| 59 | |
| 60 | def _gen(): |
| 61 | assembled_completion = None |
| 62 | for chunk in chat_completion: |
| 63 | assembled_completion = merge_openai_chunks( |
| 64 | assembled_completion, chunk |
| 65 | ) |
| 66 | |
| 67 | yield chunk |
| 68 | |
| 69 | received_at = int(time.time() * 1000) |
| 70 | |
| 71 | report( |
| 72 | configured_client=self.openpipe_reporting_client, |
| 73 | openpipe_options=openpipe_options, |
| 74 | requested_at=requested_at, |
| 75 | received_at=received_at, |
| 76 | req_payload=kwargs, |
| 77 | resp_payload=get_chat_completion_json(assembled_completion), |
| 78 | status_code=200, |
| 79 | ) |
| 80 | |
| 81 | return _gen() |
| 82 | else: |
| 83 | received_at = int(time.time() * 1000) |