1+ import logging
2+ import os
3+ from pathlib import Path
4+ from typing import List , Optional
5+
6+ from .base import CLITool
7+
8+ logger = logging .getLogger (__name__ )
9+
10+
11+ class KimiTool (CLITool ):
12+ """Moonshot AI Kimi CLI wrapper."""
13+
14+ command_name = "kimi"
15+ tool_key = "kimi"
16+ install_description = "Moonshot AI Kimi CLI"
17+
18+ def _write_config (
19+ self ,
20+ * ,
21+ endpoint_config : dict ,
22+ model : str ,
23+ provider_name : str = "kimi" ,
24+ ) -> Path :
25+ """Write kimi configuration to ~/.kimi/config.toml."""
26+ config_path = Path .home () / ".kimi" / "config.toml"
27+ config_path .parent .mkdir (parents = True , exist_ok = True )
28+
29+ # Create TOML content
30+ toml_content = f'''default_model = "{ model }
31+
32+ [providers.{ provider_name }
33+ type = "kimi"
34+ base_url = "{ endpoint_config ['endpoint' ]}
35+ api_key = "{ endpoint_config ['actual_api_key' ]}
36+
37+ [models.{ model }
38+ provider = "{ provider_name }
39+ model = "{ model }
40+ max_context_size = 262144
41+ '''
42+
43+ # Write to config file
44+ with open (config_path , "w" , encoding = "utf-8" ) as f :
45+ f .write (toml_content )
46+
47+ return config_path
48+
49+ def run (self , args : Optional [List [str ]] = None ) -> int :
50+ """
51+ Run the Moonshot AI Kimi CLI tool with the specified arguments.
52+
53+ Args:
54+ args: List of arguments to pass to the Kimi CLI
55+
56+ Returns:
57+ Exit code of the Kimi CLI process
58+ """
59+ args = args or []
60+
61+ # Load environment variables first
62+ self ._load_environment ()
63+
64+ # Check if the tool is installed and prompt for upgrade if needed
65+ if not self ._ensure_tool_installed (
66+ self .command_name , self .tool_key , self .install_description
67+ ):
68+ return 1
69+
70+ # Set up endpoint and model using the endpoint manager
71+ success , result = self ._setup_endpoint_and_models (
72+ "kimi" , select_multiple = False
73+ )
74+ if not success or result is None :
75+ return 1
76+
77+ # Extract endpoint configuration and selected model
78+ endpoint_config , _ , selected_models = result
79+
80+ # Handle the case where selected_models might be a string or tuple
81+ if isinstance (selected_models , str ):
82+ model = selected_models
83+ elif isinstance (selected_models , tuple ) and len (selected_models ) > 0 :
84+ model = selected_models [0 ] # Take the first model if multiple
85+ else :
86+ return self ._handle_error ("No model selected" )
87+
88+ if not endpoint_config or not model :
89+ return 1
90+
91+ # Write configuration to ~/.kimi/config.toml
92+ config_path = None
93+ try :
94+ config_path = self ._write_config (
95+ endpoint_config = endpoint_config ,
96+ model = model ,
97+ provider_name = "kimi"
98+ )
99+ print (f"[code-assistant-manager] Updated kimi config: { config_path } )
100+ except Exception as e :
101+ error_path = config_path or "~/.kimi/config.toml"
102+ return self ._handle_error (f"Failed to write { error_path } , e )
103+
104+ # Set up environment variables for Kimi
105+ env = os .environ .copy ()
106+ # Set TLS environment for Node.js
107+ self ._set_node_tls_env (env )
108+
109+ # Execute the Kimi CLI - it will read from the config file
110+ command = ["kimi" ] + args
111+
112+ # Display the complete command
113+ args_str = " " .join (args ) if args else ""
114+ command_str = f"kimi { args_str } .strip ()
115+ print ("" )
116+ print ("Complete command to execute:" )
117+ print (command_str )
118+ print ("" )
119+
120+ return self ._run_tool_with_env (command , env , "kimi" , interactive = True )
0 commit comments