| 20 | pass |
| 21 | |
| 22 | class Session: |
| 23 | def __init__(self, model_inference, history=None) -> None: |
| 24 | self.history: list[dict] = history or [] |
| 25 | self.exception_raised = False |
| 26 | self.model_inference = self.wrap_inference(model_inference) |
| 27 | |
| 28 | def inject(self, message: dict) -> None: |
| 29 | assert isinstance(message, dict) |
| 30 | assert "role" in message and "content" in message |
| 31 | assert isinstance(message["role"], str) |
| 32 | assert isinstance(message["content"], str) |
| 33 | assert message["role"] in ["user", "agent"] |
| 34 | self.history.append(message) |
| 35 | |
| 36 | def action(self, extend_messages: List[dict] = None) -> str: |
| 37 | extend = [] |
| 38 | if extend_messages: |
| 39 | if isinstance(extend_messages, list): |
| 40 | extend.extend(extend_messages) |
| 41 | elif isinstance(extend_messages, dict): |
| 42 | extend.append(extend_messages) |
| 43 | else: |
| 44 | raise Exception("Invalid extend_messages") |
| 45 | result = self.model_inference(self.history + extend) |
| 46 | self.history.extend(extend) |
| 47 | self.history.append({"role": "agent", "content": result}) |
| 48 | return result |
| 49 | |
| 50 | def _calc_segments(self, msg: str): |
| 51 | segments = 0 |
| 52 | current_segment = "" |
| 53 | inside_word = False |
| 54 | |
| 55 | for char in msg: |
| 56 | if char.isalpha(): |
| 57 | current_segment += char |
| 58 | if not inside_word: |
| 59 | inside_word = True |
| 60 | if len(current_segment) >= 7: |
| 61 | segments += 1 |
| 62 | current_segment = "" |
| 63 | inside_word = False |
| 64 | else: |
| 65 | if inside_word: |
| 66 | segments += 1 |
| 67 | current_segment = "" |
| 68 | inside_word = False |
| 69 | if char not in [" ", "\n"]: |
| 70 | segments += 1 |
| 71 | |
| 72 | if len(current_segment) > 0: |
| 73 | segments += 1 |
| 74 | |
| 75 | return segments |
| 76 | |
| 77 | def wrap_inference(self, inference_function: Callable[[List[dict]], str]) -> Callable[[List[dict]], str]: |
| 78 | def _func(history: List[dict]) -> str: |
| 79 | if self.exception_raised: |