Tracks and manages multimodal data for both sync and async processing.
| 449 | |
| 450 | |
| 451 | class MultimodalDataTracker: |
| 452 | """Tracks and manages multimodal data for both sync and async processing.""" |
| 453 | |
| 454 | def __init__( |
| 455 | self, |
| 456 | model_type: str, |
| 457 | multimodal_server_config: Optional[MultimodalServerConfig] = None): |
| 458 | self._model_type = model_type |
| 459 | self._data = defaultdict[str, list](list) |
| 460 | self._embeddings = defaultdict[str, list](list) |
| 461 | self._placeholder_counts = defaultdict[str, int](int) |
| 462 | self._multimodal_server_config = multimodal_server_config if multimodal_server_config is not None else MultimodalServerConfig( |
| 463 | ) |
| 464 | |
| 465 | async def retrieve_all_async( |
| 466 | self |
| 467 | ) -> tuple[Optional[Dict[str, List[Any]]], Optional[Dict[str, List[Any]]]]: |
| 468 | """Retrieve all collected multimodal data and embeddings.""" |
| 469 | |
| 470 | async def _retrieve( |
| 471 | data: Optional[dict[str, |
| 472 | list]]) -> Optional[Dict[str, List[Any]]]: |
| 473 | if not data: |
| 474 | return None |
| 475 | return { |
| 476 | modality: await asyncio.gather(*items) |
| 477 | for modality, items in data.items() if items |
| 478 | } |
| 479 | |
| 480 | return await _retrieve(self._data), await _retrieve(self._embeddings) |
| 481 | |
| 482 | def retrieve_all_sync( |
| 483 | self |
| 484 | ) -> tuple[Optional[Dict[str, List[Any]]], Optional[Dict[str, List[Any]]]]: |
| 485 | """Retrieve all collected multimodal data and embeddings.""" |
| 486 | |
| 487 | def _retrieve( |
| 488 | data: Optional[dict[str, |
| 489 | list]]) -> Optional[Dict[str, List[Any]]]: |
| 490 | if not data: |
| 491 | return None |
| 492 | return { |
| 493 | modality: items |
| 494 | for modality, items in data.items() if items |
| 495 | } |
| 496 | |
| 497 | return _retrieve(self._data), _retrieve(self._embeddings) |
| 498 | |
| 499 | def add_data(self, |
| 500 | media_type: str, |
| 501 | data: Union[Coroutine, Any], |
| 502 | *, |
| 503 | is_embedding: bool = False): |
| 504 | current_count = len(self._data[media_type]) + len( |
| 505 | self._embeddings[media_type]) + 1 |
| 506 | placeholder = retrieve_multimodal_placeholder(self._model_type, |
| 507 | media_type, current_count) |
| 508 | (self._embeddings |
no outgoing calls
no test coverage detected