Parse the completion response to a GeneratorOutput.
(
self,
completion: Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]],
)
| 195 | return client |
| 196 | |
| 197 | def parse_chat_completion( |
| 198 | self, |
| 199 | completion: Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]], |
| 200 | ) -> "GeneratorOutput": |
| 201 | """Parse the completion response to a GeneratorOutput.""" |
| 202 | try: |
| 203 | # If the completion is already a GeneratorOutput, return it directly (prevent recursion) |
| 204 | if isinstance(completion, GeneratorOutput): |
| 205 | return completion |
| 206 | |
| 207 | # Check if it's a ChatCompletion object (non-streaming response) |
| 208 | if hasattr(completion, 'choices') and hasattr(completion, 'usage'): |
| 209 | # ALWAYS extract the string content directly |
| 210 | try: |
| 211 | # Direct extraction of message content |
| 212 | if (hasattr(completion, 'choices') and |
| 213 | len(completion.choices) > 0 and |
| 214 | hasattr(completion.choices[0], 'message') and |
| 215 | hasattr(completion.choices[0].message, 'content')): |
| 216 | |
| 217 | content = completion.choices[0].message.content |
| 218 | if isinstance(content, str): |
| 219 | parsed_data = content |
| 220 | else: |
| 221 | parsed_data = str(content) |
| 222 | else: |
| 223 | # Fallback: convert entire completion to string |
| 224 | parsed_data = str(completion) |
| 225 | |
| 226 | except Exception as e: |
| 227 | # Ultimate fallback |
| 228 | parsed_data = str(completion) |
| 229 | |
| 230 | return GeneratorOutput( |
| 231 | data=parsed_data, |
| 232 | usage=CompletionUsage( |
| 233 | completion_tokens=completion.usage.completion_tokens, |
| 234 | prompt_tokens=completion.usage.prompt_tokens, |
| 235 | total_tokens=completion.usage.total_tokens, |
| 236 | ), |
| 237 | raw_response=str(completion), |
| 238 | ) |
| 239 | else: |
| 240 | # Handle streaming response - collect all content parts into a single string |
| 241 | content_parts = [] |
| 242 | usage_info = None |
| 243 | for chunk in completion: |
| 244 | if chunk.choices[0].delta.content: |
| 245 | content_parts.append(chunk.choices[0].delta.content) |
| 246 | # Try to get usage info from the last chunk |
| 247 | if hasattr(chunk, 'usage') and chunk.usage: |
| 248 | usage_info = chunk.usage |
| 249 | |
| 250 | # Join all content parts into a single string |
| 251 | full_content = ''.join(content_parts) |
| 252 | |
| 253 | # Create usage object |
| 254 | usage = None |