Skip to content
Merged
Show file tree
Hide file tree
Changes from 10 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions cli/CLAUDE.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
- use uv for running python scripts
- Write minimal comments without extra explanations
2 changes: 2 additions & 0 deletions cli/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ dependencies = [
"pyyaml>=6.0.2",
"openai>=1.98.0",
"anthropic>=0.60.0",
"docker>=7.1.0",
]

[project.urls]
Expand All @@ -37,6 +38,7 @@ infragpt = "infragpt.main:cli"




[dependency-groups]
dev = [
"pytest>=8.4.1",
Expand Down
61 changes: 2 additions & 59 deletions cli/src/infragpt/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,11 @@
from .history import log_interaction
from .tools import ToolExecutionCancelled

# Import prompt_toolkit for better input handling
import pathlib
from prompt_toolkit import PromptSession
from prompt_toolkit.history import FileHistory


# Create a simple system prompt instead of using the old prompts module
def get_system_prompt():
return """You are an intelligent shell operations assistant. You help users with:

Expand All @@ -36,7 +34,6 @@ def get_system_prompt():
Be concise but thorough in your responses."""


# Initialize console for rich output
console = Console()


Expand Down Expand Up @@ -68,25 +65,18 @@ def add_message(
"execute_shell_command" # Add tool name for compatibility
)

# Keep system message separate
if role == "system":
self.system_message = message_dict
else:
self.messages.append(message_dict)

# Maintain context window
if len(self.messages) > self.max_messages:
self.messages = self.messages[-self.max_messages :]

def get_context_messages(self) -> List[Dict[str, Any]]:
"""Get messages formatted for LLM API."""
context = []

# Add system message first
if self.system_message:
context.append(self.system_message)

# Add conversation messages
context.extend(self.messages)

return context
Expand All @@ -106,35 +96,22 @@ def __init__(self, model_string: str, api_key: str, verbose: bool = False):
self.verbose = verbose
self.context = ConversationContext()

# Create LLM adapter
self.llm_adapter = get_llm_adapter(
model_string=model_string, api_key=api_key, verbose=verbose
)

# Initialize command history session
self._setup_command_history()

# Don't set up custom signal handlers - let prompt_toolkit handle them

# Initialize system message
self._initialize_system_message()

def _setup_command_history(self):
"""Setup command history with persistent storage."""
try:
# Create history directory following InfraGPT conventions
history_dir = pathlib.Path.home() / ".infragpt"
history_dir.mkdir(exist_ok=True)

# History file for command-line input
history_file = history_dir / "history"

# WARNING: FileHistory writes clear-text inputs to disk.
# This is a security risk if users enter secrets at the prompt.
# Create PromptSession with FileHistory
# WARNING: FileHistory writes clear-text inputs to disk - security risk if users enter secrets
self.prompt_session = PromptSession(history=FileHistory(str(history_file)))

# Best-effort: restrict file permissions to user-only
try:
history_file.chmod(0o600)
except Exception:
Expand All @@ -144,7 +121,6 @@ def _setup_command_history(self):
console.print(f"[dim]Command history: {history_file}[/dim]")

except Exception as e:
# Fallback to no history if setup fails
self.prompt_session = PromptSession()
if self.verbose:
console.print(
Expand All @@ -167,8 +143,6 @@ def run_interactive_session(self):
)

console.print(f"[yellow]Model:[/yellow] [bold]{self.model_string}[/bold]")

# Validate API key
console.print("[dim]Validating API key...[/dim]")
try:
if self.llm_adapter.validate_api_key():
Expand All @@ -180,34 +154,26 @@ def run_interactive_session(self):
console.print(f"[red]✗ API key validation failed: {e}[/red]")
return

# Show initial prompt
console.print("[bold cyan]What would you like me to help with?[/bold cyan]")
console.print(
"[dim]Press Ctrl+D to exit, Ctrl+C to interrupt operations[/dim]\n"
)

while True:
try:
# Get user input
user_input = self._get_user_input()
if not user_input:
continue # Go back to prompt for empty input
continue

# Check for exit commands
if user_input.lower() in ["exit", "quit", "bye"]:
break

# Add user message to context
self.context.add_message("user", user_input)

# Process with LLM
self._process_user_input(user_input)

except KeyboardInterrupt:
# This shouldn't happen with prompt_toolkit, but handle just in case
continue
except EOFError:
# Ctrl+D - exit the application
console.print("\n[dim]EOF received (Ctrl+D). Exiting...[/dim]")
break

Expand All @@ -218,19 +184,15 @@ def _get_user_input(self) -> str:
try:
return self.prompt_session.prompt("> ")
except KeyboardInterrupt:
# Ctrl+C should just return empty string to continue
return ""
except EOFError:
# Ctrl+D - let this propagate for proper exit handling
raise

def _process_user_input(self, user_input: str):
"""Process user input with direct SDK streaming and tool execution."""
try:
# Get context messages
messages = self.context.get_context_messages()

# Debug: Show message structure if verbose
if self.verbose:
console.print(f"[dim]Context has {len(messages)} messages[/dim]")
for i, msg in enumerate(messages):
Expand All @@ -242,21 +204,15 @@ def _process_user_input(self, user_input: str):
f"[dim] {i}: {role} (content: {content_len} chars, tools: {has_tools}, tool_id: {has_tool_id})[/dim]"
)

# Show thinking and stream response
console.print("\n[dim]Thinking...[/dim]")

response_content = ""
first_content = True

# Stream response using new adapter with interrupt checking
# Note: stream_with_tools already handles the complete tool execution cycle
# including getting the final response after tool execution
try:
for chunk in self.llm_adapter.stream_with_tools(messages):
# Handle content streaming
if chunk.content:
if first_content:
# Clear thinking message and show A: prefix
console.print(
"\033[1A\033[K", end=""
) # Move up and clear line
Expand All @@ -266,7 +222,6 @@ def _process_user_input(self, user_input: str):
response_content += chunk.content
console.print(chunk.content, end="")

# Handle finish reason
if chunk.finish_reason:
if self.verbose:
console.print(
Expand All @@ -276,23 +231,15 @@ def _process_user_input(self, user_input: str):
console.print("\n[yellow]Operation cancelled by user.[/yellow]")
return

# Add newline after streaming
if response_content:
console.print()

# Add assistant message to context if we have content
if response_content:
self.context.add_message("assistant", response_content)

# Log interaction
self._log_interaction(user_input, response_content)

except ToolExecutionCancelled:
# User cancelled tool execution - just return to prompt
# No need to print anything - the tool already printed a message
return
except KeyboardInterrupt:
# Ctrl+C during streaming - just return to prompt (message already printed by LLM adapter)
return
except Exception as e:
console.print(f"[bold red]Error processing input:[/bold red] {e}")
Expand All @@ -304,19 +251,15 @@ def _process_user_input(self, user_input: str):
def _log_interaction(self, user_input: str, response: str):
"""Log the interaction for history. Sensitive fields are excluded explicitly."""
try:
# Only safe fields, do not include api_key or any secrets
interaction_data = {
"user_input": user_input,
"assistant_response": response,
"model": self.model_string,
"timestamp": datetime.utcnow().isoformat() + "Z",
# Not persisted: allow-list will drop this; used only for error reporting in history.py
"verbose": self.verbose,
}
# Note: self.api_key is NOT logged ever
log_interaction("agent_conversation_v2", interaction_data)
except Exception as e:
# Don't let logging failures interrupt the session
if self.verbose:
console.print(f"[dim]Warning: Could not log interaction: {e}[/dim]")

Expand Down
Loading
Loading