@@ -12,24 +12,103 @@ import (
1212 "github.com/spf13/cobra"
1313
1414 "github.com/inercia/MCPShell/pkg/agent"
15- "github.com/inercia/MCPShell/pkg/common"
1615)
1716
17+ // buildAgentConfig creates an AgentConfig by merging command-line flags with configuration file
18+ func buildAgentConfig () (agent.AgentConfig , error ) {
19+ // Load configuration from file
20+ config , err := agent .GetConfig ()
21+ if err != nil {
22+ return agent.AgentConfig {}, fmt .Errorf ("failed to load config: %w" , err )
23+ }
24+
25+ // Start with default model from config file
26+ var modelConfig agent.ModelConfig
27+ if defaultModel := config .GetDefaultModel (); defaultModel != nil {
28+ modelConfig = * defaultModel
29+ }
30+
31+ // Override with command-line flags if provided
32+ if agentModel != "" {
33+ // Check if the specified model exists in config
34+ if configModel := config .GetModelByName (agentModel ); configModel != nil {
35+ modelConfig = * configModel
36+ } else {
37+ // Use command-line model name if not found in config
38+ modelConfig .Model = agentModel
39+ }
40+ }
41+
42+ // Merge system prompts from config file and command-line
43+ if agentSystemPrompt != "" {
44+ // Join system prompts from config with command-line system prompt
45+ var allSystemPrompts []string
46+
47+ // Add existing system prompts from config
48+ if modelConfig .Prompts .HasSystemPrompts () {
49+ allSystemPrompts = append (allSystemPrompts , modelConfig .Prompts .System ... )
50+ }
51+
52+ // Add command-line system prompt
53+ allSystemPrompts = append (allSystemPrompts , agentSystemPrompt )
54+
55+ // Update the prompts config with merged system prompts
56+ modelConfig .Prompts .System = allSystemPrompts
57+ // Clear user prompts as they should be ignored from config
58+ modelConfig .Prompts .User = nil
59+ } else {
60+ // No command-line system prompt provided, but still clear user prompts from config
61+ modelConfig .Prompts .User = nil
62+ }
63+ if agentOpenAIApiKey != "" {
64+ modelConfig .APIKey = agentOpenAIApiKey
65+ }
66+ if agentOpenAIApiURL != "" {
67+ modelConfig .APIURL = agentOpenAIApiURL
68+ }
69+
70+ // If no API key is set, try environment variable or handle template value
71+ switch modelConfig .APIKey {
72+ case "" :
73+ modelConfig .APIKey = os .Getenv ("OPENAI_API_KEY" )
74+ case "${OPENAI_API_KEY}" :
75+ // Handle environment variable substitution
76+ modelConfig .APIKey = os .Getenv ("OPENAI_API_KEY" )
77+ }
78+
79+ return agent.AgentConfig {
80+ ToolsFile : toolsFile ,
81+ UserPrompt : agentUserPrompt ,
82+ Once : agentOnce ,
83+ Version : version ,
84+ ModelConfig : modelConfig ,
85+ }, nil
86+ }
87+
1888// agentCommand is a command that executes the MCPShell as an agent
1989var agentCommand = & cobra.Command {
2090 Use : "agent" ,
2191 Short : "Execute the MCPShell as an agent" ,
2292 Long : `
2393
2494The agent command will execute the MCPShell as an agent, connecting to a remote LLM.
25- For example, you can do
2695
27- $ mcpshell agent --configfile=examples/config.yaml \
96+ Configuration is loaded from ~/.mcpshell/agent.yaml and can be overridden with command-line flags.
97+ The configuration file should contain model definitions with their API keys and prompts.
98+
99+ For example, you can do:
100+
101+ $ mcpshell agent --tools=examples/config.yaml \
28102 --model "gpt-4o" \
29103 --system-prompt "You are a helpful assistant that debugs performance issues" \
30- --user-prompt "I am having trouble with my computer. It is slow and I think it is due to the CPU usage."
104+ --user-prompt "I am having trouble with my computer. It is slow and I think it is due to the CPU usage."
31105
32- and the agent will try to debug the issue with the given tools.
106+ If a model is configured as default in the agent configuration file, you can omit the --model flag:
107+
108+ $ mcpshell agent --tools=examples/config.yaml \
109+ --user-prompt "I am having trouble with my computer. It is slow and I think it is due to the CPU usage."
110+
111+ The agent will try to debug the issue with the given tools.
33112` ,
34113 Args : cobra .NoArgs ,
35114 PreRunE : func (cmd * cobra.Command , args []string ) error {
@@ -39,136 +118,96 @@ and the agent will try to debug the issue with the given tools.
39118 return err
40119 }
41120
42- // Setup panic handler
43- defer common .RecoverPanic ()
44-
45- logger .Info ("Starting MCPShell agent" )
46-
47- // Check if config file is provided
48- if configFile == "" {
49- logger .Error ("Configuration file is required" )
50- return fmt .Errorf ("configuration file is required. Use --config or -c flag to specify the path" )
51- }
52-
53- // Check if model is provided
54- if agentModel == "" {
55- logger .Error ("LLM model is required" )
56- return fmt .Errorf ("LLM model is required. Use --model flag to specify the model" )
121+ // Build agent configuration
122+ agentConfig , err := buildAgentConfig ()
123+ if err != nil {
124+ return err
57125 }
58126
59- // Check if API key is provided or in environment
60- if agentOpenAIApiKey == "" {
61- agentOpenAIApiKey = os .Getenv ("OPENAI_API_KEY" )
62- if agentOpenAIApiKey == "" {
63- logger .Error ("OpenAI API key is required" )
64- return fmt .Errorf ("OpenAI API key is required. Use --api-key flag or set OPENAI_API_KEY environment variable" )
65- }
127+ // Validate agent configuration
128+ agentInstance := agent .New (agentConfig , logger )
129+ if err := agentInstance .Validate (); err != nil {
130+ return err
66131 }
67132
68133 return nil
69134 },
70135 RunE : func (cmd * cobra.Command , args []string ) error {
71- logger := common .GetLogger ()
72-
73- agentConfig := agent.AgentConfig {
74- ConfigFile : configFile ,
75- Model : agentModel ,
76- SystemPrompt : agentSystemPrompt ,
77- UserPrompt : agentUserPrompt ,
78- OpenAIApiKey : agentOpenAIApiKey ,
79- OpenAIApiURL : agentOpenAIApiURL ,
80- Once : agentOnce ,
81- Version : version ,
136+ // Initialize logger
137+ logger , err := initLogger ()
138+ if err != nil {
139+ return err
82140 }
83141
84- a := agent . New ( agentConfig , logger )
85-
86- if err := a . Validate (); err != nil {
87- return fmt . Errorf ( "agent validation failed: %w" , err )
142+ // Build agent configuration
143+ agentConfig , err := buildAgentConfig ()
144+ if err != nil {
145+ return err
88146 }
89147
148+ // Create agent instance
149+ agentInstance := agent .New (agentConfig , logger )
150+
151+ // Create channels for user input and agent output
152+ userInput := make (chan string )
153+ agentOutput := make (chan string )
154+
90155 ctx , cancel := context .WithCancel (context .Background ())
91156 defer cancel ()
92157
93- // Handle Ctrl+C (SIGINT) and SIGTERM to gracefully shut down
94- sigChan := make (chan os.Signal , 1 )
95- signal .Notify (sigChan , syscall .SIGINT , syscall .SIGTERM )
96- go func () {
97- <- sigChan
98- logger .Info ("Received interrupt signal, cancelling agent context..." )
99- cancel ()
100- }()
101-
102- userInputChan := make (chan string )
103- agentOutputChan := make (chan string )
158+ // Setup signal handling for graceful shutdown
159+ signalChan := make (chan os.Signal , 1 )
160+ signal .Notify (signalChan , os .Interrupt , syscall .SIGTERM )
104161
105162 var wg sync.WaitGroup
163+ wg .Add (1 )
164+ go func () {
165+ defer wg .Done ()
166+ select {
167+ case <- signalChan :
168+ logger .Info ("Received interrupt signal, shutting down..." )
169+ cancel ()
170+ case <- ctx .Done ():
171+ }
172+ }()
106173
107- // Goroutine to read from stdin and send to userInputChan
108- if ! agentOnce {
174+ // Start a goroutine to read user input only when not in --once mode
175+ if ! agentConfig . Once {
109176 wg .Add (1 )
110177 go func () {
111178 defer wg .Done ()
112- defer close (userInputChan )
113179 scanner := bufio .NewScanner (os .Stdin )
114- for scanner . Scan () {
180+ for {
115181 select {
116- case userInputChan <- scanner .Text ():
117182 case <- ctx .Done ():
118- logger .Info ("Context cancelled, stopping stdin reader." )
119183 return
184+ default :
185+ if scanner .Scan () {
186+ userInput <- scanner .Text ()
187+ } else {
188+ close (userInput )
189+ return
190+ }
120191 }
121192 }
122- if err := scanner .Err (); err != nil {
123- logger .Error ("Error reading from stdin: %v" , err )
124- }
125- logger .Info ("Stdin scanner finished." )
126193 }()
127- } else {
128- // In one-shot mode, we'll close the channel when Run completes
129- logger .Info ("One-shot mode, skipping stdin reader." )
130194 }
131195
132- // Goroutine to read from agentOutputChan and print to stdout
196+ // Start the agent
133197 wg .Add (1 )
134198 go func () {
135199 defer wg .Done ()
136- for {
137- select {
138- case output , ok := <- agentOutputChan :
139- if ! ok {
140- logger .Info ("Agent output channel closed, stdout writer finishing." )
141- return
142- }
143- fmt .Print (output )
144- case <- ctx .Done ():
145- logger .Info ("Context cancelled, stopping stdout writer." )
146- for output := range agentOutputChan {
147- fmt .Print (output )
148- }
149- return
150- }
200+ if err := agentInstance .Run (ctx , userInput , agentOutput ); err != nil {
201+ logger .Error ("Agent encountered an error: %v" , err )
151202 }
152203 }()
153204
154- err := a .Run (ctx , userInputChan , agentOutputChan )
155-
156- cancel ()
157-
158- logger .Info ("Waiting for I/O goroutines to finish..." )
159- wg .Wait ()
160- logger .Info ("All goroutines finished." )
161-
162- if err != nil {
163- if err == context .Canceled || err == context .DeadlineExceeded {
164- logger .Info ("Agent run was cancelled: %v" , err )
165- return nil
166- }
167- logger .Error ("Agent execution failed: %v" , err )
168- return fmt .Errorf ("agent execution failed: %w" , err )
205+ // Print agent output
206+ for output := range agentOutput {
207+ fmt .Println (output )
169208 }
170209
171- logger . Info ( "Agent finished successfully." )
210+ wg . Wait ( )
172211 return nil
173212 },
174213}
@@ -186,7 +225,6 @@ func init() {
186225 agentCommand .Flags ().StringVarP (& agentOpenAIApiURL , "openai-api-url" , "b" , "" , "Base URL for the OpenAI API (optional)" )
187226 agentCommand .Flags ().BoolVarP (& agentOnce , "once" , "o" , false , "Exit after receiving a final response from the LLM (one-shot mode)" )
188227
189- // Mark required flags
190- _ = agentCommand .MarkFlagRequired ("config" )
191- _ = agentCommand .MarkFlagRequired ("model" )
228+ // Add config subcommand
229+ agentCommand .AddCommand (agentConfigCommand )
192230}
0 commit comments