Skip to content

Commit f9ad052

Browse files
committed
fix: add close and aclose methods to GeneratorWrapper
Currently, close and aclose methods are missing on GeneratorWrapper, making it impossible to close generators from client code without accessing the .generator directly, and missing span uploads in the process. This commit adds the missing methods for this usecase.
1 parent 3e530af commit f9ad052

1 file changed

Lines changed: 22 additions & 0 deletions

File tree

langfuse/_client/observe.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -593,6 +593,17 @@ def __next__(self) -> Any:
593593

594594
raise
595595

596+
def close(self) -> None:
597+
output: Any = self.items
598+
599+
if self.transform_fn is not None:
600+
output = self.transform_fn(self.items)
601+
elif all(isinstance(item, str) for item in self.items):
602+
output = "".join(self.items)
603+
604+
self.span.update(output=output).end()
605+
self.generator.close()
606+
596607

597608
class _ContextPreservedAsyncGeneratorWrapper:
598609
"""Async generator wrapper that ensures each iteration runs in preserved context."""
@@ -659,3 +670,14 @@ async def __anext__(self) -> Any:
659670
).end()
660671

661672
raise
673+
674+
async def aclose(self) -> None:
675+
output: Any = self.items
676+
677+
if self.transform_fn is not None:
678+
output = self.transform_fn(self.items)
679+
elif all(isinstance(item, str) for item in self.items):
680+
output = "".join(self.items)
681+
682+
self.span.update(output=output).end()
683+
await self.generator.aclose()

0 commit comments

Comments
 (0)