|
1 | 1 | import json |
2 | 2 | from pathlib import Path |
3 | 3 | import traceback |
4 | | -from typing import List, Sequence, Tuple, AsyncGenerator, Optional, Dict, Any |
| 4 | +from typing import List, Sequence, Tuple, AsyncGenerator, Optional, Dict, Any, Mapping |
5 | 5 | from datetime import datetime |
6 | 6 | from loguru import logger |
7 | 7 | import asyncio |
8 | 8 |
|
9 | 9 | from autogen_agentchat.agents import BaseChatAgent |
10 | 10 | from autogen_agentchat.base import Response |
| 11 | +from autogen_agentchat.state import BaseState |
11 | 12 | from autogen_agentchat.messages import ( |
12 | 13 | BaseChatMessage, |
13 | 14 | TextMessage, |
|
25 | 26 | UserMessage, |
26 | 27 | ) |
27 | 28 | from autogen_core.model_context import TokenLimitedChatCompletionContext |
28 | | -from pydantic import BaseModel |
| 29 | +from pydantic import BaseModel, Field |
29 | 30 | from typing_extensions import Self |
30 | 31 | from autogen_core.code_executor import CodeExecutor |
31 | 32 | from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor |
@@ -69,6 +70,27 @@ class FileSurferConfig(BaseModel): |
69 | 70 | use_local_executor: bool = False |
70 | 71 |
|
71 | 72 |
|
| 73 | +class FileSurferState(BaseState): |
| 74 | + """ |
| 75 | + State class for saving and loading the FileSurfer's state. |
| 76 | +
|
| 77 | + Attributes: |
| 78 | + chat_history (List[LLMMessage]): List of chat messages exchanged with the model. |
| 79 | + type (str): The type of the state. Default: FileSurferState. |
| 80 | + browser_path (str, optional): The current path of the file browser. |
| 81 | + browser_title (str, optional): The current title/filename of the file browser. |
| 82 | + viewport_current_page (int): The current page in the viewport. |
| 83 | + approved_files (List[str]): List of previously approved file paths. |
| 84 | + """ |
| 85 | + |
| 86 | + chat_history: List[LLMMessage] = [] |
| 87 | + type: str = Field(default="FileSurferState") |
| 88 | + browser_path: str | None = None |
| 89 | + browser_title: str | None = None |
| 90 | + viewport_current_page: int = 0 |
| 91 | + approved_files: List[str] = [] |
| 92 | + |
| 93 | + |
72 | 94 | class FileSurfer(BaseChatAgent, Component[FileSurferConfig]): |
73 | 95 | """ |
74 | 96 | An agent that can handle files using Markitdown and a code executor. |
@@ -608,6 +630,68 @@ def _get_browser_state(self) -> Tuple[str, str]: |
608 | 630 |
|
609 | 631 | return (header, self._browser.viewport) |
610 | 632 |
|
| 633 | + async def save_state(self) -> Mapping[str, Any]: |
| 634 | + """ |
| 635 | + Save the current state of the FileSurfer. |
| 636 | +
|
| 637 | + Returns: |
| 638 | + A dictionary containing the chat history and browser state |
| 639 | + """ |
| 640 | + if not self.did_lazy_init: |
| 641 | + state = FileSurferState( |
| 642 | + chat_history=self._chat_history, |
| 643 | + browser_path=None, # No browser state when not initialized |
| 644 | + browser_title=None, |
| 645 | + viewport_current_page=0, |
| 646 | + approved_files=list(self._approved_files), |
| 647 | + ) |
| 648 | + else: |
| 649 | + state = FileSurferState( |
| 650 | + chat_history=self._chat_history, |
| 651 | + browser_path=self._browser.path, |
| 652 | + browser_title=self._browser.page_title, |
| 653 | + viewport_current_page=self._browser.viewport_current_page, |
| 654 | + approved_files=list(self._approved_files), |
| 655 | + ) |
| 656 | + return state.model_dump() |
| 657 | + |
| 658 | + async def load_state(self, state: Mapping[str, Any]) -> None: |
| 659 | + """ |
| 660 | + Load a previously saved state. |
| 661 | +
|
| 662 | + Args: |
| 663 | + state: Dictionary containing the state to load |
| 664 | + """ |
| 665 | + # Validate and convert the state to a FileSurferState |
| 666 | + file_surfer_state = FileSurferState.model_validate(state) |
| 667 | + |
| 668 | + # Update the chat history |
| 669 | + self._chat_history = file_surfer_state.chat_history |
| 670 | + |
| 671 | + # Update approved files |
| 672 | + self._approved_files = set(file_surfer_state.approved_files) |
| 673 | + |
| 674 | + # Restore browser state if available and agent is initialized |
| 675 | + if self.did_lazy_init and file_surfer_state.browser_path is not None: |
| 676 | + # Restore browser path and page if possible |
| 677 | + try: |
| 678 | + await self._browser.open_path(file_surfer_state.browser_path) |
| 679 | + # Try to navigate to the correct page in viewport |
| 680 | + if file_surfer_state.viewport_current_page > 0: |
| 681 | + target_page = file_surfer_state.viewport_current_page |
| 682 | + current_page = self._browser.viewport_current_page |
| 683 | + if target_page != current_page: |
| 684 | + # Navigate to the correct page |
| 685 | + diff = target_page - current_page |
| 686 | + if diff > 0: |
| 687 | + for _ in range(diff): |
| 688 | + self._browser.page_down() |
| 689 | + else: |
| 690 | + for _ in range(-diff): |
| 691 | + self._browser.page_up() |
| 692 | + except Exception as e: |
| 693 | + logger.warning(f"Failed to restore browser state: {e}") |
| 694 | + |
611 | 695 | async def close(self) -> None: |
612 | 696 | """Close the FileSurfer agent.""" |
613 | 697 | logger.info("Closing FileSurfer...") |
|
0 commit comments