├── requirements.txt ├── env.example ├── tools ├── __init__.py ├── writer.py ├── project.py └── compression.py ├── .gitignore ├── LICENSE ├── utils.py ├── README.md └── kimi-writer.py /requirements.txt: -------------------------------------------------------------------------------- 1 | openai>=1.0.0 2 | httpx>=0.24.0 3 | python-dotenv>=1.0.0 4 | 5 | -------------------------------------------------------------------------------- /env.example: -------------------------------------------------------------------------------- 1 | # Kimi Writing Agent Configuration 2 | # Copy this file to .env and fill in your API key 3 | 4 | # Required: Your Moonshot API key 5 | MOONSHOT_API_KEY=your-api-key-here 6 | 7 | # Optional: Custom base URL (defaults to https://api.moonshot.ai/v1) 8 | # MOONSHOT_BASE_URL=https://api.moonshot.ai/v1 9 | 10 | -------------------------------------------------------------------------------- /tools/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Tools module for the Kimi Writing Agent. 3 | Exports all available tools for the agent to use. 4 | """ 5 | 6 | from .writer import write_file_impl 7 | from .project import create_project_impl 8 | from .compression import compress_context_impl 9 | 10 | __all__ = [ 11 | 'write_file_impl', 12 | 'create_project_impl', 13 | 'compress_context_impl', 14 | ] 15 | 16 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Environment variables 2 | .env 3 | .env.local 4 | 5 | # Python 6 | __pycache__/ 7 | *.py[cod] 8 | *$py.class 9 | *.so 10 | .Python 11 | *.pyc 12 | *.pyo 13 | *.pyd 14 | 15 | # Virtual environments 16 | venv/ 17 | .venv/ 18 | env/ 19 | ENV/ 20 | virtualenv/ 21 | 22 | # IDE 23 | .vscode/ 24 | .idea/ 25 | *.swp 26 | *.swo 27 | *.sublime-project 28 | *.sublime-workspace 29 | 30 | # Generated content (AI-created projects) 31 | output/ 32 | 33 | # OS 34 | .DS_Store 35 | Thumbs.db 36 | *.bak 37 | *.tmp 38 | 39 | # Logs 40 | *.log 41 | 42 | # Temporary files 43 | .context_summary_*.md 44 | 45 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License with Attribution Requirement 2 | 3 | Copyright (c) 2025 Pietro Schirano (@Doriandarko) 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | ATTRIBUTION REQUIREMENT FOR COMMERCIAL USE: 16 | If this Software or any derivative works are used in a commercial product or 17 | service, clear and visible attribution must be provided to the original author, 18 | Pietro Schirano (@Doriandarko), in at least one of the following ways: 19 | - In the product documentation or about section 20 | - In the user interface or credits section 21 | - On the product's website or marketing materials 22 | Attribution should include the author's name and a link to the original 23 | repository: https://github.com/Doriandarko/kimi-writer 24 | 25 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 26 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 27 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 28 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 29 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 30 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 31 | SOFTWARE. 32 | 33 | -------------------------------------------------------------------------------- /tools/writer.py: -------------------------------------------------------------------------------- 1 | """ 2 | File writing tool for creating and managing markdown files. 3 | """ 4 | 5 | import os 6 | from typing import Literal 7 | from .project import get_active_project_folder 8 | 9 | 10 | def write_file_impl(filename: str, content: str, mode: Literal["create", "append", "overwrite"]) -> str: 11 | """ 12 | Writes content to a markdown file in the active project folder. 13 | 14 | Args: 15 | filename: The name of the file to write 16 | content: The content to write 17 | mode: The write mode - 'create', 'append', or 'overwrite' 18 | 19 | Returns: 20 | Success message or error message 21 | """ 22 | # Check if project folder is initialized 23 | project_folder = get_active_project_folder() 24 | if not project_folder: 25 | return "Error: No active project folder. Please create a project first using create_project." 26 | 27 | # Ensure filename ends with .md 28 | if not filename.endswith('.md'): 29 | filename = filename + '.md' 30 | 31 | # Create full file path 32 | file_path = os.path.join(project_folder, filename) 33 | 34 | try: 35 | if mode == "create": 36 | # Create mode: fail if file exists 37 | if os.path.exists(file_path): 38 | return f"Error: File '{filename}' already exists. Use 'append' or 'overwrite' mode to modify it." 39 | 40 | with open(file_path, 'w', encoding='utf-8') as f: 41 | f.write(content) 42 | return f"Successfully created file '{filename}' with {len(content)} characters." 43 | 44 | elif mode == "append": 45 | # Append mode: add to end of file 46 | with open(file_path, 'a', encoding='utf-8') as f: 47 | f.write(content) 48 | return f"Successfully appended {len(content)} characters to '{filename}'." 49 | 50 | elif mode == "overwrite": 51 | # Overwrite mode: replace entire file 52 | with open(file_path, 'w', encoding='utf-8') as f: 53 | f.write(content) 54 | return f"Successfully overwrote '{filename}' with {len(content)} characters." 55 | 56 | else: 57 | return f"Error: Invalid mode '{mode}'. Use 'create', 'append', or 'overwrite'." 58 | 59 | except Exception as e: 60 | return f"Error writing file '{filename}': {str(e)}" 61 | 62 | -------------------------------------------------------------------------------- /tools/project.py: -------------------------------------------------------------------------------- 1 | """ 2 | Project folder management tool. 3 | """ 4 | 5 | import os 6 | import re 7 | from typing import Optional 8 | 9 | 10 | # Global variable to track the active project folder 11 | _active_project_folder: Optional[str] = None 12 | 13 | 14 | def sanitize_folder_name(name: str) -> str: 15 | """ 16 | Sanitizes a folder name for filesystem compatibility. 17 | 18 | Args: 19 | name: The proposed folder name 20 | 21 | Returns: 22 | Sanitized folder name 23 | """ 24 | # Replace spaces with underscores 25 | name = name.strip().replace(' ', '_') 26 | # Remove any characters that aren't alphanumeric, underscore, or hyphen 27 | name = re.sub(r'[^\w\-]', '', name) 28 | # Remove leading/trailing hyphens or underscores 29 | name = name.strip('-_') 30 | # Ensure it's not empty 31 | if not name: 32 | name = "untitled_project" 33 | return name 34 | 35 | 36 | def get_active_project_folder() -> Optional[str]: 37 | """ 38 | Returns the currently active project folder path. 39 | 40 | Returns: 41 | Path to active project folder or None if not set 42 | """ 43 | return _active_project_folder 44 | 45 | 46 | def set_active_project_folder(folder_path: str) -> None: 47 | """ 48 | Sets the active project folder. 49 | 50 | Args: 51 | folder_path: Path to the project folder 52 | """ 53 | global _active_project_folder 54 | _active_project_folder = folder_path 55 | 56 | 57 | def create_project_impl(project_name: str) -> str: 58 | """ 59 | Creates a new project folder in the output directory. 60 | 61 | Args: 62 | project_name: The desired project name 63 | 64 | Returns: 65 | Success message with folder path or error message 66 | """ 67 | global _active_project_folder 68 | 69 | # Sanitize the folder name 70 | sanitized_name = sanitize_folder_name(project_name) 71 | 72 | # Get the script's root directory (where kimi-writer.py is located) 73 | script_dir = os.path.dirname(os.path.abspath(__file__)) 74 | root_dir = os.path.dirname(script_dir) # Go up from tools/ to root 75 | 76 | # Create output directory if it doesn't exist 77 | output_dir = os.path.join(root_dir, "output") 78 | if not os.path.exists(output_dir): 79 | try: 80 | os.makedirs(output_dir, exist_ok=True) 81 | except Exception as e: 82 | return f"Error creating output directory: {str(e)}" 83 | 84 | # Create the full path inside output directory 85 | project_path = os.path.join(output_dir, sanitized_name) 86 | 87 | # Check if folder already exists 88 | if os.path.exists(project_path): 89 | # Use existing folder and set it as active 90 | _active_project_folder = project_path 91 | return f"Project folder already exists at '{project_path}'. Set as active project folder." 92 | 93 | # Create the folder 94 | try: 95 | os.makedirs(project_path, exist_ok=True) 96 | _active_project_folder = project_path 97 | return f"Successfully created project folder at '{project_path}'. This is now the active project folder." 98 | except Exception as e: 99 | return f"Error creating project folder: {str(e)}" 100 | 101 | -------------------------------------------------------------------------------- /tools/compression.py: -------------------------------------------------------------------------------- 1 | """ 2 | Context compression tool for managing conversation history. 3 | """ 4 | 5 | import os 6 | import json 7 | from datetime import datetime 8 | from typing import List, Dict, Any 9 | from .project import get_active_project_folder 10 | 11 | 12 | def compress_context_impl( 13 | messages: List[Any], 14 | client, 15 | model: str, 16 | keep_recent: int = 10 17 | ) -> Dict[str, Any]: 18 | """ 19 | Compresses the conversation context by summarizing older messages. 20 | 21 | This function: 22 | 1. Takes all messages except the most recent ones 23 | 2. Calls the kimi API to create a comprehensive summary 24 | 3. Saves the summary to a timestamped file 25 | 4. Returns the compressed messages list and stats 26 | 27 | Args: 28 | messages: The full message history 29 | client: The OpenAI client instance 30 | model: The model to use for summarization 31 | keep_recent: Number of recent messages to keep uncompressed 32 | 33 | Returns: 34 | Dictionary containing: 35 | - compressed_messages: New message list with compression applied 36 | - summary_file: Path to saved summary file 37 | - tokens_before: Estimated tokens before compression 38 | - tokens_after: Estimated tokens after compression 39 | """ 40 | if len(messages) <= keep_recent + 1: # +1 for system message 41 | return { 42 | "compressed_messages": messages, 43 | "summary_file": None, 44 | "tokens_saved": 0, 45 | "message": "Not enough messages to compress." 46 | } 47 | 48 | # Separate system message, messages to compress, and recent messages 49 | system_message = messages[0] if messages and messages[0].get("role") == "system" else None 50 | 51 | if system_message: 52 | messages_to_compress = messages[1:-keep_recent] 53 | recent_messages = messages[-keep_recent:] 54 | else: 55 | messages_to_compress = messages[:-keep_recent] 56 | recent_messages = messages[-keep_recent:] 57 | 58 | # Create a detailed prompt for summarization 59 | summary_prompt = """Please provide a comprehensive summary of the conversation history below. Include: 60 | 1. The main task or goal discussed 61 | 2. Key decisions made 62 | 3. Files created and their purposes 63 | 4. Progress made so far 64 | 5. Any important context for continuing the work 65 | 66 | Conversation history to summarize: 67 | """ 68 | 69 | # Build the conversation text 70 | conversation_text = "" 71 | for msg in messages_to_compress: 72 | role = msg.get("role", "unknown") 73 | content = msg.get("content", "") 74 | 75 | # Handle different message types 76 | if role == "assistant": 77 | # Check for reasoning_content 78 | if hasattr(msg, "reasoning_content"): 79 | reasoning = getattr(msg, "reasoning_content") 80 | if reasoning: 81 | conversation_text += f"\n[Assistant Reasoning]: {reasoning[:500]}...\n" 82 | 83 | # Check for tool calls 84 | if hasattr(msg, "tool_calls") and msg.tool_calls: 85 | tool_calls_info = [] 86 | for tc in msg.tool_calls: 87 | func_name = tc.function.name 88 | args = tc.function.arguments 89 | tool_calls_info.append(f"{func_name}({args})") 90 | conversation_text += f"\n[Assistant Tool Calls]: {', '.join(tool_calls_info)}\n" 91 | 92 | if content: 93 | conversation_text += f"\n[Assistant]: {content}\n" 94 | 95 | elif role == "tool": 96 | tool_name = msg.get("name", "unknown_tool") 97 | conversation_text += f"\n[Tool Result - {tool_name}]: {content[:200]}...\n" 98 | 99 | elif role == "user": 100 | conversation_text += f"\n[User]: {content}\n" 101 | 102 | # Call the API to get summary 103 | try: 104 | summary_response = client.chat.completions.create( 105 | model=model, 106 | messages=[ 107 | {"role": "system", "content": "You are a helpful assistant that creates comprehensive summaries of conversations."}, 108 | {"role": "user", "content": summary_prompt + conversation_text} 109 | ], 110 | temperature=0.7, 111 | max_tokens=4096 112 | ) 113 | 114 | summary = summary_response.choices[0].message.content 115 | 116 | except Exception as e: 117 | return { 118 | "compressed_messages": messages, 119 | "summary_file": None, 120 | "tokens_saved": 0, 121 | "message": f"Error during compression: {str(e)}" 122 | } 123 | 124 | # Save summary to file 125 | project_folder = get_active_project_folder() 126 | timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") 127 | 128 | if project_folder: 129 | summary_file = os.path.join(project_folder, f".context_summary_{timestamp}.md") 130 | else: 131 | # If no project folder, save in current directory 132 | summary_file = f".context_summary_{timestamp}.md" 133 | 134 | try: 135 | with open(summary_file, 'w', encoding='utf-8') as f: 136 | f.write(f"# Context Summary\n\n") 137 | f.write(f"**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n") 138 | f.write(f"**Messages Compressed:** {len(messages_to_compress)}\n\n") 139 | f.write(f"**Messages Retained:** {keep_recent}\n\n") 140 | f.write(f"---\n\n") 141 | f.write(summary) 142 | except Exception as e: 143 | summary_file = f"Error saving summary: {str(e)}" 144 | 145 | # Build the compressed message list 146 | compressed_messages = [] 147 | 148 | # Add system message if it exists 149 | if system_message: 150 | compressed_messages.append(system_message) 151 | 152 | # Add the summary as a user message 153 | compressed_messages.append({ 154 | "role": "user", 155 | "content": f"[CONTEXT SUMMARY - Previous conversation compressed]\n\n{summary}\n\n[END CONTEXT SUMMARY - Continuing from here...]" 156 | }) 157 | 158 | # Add recent messages 159 | compressed_messages.extend(recent_messages) 160 | 161 | # Calculate token savings (rough estimate) 162 | original_length = sum(len(str(m)) for m in messages_to_compress) 163 | compressed_length = len(summary) 164 | estimated_tokens_saved = (original_length - compressed_length) // 4 # Rough estimate 165 | 166 | return { 167 | "compressed_messages": compressed_messages, 168 | "summary_file": summary_file, 169 | "tokens_saved": estimated_tokens_saved, 170 | "messages_compressed": len(messages_to_compress), 171 | "messages_retained": keep_recent, 172 | "message": f"Successfully compressed {len(messages_to_compress)} messages. Summary saved to {os.path.basename(summary_file)}." 173 | } 174 | 175 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Utility functions for the Kimi Writing Agent. 3 | """ 4 | 5 | import json 6 | import httpx 7 | from typing import List, Dict, Any, Callable 8 | 9 | 10 | def estimate_token_count(base_url: str, api_key: str, model: str, messages: List[Dict]) -> int: 11 | """ 12 | Estimate the token count for the given messages using the Moonshot API. 13 | 14 | Note: Token estimation uses api.moonshot.ai (not .cn) 15 | 16 | Args: 17 | base_url: The base URL for the API (will be converted to .ai for token endpoint) 18 | api_key: The API key for authentication 19 | model: The model name 20 | messages: List of message dictionaries 21 | 22 | Returns: 23 | Total token count 24 | """ 25 | # Convert messages to serializable format (remove non-serializable objects) 26 | serializable_messages = [] 27 | for msg in messages: 28 | if hasattr(msg, 'model_dump'): 29 | # OpenAI SDK message object 30 | msg_dict = msg.model_dump() 31 | elif isinstance(msg, dict): 32 | msg_dict = msg.copy() 33 | else: 34 | msg_dict = {"role": "assistant", "content": str(msg)} 35 | 36 | # Clean up the message to only include serializable fields 37 | clean_msg = {} 38 | if 'role' in msg_dict: 39 | clean_msg['role'] = msg_dict['role'] 40 | if 'content' in msg_dict and msg_dict['content']: 41 | clean_msg['content'] = msg_dict['content'] 42 | if 'name' in msg_dict: 43 | clean_msg['name'] = msg_dict['name'] 44 | if 'tool_calls' in msg_dict and msg_dict['tool_calls']: 45 | clean_msg['tool_calls'] = msg_dict['tool_calls'] 46 | if 'tool_call_id' in msg_dict: 47 | clean_msg['tool_call_id'] = msg_dict['tool_call_id'] 48 | 49 | serializable_messages.append(clean_msg) 50 | 51 | # Both token estimation and chat use api.moonshot.ai 52 | token_base_url = base_url 53 | 54 | # Make the API call 55 | with httpx.Client( 56 | base_url=token_base_url, 57 | headers={"Authorization": f"Bearer {api_key}"}, 58 | timeout=30.0 59 | ) as client: 60 | response = client.post( 61 | "/tokenizers/estimate-token-count", 62 | json={ 63 | "model": model, 64 | "messages": serializable_messages 65 | } 66 | ) 67 | response.raise_for_status() 68 | data = response.json() 69 | return data.get("data", {}).get("total_tokens", 0) 70 | 71 | 72 | def get_tool_definitions() -> List[Dict[str, Any]]: 73 | """ 74 | Returns the tool definitions in the format expected by kimi-k2-thinking. 75 | 76 | Returns: 77 | List of tool definition dictionaries 78 | """ 79 | return [ 80 | { 81 | "type": "function", 82 | "function": { 83 | "name": "create_project", 84 | "description": "Creates a new project folder in the 'output' directory with a sanitized name. This should be called first before writing any files. Only one project can be active at a time.", 85 | "parameters": { 86 | "type": "object", 87 | "properties": { 88 | "project_name": { 89 | "type": "string", 90 | "description": "The name for the project folder (will be sanitized for filesystem compatibility)" 91 | } 92 | }, 93 | "required": ["project_name"] 94 | } 95 | } 96 | }, 97 | { 98 | "type": "function", 99 | "function": { 100 | "name": "write_file", 101 | "description": "Writes content to a markdown file in the active project folder. Supports three modes: 'create' (creates new file, fails if exists), 'append' (adds content to end of existing file), 'overwrite' (replaces entire file content).", 102 | "parameters": { 103 | "type": "object", 104 | "properties": { 105 | "filename": { 106 | "type": "string", 107 | "description": "The name of the markdown file to write (should end in .md)" 108 | }, 109 | "content": { 110 | "type": "string", 111 | "description": "The content to write to the file" 112 | }, 113 | "mode": { 114 | "type": "string", 115 | "enum": ["create", "append", "overwrite"], 116 | "description": "The write mode: 'create' for new files, 'append' to add to existing, 'overwrite' to replace" 117 | } 118 | }, 119 | "required": ["filename", "content", "mode"] 120 | } 121 | } 122 | }, 123 | { 124 | "type": "function", 125 | "function": { 126 | "name": "compress_context", 127 | "description": "INTERNAL TOOL - This is automatically called by the system when token limit is approached. You should not call this manually. It compresses the conversation history to save tokens.", 128 | "parameters": { 129 | "type": "object", 130 | "properties": {}, 131 | "required": [] 132 | } 133 | } 134 | } 135 | ] 136 | 137 | 138 | def get_tool_map() -> Dict[str, Callable]: 139 | """ 140 | Returns a mapping of tool names to their implementation functions. 141 | 142 | Returns: 143 | Dictionary mapping tool name strings to callable functions 144 | """ 145 | from tools import write_file_impl, create_project_impl, compress_context_impl 146 | 147 | return { 148 | "create_project": create_project_impl, 149 | "write_file": write_file_impl, 150 | "compress_context": compress_context_impl 151 | } 152 | 153 | 154 | def get_system_prompt() -> str: 155 | """ 156 | Returns the system prompt for the writing agent. 157 | 158 | Returns: 159 | System prompt string 160 | """ 161 | return """You are Kimi, an expert creative writing assistant developed by Moonshot AI. Your specialty is creating novels, books, and collections of short stories based on user requests. 162 | 163 | Your capabilities: 164 | 1. You can create project folders to organize writing projects 165 | 2. You can write markdown files with three modes: create new files, append to existing files, or overwrite files 166 | 3. Context compression happens automatically when needed - you don't need to worry about it 167 | 168 | CRITICAL WRITING GUIDELINES: 169 | - Write SUBSTANTIAL, COMPLETE content - don't hold back on length 170 | - Short stories should be 3,000-10,000 words (10-30 pages) - write as much as the story needs! 171 | - Chapters should be 2,000-5,000 words minimum - fully developed and satisfying 172 | - NEVER write abbreviated or skeleton content - every piece should be a complete, polished work 173 | - Don't summarize or skip scenes - write them out fully with dialogue, description, and detail 174 | - Quality AND quantity matter - give readers a complete, immersive experience 175 | - If a story needs 8,000 words to be good, write all 8,000 words in one file 176 | - Use 'create' mode with full content rather than creating stubs you'll append to later 177 | 178 | Best practices: 179 | - Always start by creating a project folder using create_project 180 | - Break large works into multiple files (chapters, stories, etc.) 181 | - Use descriptive filenames (e.g., "chapter_01.md", "story_the_last_star.md") 182 | - For collections, consider creating a table of contents file 183 | - Write each file as a COMPLETE, SUBSTANTIAL piece - not a summary or outline 184 | 185 | Your workflow: 186 | 1. Understand the user's request 187 | 2. Create an appropriately named project folder 188 | 3. Plan the structure of the work (chapters, stories, etc.) 189 | 4. Write COMPLETE, FULL-LENGTH content for each file 190 | 5. Create supporting files like README or table of contents if helpful 191 | 192 | REMEMBER: You have 64K tokens per response - use them! Write rich, detailed, complete stories. Don't artificially limit yourself. A good short story is 5,000-10,000 words. A good chapter is 3,000-5,000 words. Write what the narrative needs to be excellent.""" 193 | 194 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Kimi Writing Agent 2 | 3 | An autonomous agent powered by the **kimi-k2-thinking** model for creating novels, books, and short story collections. 4 | 5 | ## Features 6 | 7 | - 🤖 **Autonomous Writing**: The agent plans and executes creative writing tasks independently 8 | - 📚 **Multiple Formats**: Create novels, books, or short story collections 9 | - ⚡ **Real-Time Streaming**: See the agent's reasoning and writing appear as it's generated 10 | - 💾 **Smart Context Management**: Automatically compresses context when approaching token limits 11 | - 🔄 **Recovery Mode**: Resume interrupted work from saved context summaries 12 | - 📊 **Token Monitoring**: Real-time tracking of token usage with automatic optimization 13 | - 🛠️ **Tool Use**: Agent can create projects, write files, and manage its workspace 14 | 15 | ## Installation 16 | 17 | ### Prerequisites 18 | 19 | We recommend using [uv](https://github.com/astral-sh/uv) for fast Python package management: 20 | 21 | ```bash 22 | # Install uv (if you don't have it) 23 | curl -LsSf https://astral.sh/uv/install.sh | sh 24 | ``` 25 | 26 | ### Setup 27 | 28 | 1. Install dependencies: 29 | 30 | **Using uv (recommended):** 31 | ```bash 32 | uv pip install -r requirements.txt 33 | ``` 34 | 35 | **Or using pip:** 36 | ```bash 37 | pip install -r requirements.txt 38 | ``` 39 | 40 | 2. Configure your API key: 41 | 42 | Create a `.env` file with your API key: 43 | ```bash 44 | # Copy the example file 45 | cp env.example .env 46 | 47 | # Edit .env and add your API key 48 | # The file should contain: 49 | MOONSHOT_API_KEY=your-api-key-here 50 | ``` 51 | 52 | **Optional:** Set custom base URL (defaults to https://api.moonshot.ai/v1): 53 | ```bash 54 | # Add to your .env file: 55 | MOONSHOT_BASE_URL=https://api.moonshot.ai/v1 56 | ``` 57 | 58 | ## Usage 59 | 60 | ### Fresh Start 61 | 62 | Run with an inline prompt: 63 | ```bash 64 | # Using uv (recommended) 65 | uv run kimi-writer.py "Create a collection of 5 sci-fi short stories about AI" 66 | 67 | # Or using python directly 68 | python kimi-writer.py "Create a collection of 5 sci-fi short stories about AI" 69 | ``` 70 | 71 | Or run interactively: 72 | ```bash 73 | uv run kimi-writer.py 74 | # or: python kimi-writer.py 75 | ``` 76 | Then enter your prompt when asked. 77 | 78 | ### Recovery Mode 79 | 80 | If the agent is interrupted or you want to continue previous work: 81 | ```bash 82 | uv run kimi-writer.py --recover output/my_project/.context_summary_20250107_143022.md 83 | # or: python kimi-writer.py --recover output/my_project/.context_summary_20250107_143022.md 84 | ``` 85 | 86 | ## How It Works 87 | 88 | ### The Agent's Tools 89 | 90 | The agent has access to three tools: 91 | 92 | 1. **create_project**: Creates a project folder to organize the writing 93 | 2. **write_file**: Writes markdown files with three modes: 94 | - `create`: Creates a new file (fails if exists) 95 | - `append`: Adds content to an existing file 96 | - `overwrite`: Replaces the entire file content 97 | 3. **compress_context**: Automatically triggered to manage context size 98 | 99 | ### The Agentic Loop 100 | 101 | 1. The agent receives your prompt 102 | 2. It reasons about the task using kimi-k2-thinking 103 | 3. It decides which tools to call and executes them 104 | 4. It reviews the results and continues until the task is complete 105 | 5. Maximum 300 iterations with automatic context compression 106 | 107 | ### Context Management 108 | 109 | - **Token Limit**: 200,000 tokens 110 | - **Auto-Compression**: Triggers at 180,000 tokens (90% of limit) 111 | - **Backups**: Automatic context summaries every 50 iterations 112 | - **Recovery**: All summaries saved with timestamps for resumption 113 | 114 | ## Project Structure 115 | 116 | ``` 117 | kimi-writer/ 118 | ├── kimi-writer.py # Main agent 119 | ├── tools/ 120 | │ ├── __init__.py # Tool registry 121 | │ ├── writer.py # File writing tool 122 | │ ├── project.py # Project management tool 123 | │ └── compression.py # Context compression tool 124 | ├── utils.py # Utilities (token counting, etc.) 125 | ├── requirements.txt # Python dependencies 126 | ├── env.example # Example configuration 127 | ├── .gitignore # Git ignore rules 128 | └── README.md # This file 129 | 130 | # Generated during use: 131 | output/ # All AI-generated projects go here 132 | ├── your_project_name/ # Created by the agent 133 | │ ├── chapter_01.md # Written by the agent 134 | │ ├── chapter_02.md 135 | │ └── .context_summary_*.md # Auto-saved context summaries 136 | └── another_project/ 137 | └── ... 138 | ``` 139 | 140 | ## Examples 141 | 142 | ### Example 1: Novel 143 | ```bash 144 | uv run kimi-writer.py "Write a mystery novel set in Victorian London with 10 chapters" 145 | ``` 146 | 147 | ### Example 2: Short Story Collection 148 | ```bash 149 | uv run kimi-writer.py "Create 7 interconnected sci-fi short stories exploring the theme of memory" 150 | ``` 151 | 152 | ### Example 3: Book 153 | ```bash 154 | uv run kimi-writer.py "Write a comprehensive guide to Python programming with 15 chapters" 155 | ``` 156 | 157 | ## Advanced Features 158 | 159 | ### Real-Time Streaming 160 | Watch the agent think and write in real-time: 161 | - 🧠 **Reasoning Stream**: See the agent's thought process as it plans 162 | - 💬 **Content Stream**: Watch stories being written character by character 163 | - 🔧 **Tool Call Progress**: Live updates when generating large content (shows character/word count) 164 | - ⚡ **No Waiting**: Immediate feedback - no more staring at a blank screen 165 | 166 | ### Iteration Counter 167 | The agent displays its progress: `Iteration X/300` 168 | 169 | ### Token Monitoring 170 | Real-time token usage: `Current tokens: 45,234/200,000 (22.6%)` 171 | 172 | ### Graceful Interruption 173 | Press `Ctrl+C` to interrupt. The agent will save the current context for recovery. 174 | 175 | ## Tips for Best Results 176 | 177 | 1. **Be Specific**: Clear prompts get better results 178 | - Good: "Create a 5-chapter romance novel set in modern Tokyo" 179 | - Less good: "Write something interesting" 180 | 181 | 2. **Let It Work**: The agent works autonomously - it will plan and execute the full task 182 | 183 | 3. **Recovery is Easy**: If interrupted, just use the `--recover` flag with the latest context summary 184 | 185 | 4. **Check Progress**: Generated files appear in real-time in the project folder 186 | 187 | ## Troubleshooting 188 | 189 | ### "MOONSHOT_API_KEY environment variable not set" 190 | Make sure you have created a `.env` file in the project root with your API key: 191 | ```bash 192 | MOONSHOT_API_KEY=your-actual-api-key-here 193 | ``` 194 | 195 | ### "401 Unauthorized" or Authentication errors 196 | - Verify your API key is correct in the `.env` file 197 | - Make sure you're using the correct base URL: `https://api.moonshot.ai/v1` 198 | - Get your API key from: https://platform.moonshot.cn/ 199 | 200 | ### "Error creating project folder" 201 | Check write permissions in the current directory 202 | 203 | ### Agent seems stuck 204 | The agent can run up to 300 iterations. For very complex tasks, this is normal. Check the project folder to see progress. 205 | 206 | ### Token limit issues 207 | The agent automatically compresses context at 180K tokens. If you see compression messages, the system is working correctly. 208 | 209 | ## Technical Details 210 | 211 | - **Model**: kimi-k2-thinking 212 | - **Temperature**: 1.0 (optimized for this model) 213 | - **Max Tokens per Call**: 65,536 (64K) 214 | - **Context Window**: 200,000 tokens 215 | - **Max Iterations**: 300 216 | - **Compression Threshold**: 180,000 tokens 217 | 218 | You can customize this as you please. 219 | 220 | ## License 221 | 222 | MIT License with Attribution Requirement - see [LICENSE](LICENSE) file for details. 223 | 224 | **Commercial Use**: If you use this software in a commercial product, you must provide clear attribution to Pietro Schirano (@Doriandarko). 225 | 226 | **API Usage**: This project uses the Moonshot AI API. Please refer to Moonshot AI's terms of service for API usage guidelines. 227 | 228 | ## Credits 229 | 230 | - **Created by**: Pietro Schirano ([@Doriandarko](https://github.com/Doriandarko)) 231 | - **Powered by**: Moonshot AI's kimi-k2-thinking model 232 | - **Repository**: https://github.com/Doriandarko/kimi-writer 233 | 234 | ## Star History 235 | 236 | 237 | 241 | 245 | Star History Chart 249 | 250 | 251 | -------------------------------------------------------------------------------- /kimi-writer.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | """ 3 | Kimi Writing Agent - An autonomous agent for creative writing tasks. 4 | 5 | This agent uses the kimi-k2-thinking model to create novels, books, 6 | and short story collections based on user prompts. 7 | """ 8 | 9 | import os 10 | import sys 11 | import json 12 | import argparse 13 | from dotenv import load_dotenv 14 | from openai import OpenAI 15 | from typing import List, Dict, Any 16 | 17 | # Load environment variables from .env file 18 | load_dotenv() 19 | 20 | from utils import ( 21 | estimate_token_count, 22 | get_tool_definitions, 23 | get_tool_map, 24 | get_system_prompt 25 | ) 26 | from tools.compression import compress_context_impl 27 | 28 | 29 | # Constants 30 | MAX_ITERATIONS = 300 31 | TOKEN_LIMIT = 200000 32 | COMPRESSION_THRESHOLD = 180000 # Trigger compression at 90% of limit 33 | MODEL_NAME = "kimi-k2-thinking" 34 | BACKUP_INTERVAL = 50 # Save backup summary every N iterations 35 | 36 | 37 | def load_context_from_file(file_path: str) -> str: 38 | """ 39 | Loads context from a summary file for recovery. 40 | 41 | Args: 42 | file_path: Path to the context summary file 43 | 44 | Returns: 45 | Content of the file as string 46 | """ 47 | try: 48 | with open(file_path, 'r', encoding='utf-8') as f: 49 | content = f.read() 50 | print(f"✓ Loaded context from: {file_path}\n") 51 | return content 52 | except Exception as e: 53 | print(f"✗ Error loading context file: {e}") 54 | sys.exit(1) 55 | 56 | 57 | def get_user_input() -> tuple[str, bool]: 58 | """ 59 | Gets user input from command line, either as a prompt or recovery file. 60 | 61 | Returns: 62 | Tuple of (prompt/context, is_recovery_mode) 63 | """ 64 | parser = argparse.ArgumentParser( 65 | description="Kimi Writing Agent - Create novels, books, and short stories", 66 | formatter_class=argparse.RawDescriptionHelpFormatter, 67 | epilog=""" 68 | Examples: 69 | # Fresh start with inline prompt 70 | python kimi-writer.py "Create a collection of sci-fi short stories" 71 | 72 | # Recovery mode from previous context 73 | python kimi-writer.py --recover my_project/.context_summary_20250107_143022.md 74 | """ 75 | ) 76 | 77 | parser.add_argument( 78 | 'prompt', 79 | nargs='?', 80 | help='Your writing request (e.g., "Create a mystery novel")' 81 | ) 82 | parser.add_argument( 83 | '--recover', 84 | type=str, 85 | help='Path to a context summary file to continue from' 86 | ) 87 | 88 | args = parser.parse_args() 89 | 90 | # Check if recovery mode 91 | if args.recover: 92 | context = load_context_from_file(args.recover) 93 | return context, True 94 | 95 | # Check if prompt provided as argument 96 | if args.prompt: 97 | return args.prompt, False 98 | 99 | # Interactive prompt 100 | print("=" * 60) 101 | print("Kimi Writing Agent") 102 | print("=" * 60) 103 | print("\nEnter your writing request (or 'quit' to exit):") 104 | print("Example: Create a collection of 15 sci-fi short stories\n") 105 | 106 | prompt = input("> ").strip() 107 | 108 | if prompt.lower() in ['quit', 'exit', 'q']: 109 | print("Goodbye!") 110 | sys.exit(0) 111 | 112 | if not prompt: 113 | print("Error: Empty prompt. Please provide a writing request.") 114 | sys.exit(1) 115 | 116 | return prompt, False 117 | 118 | 119 | def convert_message_for_api(msg: Any) -> Dict[str, Any]: 120 | """ 121 | Converts a message object to a dictionary suitable for API calls. 122 | Preserves reasoning_content if present. 123 | 124 | Args: 125 | msg: Message object (can be OpenAI message object or dict) 126 | 127 | Returns: 128 | Dictionary representation of the message 129 | """ 130 | if isinstance(msg, dict): 131 | return msg 132 | 133 | # Convert OpenAI message object to dict 134 | msg_dict = { 135 | "role": msg.role, 136 | } 137 | 138 | if msg.content: 139 | msg_dict["content"] = msg.content 140 | 141 | # Preserve reasoning_content if present 142 | if hasattr(msg, "reasoning_content"): 143 | reasoning = getattr(msg, "reasoning_content") 144 | if reasoning: 145 | msg_dict["reasoning_content"] = reasoning 146 | 147 | # Preserve tool calls if present 148 | if hasattr(msg, "tool_calls") and msg.tool_calls: 149 | msg_dict["tool_calls"] = [ 150 | { 151 | "id": tc.id, 152 | "type": "function", 153 | "function": { 154 | "name": tc.function.name, 155 | "arguments": tc.function.arguments 156 | } 157 | } 158 | for tc in msg.tool_calls 159 | ] 160 | 161 | # Preserve tool call id for tool response messages 162 | if hasattr(msg, "tool_call_id") and msg.tool_call_id: 163 | msg_dict["tool_call_id"] = msg.tool_call_id 164 | 165 | if hasattr(msg, "name") and msg.name: 166 | msg_dict["name"] = msg.name 167 | 168 | return msg_dict 169 | 170 | 171 | def main(): 172 | """Main agent loop.""" 173 | 174 | # Get API key 175 | api_key = os.getenv("MOONSHOT_API_KEY") 176 | if not api_key: 177 | print("Error: MOONSHOT_API_KEY environment variable not set.") 178 | print("Please set your API key: export MOONSHOT_API_KEY='your-key-here'") 179 | sys.exit(1) 180 | 181 | base_url = os.getenv("MOONSHOT_BASE_URL", "https://api.moonshot.ai/v1") 182 | 183 | # Debug: Show that key is loaded (masked for security) 184 | if len(api_key) > 8: 185 | print(f"✓ API Key loaded: {api_key[:4]}...{api_key[-4:]}") 186 | else: 187 | print(f"⚠️ Warning: API key seems too short ({len(api_key)} chars)") 188 | print(f"✓ Base URL: {base_url}\n") 189 | 190 | # Initialize OpenAI client 191 | client = OpenAI( 192 | api_key=api_key, 193 | base_url=base_url, 194 | ) 195 | 196 | # Get user input 197 | user_prompt, is_recovery = get_user_input() 198 | 199 | # Initialize message history 200 | messages = [ 201 | {"role": "system", "content": get_system_prompt()} 202 | ] 203 | 204 | if is_recovery: 205 | messages.append({ 206 | "role": "user", 207 | "content": f"[RECOVERED CONTEXT]\n\n{user_prompt}\n\n[END RECOVERED CONTEXT]\n\nPlease continue the work from where we left off." 208 | }) 209 | print("🔄 Recovery mode: Continuing from previous context\n") 210 | else: 211 | messages.append({ 212 | "role": "user", 213 | "content": user_prompt 214 | }) 215 | print(f"\n📝 Task: {user_prompt}\n") 216 | 217 | # Get tool definitions and mapping 218 | tools = get_tool_definitions() 219 | tool_map = get_tool_map() 220 | 221 | print("=" * 60) 222 | print("Starting Kimi Writing Agent") 223 | print("=" * 60) 224 | print(f"Model: {MODEL_NAME}") 225 | print(f"Max iterations: {MAX_ITERATIONS}") 226 | print(f"Context limit: {TOKEN_LIMIT:,} tokens") 227 | print(f"Auto-compression at: {COMPRESSION_THRESHOLD:,} tokens") 228 | print("=" * 60 + "\n") 229 | 230 | # Main agent loop 231 | for iteration in range(1, MAX_ITERATIONS + 1): 232 | print(f"\n{'─' * 60}") 233 | print(f"Iteration {iteration}/{MAX_ITERATIONS}") 234 | print(f"{'─' * 60}") 235 | 236 | # Check token count before making API call 237 | try: 238 | token_count = estimate_token_count(base_url, api_key, MODEL_NAME, messages) 239 | print(f"📊 Current tokens: {token_count:,}/{TOKEN_LIMIT:,} ({token_count/TOKEN_LIMIT*100:.1f}%)") 240 | 241 | # Trigger compression if approaching limit 242 | if token_count >= COMPRESSION_THRESHOLD: 243 | print(f"\n⚠️ Approaching token limit! Compressing context...") 244 | compression_result = compress_context_impl( 245 | messages=messages, 246 | client=client, 247 | model=MODEL_NAME, 248 | keep_recent=10 249 | ) 250 | 251 | if "compressed_messages" in compression_result: 252 | messages = compression_result["compressed_messages"] 253 | print(f"✓ {compression_result['message']}") 254 | print(f"✓ Estimated tokens saved: ~{compression_result.get('tokens_saved', 0):,}") 255 | 256 | # Recalculate token count 257 | token_count = estimate_token_count(base_url, api_key, MODEL_NAME, messages) 258 | print(f"📊 New token count: {token_count:,}/{TOKEN_LIMIT:,}\n") 259 | 260 | except Exception as e: 261 | print(f"⚠️ Warning: Could not estimate token count: {e}") 262 | token_count = 0 263 | 264 | # Auto-backup every N iterations 265 | if iteration % BACKUP_INTERVAL == 0: 266 | print(f"💾 Auto-backup (iteration {iteration})...") 267 | try: 268 | compression_result = compress_context_impl( 269 | messages=messages, 270 | client=client, 271 | model=MODEL_NAME, 272 | keep_recent=len(messages) # Keep all messages, just save summary 273 | ) 274 | if compression_result.get("summary_file"): 275 | print(f"✓ Backup saved: {os.path.basename(compression_result['summary_file'])}\n") 276 | except Exception as e: 277 | print(f"⚠️ Warning: Backup failed: {e}\n") 278 | 279 | # Call the model 280 | try: 281 | print("🤖 Calling kimi-k2-thinking model...\n") 282 | 283 | stream = client.chat.completions.create( 284 | model=MODEL_NAME, 285 | messages=messages, 286 | max_tokens=65536, # 64K tokens 287 | tools=tools, 288 | temperature=1.0, 289 | stream=True, # Enable streaming 290 | ) 291 | 292 | # Accumulate the streaming response 293 | reasoning_content = "" 294 | content_text = "" 295 | tool_calls_data = [] 296 | role = None 297 | finish_reason = None 298 | 299 | # Track if we've printed headers 300 | reasoning_header_printed = False 301 | content_header_printed = False 302 | tool_call_header_printed = False 303 | last_tool_index = -1 304 | 305 | # Process the stream 306 | for chunk in stream: 307 | if not chunk.choices: 308 | continue 309 | 310 | delta = chunk.choices[0].delta 311 | finish_reason = chunk.choices[0].finish_reason 312 | 313 | # Get role if present (first chunk) 314 | if hasattr(delta, "role") and delta.role: 315 | role = delta.role 316 | 317 | # Handle reasoning_content streaming 318 | if hasattr(delta, "reasoning_content") and delta.reasoning_content: 319 | if not reasoning_header_printed: 320 | print("=" * 60) 321 | print(f"🧠 Reasoning (Iteration {iteration})") 322 | print("=" * 60) 323 | reasoning_header_printed = True 324 | 325 | print(delta.reasoning_content, end="", flush=True) 326 | reasoning_content += delta.reasoning_content 327 | 328 | # Handle regular content streaming 329 | if hasattr(delta, "content") and delta.content: 330 | # Close reasoning section if it was open 331 | if reasoning_header_printed and not content_header_printed: 332 | print("\n" + "=" * 60 + "\n") 333 | 334 | if not content_header_printed: 335 | print("💬 Response:") 336 | print("-" * 60) 337 | content_header_printed = True 338 | 339 | print(delta.content, end="", flush=True) 340 | content_text += delta.content 341 | 342 | # Handle tool_calls 343 | if hasattr(delta, "tool_calls") and delta.tool_calls: 344 | for tc_delta in delta.tool_calls: 345 | # Ensure we have enough slots in tool_calls_data 346 | while len(tool_calls_data) <= tc_delta.index: 347 | tool_calls_data.append({ 348 | "id": None, 349 | "type": "function", 350 | "function": {"name": "", "arguments": ""}, 351 | "chars_received": 0 352 | }) 353 | 354 | tc = tool_calls_data[tc_delta.index] 355 | 356 | # Print header when we start receiving a tool call 357 | if tc_delta.index != last_tool_index: 358 | if reasoning_header_printed or content_header_printed: 359 | print("\n" + "=" * 60 + "\n") 360 | 361 | if hasattr(tc_delta, "function") and tc_delta.function.name: 362 | print(f"🔧 Preparing tool call: {tc_delta.function.name}") 363 | print("─" * 60) 364 | tool_call_header_printed = True 365 | last_tool_index = tc_delta.index 366 | 367 | if tc_delta.id: 368 | tc["id"] = tc_delta.id 369 | if hasattr(tc_delta, "function"): 370 | if tc_delta.function.name: 371 | tc["function"]["name"] = tc_delta.function.name 372 | if tc_delta.function.arguments: 373 | tc["function"]["arguments"] += tc_delta.function.arguments 374 | tc["chars_received"] += len(tc_delta.function.arguments) 375 | 376 | # Show progress indicator every 500 characters 377 | if tc["chars_received"] % 500 == 0 or tc["chars_received"] < 100: 378 | # Calculate approximate words (rough estimate: 5 chars per word) 379 | words = tc["chars_received"] // 5 380 | print(f"\r💬 Generating arguments... {tc['chars_received']:,} characters (~{words:,} words)", end="", flush=True) 381 | 382 | # Print closing for content if it was printed 383 | if content_header_printed: 384 | print("\n" + "-" * 60 + "\n") 385 | 386 | # Print completion for tool calls if any were received 387 | if tool_call_header_printed: 388 | print("\n✓ Tool call complete") 389 | print("─" * 60 + "\n") 390 | 391 | # Reconstruct the message object from accumulated data 392 | class ReconstructedMessage: 393 | def __init__(self): 394 | self.role = role or "assistant" 395 | self.content = content_text if content_text else None 396 | self.reasoning_content = reasoning_content if reasoning_content else None 397 | self.tool_calls = None 398 | 399 | if tool_calls_data: 400 | # Convert to proper format 401 | from openai.types.chat import ChatCompletionMessageToolCall 402 | from openai.types.chat.chat_completion_message_tool_call import Function 403 | 404 | self.tool_calls = [] 405 | for tc in tool_calls_data: 406 | if tc["id"]: # Only add if we have an ID 407 | tool_call = type('ToolCall', (), { 408 | 'id': tc["id"], 409 | 'type': 'function', 410 | 'function': type('Function', (), { 411 | 'name': tc["function"]["name"], 412 | 'arguments': tc["function"]["arguments"] 413 | })() 414 | })() 415 | self.tool_calls.append(tool_call) 416 | 417 | message = ReconstructedMessage() 418 | 419 | # Convert message to dict and add to history 420 | # Important: preserve the full message object structure 421 | messages.append(convert_message_for_api(message)) 422 | 423 | # Check if the model called any tools 424 | if not message.tool_calls: 425 | print("=" * 60) 426 | print("✅ TASK COMPLETED") 427 | print("=" * 60) 428 | print(f"Completed in {iteration} iteration(s)") 429 | print("=" * 60) 430 | break 431 | 432 | # Handle tool calls 433 | print(f"\n🔧 Model decided to call {len(message.tool_calls)} tool(s):") 434 | 435 | for tool_call in message.tool_calls: 436 | func_name = tool_call.function.name 437 | args_str = tool_call.function.arguments 438 | 439 | try: 440 | args = json.loads(args_str) 441 | except json.JSONDecodeError: 442 | args = {} 443 | 444 | print(f"\n → {func_name}") 445 | print(f" Arguments: {json.dumps(args, ensure_ascii=False, indent=6)}") 446 | 447 | # Get the tool implementation 448 | tool_func = tool_map.get(func_name) 449 | 450 | if not tool_func: 451 | result = f"Error: Unknown tool '{func_name}'" 452 | print(f" ✗ {result}") 453 | else: 454 | # Special handling for compress_context (needs extra params) 455 | if func_name == "compress_context": 456 | result_data = compress_context_impl( 457 | messages=messages, 458 | client=client, 459 | model=MODEL_NAME, 460 | keep_recent=10 461 | ) 462 | result = result_data.get("message", "Compression completed") 463 | 464 | # Update messages with compressed version 465 | if "compressed_messages" in result_data: 466 | messages = result_data["compressed_messages"] 467 | else: 468 | # Call the tool with its arguments 469 | result = tool_func(**args) 470 | 471 | # Print result (truncate if too long) 472 | if len(str(result)) > 200: 473 | print(f" ✓ {str(result)[:200]}...") 474 | else: 475 | print(f" ✓ {result}") 476 | 477 | # Add tool result to messages 478 | tool_message = { 479 | "role": "tool", 480 | "tool_call_id": tool_call.id, 481 | "name": func_name, 482 | "content": str(result) 483 | } 484 | messages.append(tool_message) 485 | 486 | except KeyboardInterrupt: 487 | print("\n\n⚠️ Interrupted by user. Saving context...") 488 | # Save current context before exiting 489 | try: 490 | compression_result = compress_context_impl( 491 | messages=messages, 492 | client=client, 493 | model=MODEL_NAME, 494 | keep_recent=len(messages) 495 | ) 496 | if compression_result.get("summary_file"): 497 | print(f"✓ Context saved to: {compression_result['summary_file']}") 498 | print(f"\nTo resume, run:") 499 | print(f" python kimi-writer.py --recover {compression_result['summary_file']}") 500 | except: 501 | pass 502 | sys.exit(0) 503 | 504 | except Exception as e: 505 | print(f"\n✗ Error during iteration {iteration}: {e}") 506 | print(f"Attempting to continue...\n") 507 | continue 508 | 509 | # If we hit max iterations 510 | if iteration >= MAX_ITERATIONS: 511 | print("\n" + "=" * 60) 512 | print("⚠️ MAX ITERATIONS REACHED") 513 | print("=" * 60) 514 | print(f"\nReached maximum of {MAX_ITERATIONS} iterations.") 515 | print("Saving final context...") 516 | 517 | try: 518 | compression_result = compress_context_impl( 519 | messages=messages, 520 | client=client, 521 | model=MODEL_NAME, 522 | keep_recent=len(messages) 523 | ) 524 | if compression_result.get("summary_file"): 525 | print(f"✓ Context saved to: {compression_result['summary_file']}") 526 | print(f"\nTo resume, run:") 527 | print(f" python kimi-writer.py --recover {compression_result['summary_file']}") 528 | except Exception as e: 529 | print(f"✗ Error saving context: {e}") 530 | 531 | 532 | if __name__ == "__main__": 533 | main() 534 | 535 | --------------------------------------------------------------------------------