New LogMessage design, plus Response.json() method

This commit is contained in:
Simon Willison 2023-07-03 06:46:51 -07:00
parent 61dd8afc60
commit b1c51df3f1
2 changed files with 14 additions and 6 deletions

View file

@ -65,7 +65,6 @@ class ChatResponse(Response):
def __init__(self, prompt, model, stream, key):
super().__init__(prompt, model, stream)
self.key = key
self._response_json = None
self._prompt_json = None
def iter_prompt(self):
@ -103,10 +102,11 @@ class ChatResponse(Response):
model=self.prompt.model.model_id,
prompt=self.prompt.prompt,
system=self.prompt.system,
options=not_nulls(self.prompt.options),
options_json=not_nulls(self.prompt.options),
prompt_json=self._prompt_json,
response=self.text(),
response_json=self._response_json,
response_json=self.json(),
reply_to_id=None, # TODO
chat_id=None, # TODO
)

View file

@ -32,11 +32,14 @@ class LogMessage:
model: str # Actually the model.model_id string
prompt: str # Simplified string version of prompt
system: Optional[str] # Simplified string of system prompt
options: Dict[str, Any] # Any options e.g. temperature
options_json: Dict[str, Any] # Any options e.g. temperature
prompt_json: Optional[Dict[str, Any]] # Detailed JSON of prompt
response: str # Simplified string version of response
response_json: Dict[str, Any] # Detailed JSON of response
chat_id: Optional[int] # ID of chat, if this is part of one
response_json: Optional[Dict[str, Any]] # Detailed JSON of response
reply_to_id: Optional[int] # ID of message this is a reply to
chat_id: Optional[
int
] # ID of chat this is a part of (ID of first message in thread)
class Response(ABC):
@ -46,6 +49,7 @@ class Response(ABC):
self.stream = stream
self._chunks: List[str] = []
self._done = False
self._response_json = None
def reply(self, prompt, system=None, **options):
new_prompt = [self.prompt.prompt, self.text(), prompt]
@ -83,6 +87,10 @@ class Response(ABC):
self._force()
return "".join(self._chunks)
def json(self) -> Optional[Dict[str, Any]]:
self._force()
return self._response_json
def duration_ms(self) -> int:
self._force()
return int((self._end - self._start) * 1000)