22from abc import ABC , abstractmethod
33from typing import Generic , TypeVar
44
5+ from guidellm .backend .response import ResponseSummary
56from guidellm .request .request import GenerationRequest
67
78__all__ = ["RequestSession" ]
@@ -26,8 +27,8 @@ def push_response(self, response: ResponseT) -> None: ...
2627 def complete (self ) -> bool : ...
2728
2829
29- # FIXME: Dummy implementation
30- class GenerativeRequestSession (RequestSession [GenerationRequest , str ]):
30+ # FIXME: Bad implementation. Can only handle string requests
31+ class GenerativeRequestSession (RequestSession [GenerationRequest , ResponseSummary ]):
3132 def __init__ (self , prompts : list [GenerationRequest ]) -> None :
3233 if not prompts :
3334 raise ValueError ("Prompts cannot be empty" )
@@ -52,9 +53,13 @@ def get_request(self) -> GenerationRequest:
5253
5354 return base_request
5455
55- def push_response (self , response : str ) -> None :
56+ def push_response (self , response : ResponseSummary ) -> None :
5657 if len (self .responses ) < len (self .prompts ):
57- self .responses .append (response )
58+ if response .response_output_tokens is not None :
59+ self .prompts [len (self .responses )].constraints ["output_tokens" ] = (
60+ response .response_output_tokens
61+ )
62+ self .responses .append (response .value )
5863 else :
5964 raise ValueError ("Response list full" )
6065
0 commit comments