2
2
from abc import ABC , abstractmethod
3
3
from typing import Generic , TypeVar
4
4
5
+ from guidellm .backend .response import ResponseSummary
5
6
from guidellm .request .request import GenerationRequest
6
7
7
8
__all__ = ["RequestSession" ]
@@ -26,8 +27,8 @@ def push_response(self, response: ResponseT) -> None: ...
26
27
def complete (self ) -> bool : ...
27
28
28
29
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 ]):
31
32
def __init__ (self , prompts : list [GenerationRequest ]) -> None :
32
33
if not prompts :
33
34
raise ValueError ("Prompts cannot be empty" )
@@ -52,9 +53,13 @@ def get_request(self) -> GenerationRequest:
52
53
53
54
return base_request
54
55
55
- def push_response (self , response : str ) -> None :
56
+ def push_response (self , response : ResponseSummary ) -> None :
56
57
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 )
58
63
else :
59
64
raise ValueError ("Response list full" )
60
65
0 commit comments