Manages orchestrator conversation history and usage tracking.
| 21 | |
| 22 | |
| 23 | class Session: |
| 24 | """Manages orchestrator conversation history and usage tracking.""" |
| 25 | |
| 26 | def __init__(self, config: LoopConfig): |
| 27 | self.config = config |
| 28 | self.messages: list[dict[str, Any]] = [] |
| 29 | self.turn_count: int = 0 |
| 30 | self.total_tokens: int = 0 |
| 31 | self.prompt_tokens: int = 0 |
| 32 | self.completion_tokens: int = 0 |
| 33 | |
| 34 | def add_system_message(self, content: str) -> None: |
| 35 | """Set the system message (called once at start).""" |
| 36 | # Replace any existing system message |
| 37 | self.messages = [m for m in self.messages if m.get("role") != "system"] |
| 38 | self.messages.insert(0, {"role": "system", "content": content}) |
| 39 | |
| 40 | def add_user_message(self, content: str) -> None: |
| 41 | """Add user message to history.""" |
| 42 | self.messages.append({"role": "user", "content": content}) |
| 43 | |
| 44 | def add_assistant_message( |
| 45 | self, content: str, tool_calls: Optional[list] = None |
| 46 | ) -> None: |
| 47 | """Add assistant message, optionally with tool calls.""" |
| 48 | msg: dict[str, Any] = {"role": "assistant", "content": content} |
| 49 | if tool_calls: |
| 50 | msg["tool_calls"] = tool_calls |
| 51 | self.messages.append(msg) |
| 52 | |
| 53 | def add_tool_results(self, results: list[dict]) -> None: |
| 54 | """Add tool result messages to history.""" |
| 55 | for result in results: |
| 56 | self.messages.append(result) |
| 57 | |
| 58 | def track_usage(self, usage: dict[str, int]) -> None: |
| 59 | """Accumulate token usage from an LLM response.""" |
| 60 | self.total_tokens += usage.get("total_tokens", 0) |
| 61 | self.prompt_tokens += usage.get("prompt_tokens", 0) |
| 62 | self.completion_tokens += usage.get("completion_tokens", 0) |
| 63 | |
| 64 | def get_usage_summary(self) -> dict[str, Any]: |
| 65 | """Return a summary of token usage.""" |
| 66 | return { |
| 67 | "turns": self.turn_count, |
| 68 | "total_tokens": self.total_tokens, |
| 69 | "prompt_tokens": self.prompt_tokens, |
| 70 | "completion_tokens": self.completion_tokens, |
| 71 | } |