(self, project: str, token_usage: int)
| 154 | return None |
| 155 | |
| 156 | def update_token_usage(self, project: str, token_usage: int): |
| 157 | with Session(self.engine) as session: |
| 158 | agent_state = session.query(AgentStateModel).filter(AgentStateModel.project == project).first() |
| 159 | if agent_state: |
| 160 | state_stack = json.loads(agent_state.state_stack_json) |
| 161 | state_stack[-1]["token_usage"] += token_usage |
| 162 | agent_state.state_stack_json = json.dumps(state_stack) |
| 163 | session.commit() |
| 164 | else: |
| 165 | state_stack = [self.new_state()] |
| 166 | state_stack[-1]["token_usage"] = token_usage |
| 167 | agent_state = AgentStateModel(project=project, state_stack_json=json.dumps(state_stack)) |
| 168 | session.add(agent_state) |
| 169 | session.commit() |
| 170 | |
| 171 | def get_latest_token_usage(self, project: str): |
| 172 | with Session(self.engine) as session: |
no test coverage detected