|
50 | 50 | set_captured_response_headers, |
51 | 51 | ) |
52 | 52 | from uipath.llm_client.utils.sampling import ( |
| 53 | + disabled_fields_stripped, |
53 | 54 | disabled_params_from_model_details, |
54 | 55 | strip_disabled_kwargs, |
55 | 56 | ) |
@@ -422,7 +423,15 @@ def _generate( |
422 | 423 | ) |
423 | 424 | set_captured_response_headers({}) |
424 | 425 | try: |
425 | | - result = self._uipath_generate(messages, stop=stop, run_manager=run_manager, **kwargs) |
| 426 | + with disabled_fields_stripped( |
| 427 | + self, |
| 428 | + disabled_params=self.disabled_params, |
| 429 | + model_name=self.model_name, |
| 430 | + logger=self.logger, |
| 431 | + ): |
| 432 | + result = self._uipath_generate( |
| 433 | + messages, stop=stop, run_manager=run_manager, **kwargs |
| 434 | + ) |
426 | 435 | self._inject_gateway_headers(result.generations) |
427 | 436 | return result |
428 | 437 | finally: |
@@ -453,9 +462,15 @@ async def _agenerate( |
453 | 462 | ) |
454 | 463 | set_captured_response_headers({}) |
455 | 464 | try: |
456 | | - result = await self._uipath_agenerate( |
457 | | - messages, stop=stop, run_manager=run_manager, **kwargs |
458 | | - ) |
| 465 | + with disabled_fields_stripped( |
| 466 | + self, |
| 467 | + disabled_params=self.disabled_params, |
| 468 | + model_name=self.model_name, |
| 469 | + logger=self.logger, |
| 470 | + ): |
| 471 | + result = await self._uipath_agenerate( |
| 472 | + messages, stop=stop, run_manager=run_manager, **kwargs |
| 473 | + ) |
459 | 474 | self._inject_gateway_headers(result.generations) |
460 | 475 | return result |
461 | 476 | finally: |
@@ -486,14 +501,20 @@ def _stream( |
486 | 501 | ) |
487 | 502 | set_captured_response_headers({}) |
488 | 503 | try: |
489 | | - first = True |
490 | | - for chunk in self._uipath_stream( |
491 | | - messages, stop=stop, run_manager=run_manager, **kwargs |
| 504 | + with disabled_fields_stripped( |
| 505 | + self, |
| 506 | + disabled_params=self.disabled_params, |
| 507 | + model_name=self.model_name, |
| 508 | + logger=self.logger, |
492 | 509 | ): |
493 | | - if first: |
494 | | - self._inject_gateway_headers([chunk]) |
495 | | - first = False |
496 | | - yield chunk |
| 510 | + first = True |
| 511 | + for chunk in self._uipath_stream( |
| 512 | + messages, stop=stop, run_manager=run_manager, **kwargs |
| 513 | + ): |
| 514 | + if first: |
| 515 | + self._inject_gateway_headers([chunk]) |
| 516 | + first = False |
| 517 | + yield chunk |
497 | 518 | finally: |
498 | 519 | set_captured_response_headers({}) |
499 | 520 |
|
@@ -522,14 +543,20 @@ async def _astream( |
522 | 543 | ) |
523 | 544 | set_captured_response_headers({}) |
524 | 545 | try: |
525 | | - first = True |
526 | | - async for chunk in self._uipath_astream( |
527 | | - messages, stop=stop, run_manager=run_manager, **kwargs |
| 546 | + with disabled_fields_stripped( |
| 547 | + self, |
| 548 | + disabled_params=self.disabled_params, |
| 549 | + model_name=self.model_name, |
| 550 | + logger=self.logger, |
528 | 551 | ): |
529 | | - if first: |
530 | | - self._inject_gateway_headers([chunk]) |
531 | | - first = False |
532 | | - yield chunk |
| 552 | + first = True |
| 553 | + async for chunk in self._uipath_astream( |
| 554 | + messages, stop=stop, run_manager=run_manager, **kwargs |
| 555 | + ): |
| 556 | + if first: |
| 557 | + self._inject_gateway_headers([chunk]) |
| 558 | + first = False |
| 559 | + yield chunk |
533 | 560 | finally: |
534 | 561 | set_captured_response_headers({}) |
535 | 562 |
|
|
0 commit comments