project_charles / response_state_manager.py
sohojoe's picture
refactor - move ResponseState to responce_state_manager
90a9891
raw
history blame
2.47 kB
from datetime import datetime
class ResponseStepObservations:
def __init__(self, episode, step):
self.timestamp = datetime.utcnow()
self.episode = episode
self.step = step
self.llm_preview = ''
self.llm_responses = []
self.tts_raw_chunk_ids = []
def __str__(self):
state = ', '.join(f'{k}={v}' for k, v in self.__dict__.items() if k not in {'episode', 'step', 'timestamp'})
return f'episode={self.episode}, step={self.step}, timestamp={self.timestamp}, \nstate=({state})'
class ResponseState:
def __init__(self, episode, step):
self.timestamp = datetime.utcnow()
self.episode = episode
self.step = step
self.current_responses = []
self.speech_chunks_per_response = []
self.is_speaking = False
def __str__(self):
state = ', '.join(f'{k}={v}' for k, v in self.__dict__.items() if k not in {'episode', 'step'})
return f'episode={self.episode}, step={self.step}, \nstate=({state})'
class ResponseStateManager:
def __init__(self):
self.episode = 0
self.step = 0
self.response_step_obs = None
self.response_state = None
self.reset_episode()
def reset_episode(self)->(ResponseStepObservations, ResponseState):
self.episode += 1
self.step = 0
self.response_state = ResponseState(self.episode, self.step)
self.response_step_obs = ResponseStepObservations(self.episode, self.step)
return self.response_step_obs, self.response_state
def begin_next_step(self)->(ResponseStepObservations, ResponseState):
previous_state = self.response_step_obs
self.step += 1
self.response_step_obs = ResponseStepObservations(self.episode, self.step)
return previous_state, self.response_state
def set_llm_preview(self, llm_preview):
self.response_step_obs.llm_preview = llm_preview
def add_llm_response_and_clear_llm_preview(self, llm_response):
self.response_state.current_responses.append(llm_response)
self.response_state.speech_chunks_per_response.append(0)
self.response_step_obs.llm_responses.append(llm_response)
self.response_step_obs.llm_preview = ''
def add_tts_raw_chunk_id(self, chunk_id, llm_sentence_id):
self.response_state.speech_chunks_per_response[llm_sentence_id] += 1
self.response_step_obs.tts_raw_chunk_ids.append(chunk_id)