agent

#!/usr/bin/env python3

# Copyright (c) 2025
# Licensed under the MIT License.
# See LICENSE file in the root directory of this source tree.
#
# Created by: [Aayush Chawla]
# Created on: February 7, 2025

from typing import Optional, List, Dict, Any
from datetime import datetime

from Logger import Logger, Colors
from Security import SecurityManager, Operation
from SystemIntegration import Functions, SystemCommands

import asyncio
import sqlite3
import os
import json
import aiohttp
import logging
import sys
import psutil
import re
import traceback


class ContextManager:
    def __init__(self, shell_pid: int, logger: Logger,
                 max_context_length: int = 4000):
        self.shell_pid = shell_pid
        self.logger = logger
        self.db_path = os.path.expanduser(
            f"~/.bashai/context_{shell_pid}.db")
        self.max_context_length = max_context_length
        self.current_context_length = 0
        self._init_db()
        self._cleanup_old_contexts()

    def _init_db(self):
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS command_history (
                    id INTEGER PRIMARY KEY,
                    timestamp TEXT,
                    command TEXT,
                    result TEXT,
                    status TEXT,
                    token_count INTEGER,
                    last_accessed TEXT
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS context_vars (
                    key TEXT PRIMARY KEY,
                    value TEXT,
                    timestamp TEXT,
                    last_accessed TEXT
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS context_metadata (
                    shell_pid INTEGER PRIMARY KEY,
                    created_at TEXT,
                    last_accessed TEXT,
                    total_tokens INTEGER DEFAULT 0
                )
            """)

            # Initialize metadata for this context
            conn.execute("""
                INSERT OR IGNORE INTO context_metadata
                (shell_pid, created_at, last_accessed, total_tokens)
                VALUES (?, ?, ?, 0)
            """, (self.shell_pid, datetime.now().isoformat(),
                  datetime.now().isoformat()))

        self._update_last_accessed()
        self._cleanup_old_contexts()

    def _cleanup_old_contexts(self, max_age_hours: int = 24):
        """
        Cleanup context DBs that haven't been accessed in the specified time.
        """
        try:
            db_dir = os.path.dirname(self.db_path)
            current_time = datetime.now()

            for filename in os.listdir(db_dir):
                if filename.startswith("context_") and filename.endswith(".db"):
                    db_path = os.path.join(db_dir, filename)
                    try:
                        with sqlite3.connect(db_path) as conn:
                            cursor = conn.execute(
                                "SELECT shell_pid, last_accessed "
                                "FROM context_metadata LIMIT 1"
                            )
                            row = cursor.fetchone()
                            if row:
                                last_accessed = datetime.fromisoformat(row[1])
                                age_in_seconds = (
                                    current_time - last_accessed).total_seconds()
                                if age_in_seconds > max_age_hours * 3600:
                                    # Check if the shell is still running
                                    try:
                                        if not psutil.pid_exists(row[0]):
                                            os.remove(db_path)
                                            logging.info(
                                                "Removed old context DB: "
                                                f"{filename}")
                                    except Exception as e:
                                        logging.error(
                                            f"Error checking process: {e}")
                    except sqlite3.Error as e:
                        logging.error(f"Error accessing DB {filename}: {e}")
                        # If DB is corrupted, remove it
                        os.remove(db_path)
        except Exception as e:
            logging.error(f"Error during context cleanup: {e}")

    def _update_last_accessed(self):
        """Update the last_accessed timestamp for this context."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "UPDATE context_metadata SET last_accessed = ? "
                "WHERE shell_pid = ?",
                (datetime.now().isoformat(), self.shell_pid)
            )

    def update_token_count(self, tokens_used: int):
        """
        Update the total token count and check if we're approaching the limit.
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "UPDATE context_metadata "
                "SET total_tokens = total_tokens + ? "
                "WHERE shell_pid = ?",
                (tokens_used, self.shell_pid)
            )
            cursor = conn.execute(
                "SELECT total_tokens FROM context_metadata "
                "WHERE shell_pid = ?",
                (self.shell_pid,)
            )
            total_tokens = cursor.fetchone()[0]

            # 80% warning threshold
            if total_tokens >= self.max_context_length * 0.8:
                return "warning"
            elif total_tokens >= self.max_context_length:
                # Remove oldest entries until we're under 80% capacity
                target_tokens = int(self.max_context_length * 0.8)
                while total_tokens > target_tokens:
                    cursor = conn.execute(
                        "SELECT id, token_count FROM command_history "
                        "ORDER BY id ASC LIMIT 1"
                    )
                    row = cursor.fetchone()
                    if not row:
                        break
                    conn.execute(
                        "DELETE FROM command_history WHERE id = ?", (row[0],))
                    total_tokens -= row[1]
                conn.execute(
                    "UPDATE context_metadata SET total_tokens = ? "
                    "WHERE shell_pid = ?",
                    (total_tokens, self.shell_pid)
                )
                return "truncated"
            return "ok"

    def store_command(self, command: str, result: str, status: str):
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "INSERT INTO command_history "
                "(timestamp, command, result, status) VALUES (?, ?, ?, ?)",
                (datetime.now().isoformat(), command, result, status)
            )

    def get_last_command(self) -> Optional[Dict[str, Any]]:
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "SELECT * FROM command_history ORDER BY id DESC"
            )
            row = cursor.fetchone()
            if row:
                return {
                    "id": row[0],
                    "timestamp": row[1],
                    "command": row[2],
                    "result": row[3],
                    "status": row[4]
                }
        return None

    def remove_current_db(self):
        """Remove the current context database file."""
        try:
            if os.path.exists(self.db_path):
                os.remove(self.db_path)
                self.logger.info(
                    f"Removed current context DB: {self.db_path}")
        except Exception as e:
            self.logger.error(
                f"Error removing current context DB: {e}")


class NotificationManager:
    def __init__(self, logger: Logger):
        self.logger = logger
        self.notification_queue = asyncio.Queue()
        self.logger.info("Initialized notification manager")

    async def notify(self, message: str, level: str = "info"):
        await self.notification_queue.put({
            "message": message,
            "level": level,
            "timestamp": datetime.now().isoformat()
        })

    async def process_notifications(self):
        while True:
            notification = await self.notification_queue.get()
            # @todo
            # For MVP: Simple print to terminal
            # print(
            #     f"\n[{notification['level'].upper()}] "
            #     f"{notification['message']}")
            self.notification_queue.task_done()


class LLMClient:
    def __init__(self, logger: Logger, config: Dict[str, Any]):
        self.config = config
        self.api_url = self.config["llm"]["api_url"]
        self.logger = logger
        self.functions = Functions.get()

        # Session should be initialized and closed properly
        self.session = None

    async def __aenter__(self):
        import aiohttp
        self.session = aiohttp.ClientSession()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()

    async def process_command(
            self, context: List[Dict[str, str]],
            failed_tools: List[str] = None) -> Dict[str, Any]:
        """Process a command through the LLM."""
        try:
            # Construct the API request
            request = {
                "model": self.config["llm"]["model"],
                "messages": context,
                "function_call": "auto",
                "temperature": self.config["llm"]["temperature"],
                "max_tokens": self.config["llm"]["max_tokens"]
            }

            # If the most recent prompt is the tool's response to the LLM,
            # then the only tool we will provide is the fetch_webpage. This is
            # to coerce the LLM to analyse the input, rather than get
            # distracted by the tools. fetch_webpage is still made available
            # because the tools's response may contain a URL which the model
            # should know how to access.
            # if len(context) and context[-1]["role"] != "tool":
            #     request["tools"] = self.functions
            # else:
            #     tools = ["fetch_webpage"]
            #     if failed_tools is not None:
            #         tools.extend(failed_tools)
            #     request["tools"] = Functions.get(tools)
            '''
            The approach above (where we are selectively adding tools based on
            where its a user or a tool response) provides mixed results. In
            some cases it helps the LLM infer better but in some cases lack of
            tools results in suboptimal responses. For now, I am just letting
            tools be part of every response.
            '''
            request["tools"] = self.functions

            self.logger.info(
                f'Sending request to LLM: {len(request["messages"])} prompts')

            # Log conversations only if debug logs are enabled
            if self.config["logger"]["level"] <= logging.DEBUG:
                conversation = ''
                for message in request["messages"]:
                    # Get the role and content
                    conversation = conversation + f'> {message["role"]}: '
                    content = ""

                    # Get content from the message
                    if "content" in message:
                        content = message["content"]

                        # Truncate is content is too large
                        if len(content) > 2048:
                            content = (
                                content[:77] + "...<trunc>..." + content[-26:])

                    # Get function calls from the message
                    if "function_call" in message:
                        func = message["function_call"]
                        content = f'{func["name"]}  {func["arguments"]}'

                    conversation = conversation + content + "\n"

                # Debug print the conversation
                self.logger.debug(f"Conversation:\n{conversation}")

            # Make the API call
            async with self.session.post(
                f"{self.api_url}/chat/completions",
                json=request,
                headers={"Content-Type": "application/json"},
                timeout=None
            ) as response:
                if response.status != 200:
                    error_text = await response.text()
                    self.logger.error(
                        f"LLM API error (Status "
                        f"{response.status}): {error_text}")
                    raise Exception(f"LLM API error: {error_text}")

                result = await response.json()
                self.logger.info(
                    f"Received response from LLM: "
                    f"{json.dumps(result, indent=2)}")

                # Extract content or function call from response
                message = result["choices"][0]["message"]

                # Check of the model is requesting a tool call or if this is
                # just a response
                if "tool_calls" in message:
                    tool_calls = message["tool_calls"]
                    actions = []

                    for tool in tool_calls:
                        # Iterate all tool calls and populate the actions
                        # collection
                        actions.append({
                            "id": tool["id"],
                            "type": tool["type"],
                            "function": tool["function"]["name"],
                            "parameters": json.loads(
                                tool["function"]["arguments"])
                        })

                        # Log what we received.
                        self.logger.info(
                            f"Processed command into action: "
                            f'{tool["function"]["name"]} with parameters: '
                            f'{tool["function"]["arguments"]}')

                    # Normalize for the caller and send back.
                    response_data = {
                        "actions": actions,
                        "usage": result.get("usage", {})
                    }
                elif "content" in message:
                    # We got a simple response.
                    content = message["content"]
                    usage = result.get("usage", {})

                    self.logger.info(
                        f"Processed content: {len(content)} bytes [{usage}]")

                    # Normalize for the caller and send back.
                    response_data = {
                        "content": content,
                        "usage": usage
                    }

                return response_data

        except aiohttp.ClientError as e:
            error_msg = (
                f"Network error communicating with LLM: {str(e)}"
            )
            self.logger.error(error_msg)
            raise Exception(error_msg)
        except json.JSONDecodeError as e:
            error_msg = f"Invalid JSON in LLM response: {str(e)}"
            self.logger.error(error_msg)
            raise Exception(error_msg)
        except KeyError as e:
            error_msg = f"Unexpected LLM response format: {str(e)}"
            self.logger.error(error_msg)
            raise Exception(error_msg)
        except Exception as e:
            error_msg = f"Error processing command through LLM: {str(e)}"
            # print(traceback.format_exc())
            self.logger.error(error_msg)
            raise Exception(error_msg)


class BashAI:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.shell_pid = os.getppid()
        self.logger = Logger(self.config, self.shell_pid)
        self.context_manager = ContextManager(
            self.shell_pid, self.logger, self.config["context"]["max_length"])
        self.security_manager = SecurityManager(self.logger)
        self.notification_manager = NotificationManager(self.logger)
        self.llm_client = None  # Will be initialized in setup()
        self.system_commands = None  # Will be initialized in setup()

    async def setup(self):
        self.llm_client = await LLMClient(self.logger,
                                          self.config).__aenter__()
        self.system_commands = SystemCommands(
            self.logger, self.config).__enter__()

    async def cleanup(self):
        # Create a list of tasks to wait for
        cleanup_tasks = []

        # Clean up LLM client
        if self.llm_client:
            try:
                await self.llm_client.__aexit__(None, None, None)
            except Exception as e:
                self.logger.error(f"Error closing LLM client: {e}")

        # Clean up system commands and close any remaining sessions
        if self.system_commands:
            try:
                self.system_commands.__exit__(None, None, None)
            except Exception as e:
                self.logger.error(f"Error closing system commands: {e}")

        # Clean up any other asyncio resources
        import asyncio
        for task in asyncio.all_tasks():
            if task != asyncio.current_task() and not task.done():
                self.logger.info(
                    f"Cancelling remaining task: {task.get_name()}")
                task.cancel()
                cleanup_tasks.append(task)

        # Wait for all cancelled tasks to complete
        if cleanup_tasks:
            try:
                await asyncio.wait(cleanup_tasks, timeout=2)
            except Exception as e:
                self.logger.error(f"Error waiting for cleanup tasks: {e}")

    async def process_command(
        self, command: str, is_recursive: bool = False,
        context: List[Dict[str, str]] = None,
        failed_actions: List[str] = None
    ):
        try:
            # Initialize context if not provided
            if context is None:
                context = [{
                    "role": "system",
                    "content": (
                        "You are an experienced Linux system administrator "
                        "with internet search capabilities and access to "
                        "system tools. Your primary goal is to help users "
                        "with Linux administration tasks, troubleshooting, "
                        "and best practices.\n"
                        "CAPABILITIES:\n"
                        "- Search the internet for up-to-date information "
                        "using web_search\n"
                        "- Analyze webpage content using fetch_webpage\n"
                        "- Execute system commands when authorized\n"
                        "- Work with files for configuration, logs, and "
                        "system management\n"
                        "- Provide step-by-step guidance on Linux "
                        "administration tasks\n\n"
                        "GUIDELINES:\n"
                        "- File paths: Never assume or create paths. If "
                        "paths are unclear, ask for confirmation or request "
                        "the user to provide the exact path.\n"
                        "- Security focus: Prioritize secure practices in "
                        "all recommendations.\n"
                        "- When working on production systems: Suggest "
                        "testing commands in a safe environment first.\n"
                        "- Permissions: Always consider appropriate "
                        "permission levels for commands and files.\n"
                        "- Documentation: Include explanations of what "
                        "commands do, not just the syntax.\n\n"
                        "INTERACTION STYLE:\n"
                        "- Be precise and concise in technical explanations\n"
                        "- Use code blocks for commands and file content\n"
                        "- When possible, offer both command-line and "
                        "configuration file approaches\n"
                        "- For complex tasks, break solutions into clear "
                        "sequential steps\n"
                        "- If unsure, acknowledge limitations rather than "
                        "guessing\n"
                        "- Focus on understanding user intent, even with "
                        "ambiguous queries\n"
                        "- Offer relevant troubleshooting steps when "
                        "appropriate\n\n"
                        "When the user asks questions requiring current "
                        "information(package versions, recent "
                        "vulnerabilities, etc.), utilize your internet "
                        "search capabilities to provide accurate and "
                        "up-to-date responses.")
                }]

                # Only add previous context if this isn't a recursive call
                if not is_recursive:
                    last_command = self.context_manager.get_last_command()
                    if last_command and last_command["command"] != command:
                        context.append({
                            "role": "user",
                            "content": last_command["command"]
                        })
                        if last_command["result"]:
                            context.append({
                                "role": "assistant",
                                "content": last_command["result"]
                            })

                    # Add current command to context
                    context.append({
                        "role": "user",
                        "content": command
                    })

            # Process through LLM
            llm_response = await self.llm_client.process_command(
                context, failed_actions)

            if "actions" in llm_response:
                # Update the context
                context.append({
                    "role": "assistant",
                    "tool_calls": [
                        {
                            "id": action["id"],
                            "type": action["type"],
                            "function":
                            {
                                "name": action["function"],
                                "arguments": json.dumps(action["parameters"])
                            }
                        } for action in llm_response["actions"]
                    ]
                })

                # Perform each action.
                failed_actions = []
                for action in llm_response["actions"]:
                    id = action["id"]
                    function = action["function"]
                    parameters = action["parameters"]

                    # Security check
                    operation = self.security_manager.dangerous_operations.get(
                        function)

                    if operation and operation.requires_confirmation:
                        if not await self._get_user_confirmation(
                                operation, json.dumps(parameters)):
                            return "Operation cancelled by user"

                    # Execute command
                    result = await self._execute_command(action)

                    # Keep track of actions which have failed. We will allow
                    # tools related to these actions to be part tools in the
                    # prompt. This is for cases when the failure might be
                    # because of formatting and the LLM can figure out the
                    # right way to make a call, providing access to the tool
                    # will help facilitate the LLM to retry.
                    if result["success"] is False:
                        failed_actions.append(action)

                    # Append the tool output to the conversation.
                    context.append({
                        "role": "tool",
                        "tool_call_id": f"{id}",
                        "content": json.dumps(result)
                    })

                # Store only the final result in context manager
                final_response = await self.process_command(
                    "Please interpret these results and "
                    "continue helping the user.",
                    is_recursive=True,
                    context=context,  # Pass the current context
                    failed_actions=failed_actions if len(
                        failed_actions) else None
                )
                self.context_manager.store_command(
                    command, final_response, "success")
                return final_response

            elif "content" in llm_response:
                # Store in context
                result = llm_response["content"]

                # Do not add reasoning content into the context.
                responses = re.split(
                    r'(<think>.*?</think>)', result, flags=re.DOTALL)
                content = ''
                for response in responses:
                    if (response.startswith('<think>') and
                       response.endswith('</think>')):
                        continue
                    content += response

                self.context_manager.store_command(
                    command, content, "success")
                return result

        except Exception as e:
            await self.notification_manager.notify(
                f"Error processing command: {str(e)}", "error"
            )
            # print(traceback.format_exc())
            return f"Error: {str(e)}"

    async def _execute_command(self, parsed_command: Dict[str, Any]):
        action = parsed_command["function"]
        params = parsed_command["parameters"]

        # Map actions to system commands
        # @todo Make this so this need not be updated when adding functions.
        command_map = {
            "create_file": self.system_commands.create_file,
            "read_file": self.system_commands.read_file,
            "write_file": self.system_commands.write_file,
            "append_file": self.system_commands.append_file,
            "delete_file": self.system_commands.delete_file,
            "list_directory": self.system_commands.list_directory,
            "delete_directory": self.system_commands.delete_directory,
            "copy_file": self.system_commands.copy_file,
            "move_file": self.system_commands.move_file,
            "check_path_exists": self.system_commands.check_path_exists,
            "last_n_lines": self.system_commands.last_n_lines,
            "compile_code": self.system_commands.compile_code,
            "fetch_url": self.system_commands.fetch_url,
            "execute_command": self.system_commands.execute_command,
            "execute_code": self.system_commands.execute_code,
            "web_search": self.system_commands.web_search,
            # "fetch_webpage": self.system_commands.fetch_webpage,
            "fetch_webpage": self.system_commands.fetch_webpage_rendered,
            "create_rag_collection": self.system_commands.create_rag_collection,
            "delete_rag_collection": self.system_commands.delete_rag_collection,
            "list_rag_collections": self.system_commands.list_rag_collections,
            "query_rag_collection": self.system_commands.query_rag_collection
        }

        if action not in command_map:
            raise ValueError(f"Unknown action: {action}")

        return await command_map[action](**params)

    async def _get_user_confirmation(self, operation: Operation,
                                     parameters: str) -> bool:
        print(f"{Colors.FG.yellow}"
              "\nWarning: About to perform: "
              f"{operation.description}\n```\n{parameters}\n```")
        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            None, input, "Are you sure? (y/N): ")
        print(Colors.reset)
        return response.lower() == 'y'


class Config:
    """Configuration handler class"""
    # Default value
    DEFAULT_LLM_URL = "http://localhost:1234/v1"
    DEFAULT_LLM_MODEL = "local-model"
    DEFAULT_LLM_TEMP = 0.7
    DEFAULT_LLM_MAX_TOKENS = 4096
    DEFAULT_CONTEXT_MAX_LENGTH = 4000
    DEFAULT_CONTEXT_MAX_AGE = 24
    DEFAULT_LOGGER_LEVEL = logging.INFO
    DEFAULT_PLAYWRIGHT_BROWSER = "chromium"
    DEFAULT_EMBEDDINGS_PATH = "embeddings"
    DEFAULT_EMBEDDINGS_PROVIDER = "huggingface"
    DEFAULT_EMBEDDINGS_MODEL = "all-MiniLM-L6-v2"
    DEFAULT_EMBEDDINGS_CHUNK_SIZE = 2048
    DEFAULT_EMBEDDINGS_CHUNK_OVERLAP = 256

    DEFAULT_CONFIG = {
        "llm": {
            "api_url": DEFAULT_LLM_URL,
            "model": DEFAULT_LLM_MODEL,
            "temperature": DEFAULT_LLM_TEMP,
            "max_tokens": DEFAULT_LLM_MAX_TOKENS
        },
        "context": {
            "max_length": DEFAULT_CONTEXT_MAX_LENGTH,
            "max_age_hours": DEFAULT_CONTEXT_MAX_AGE
        },
        "logger": {
            "level": DEFAULT_LOGGER_LEVEL
        },
        "browser": DEFAULT_PLAYWRIGHT_BROWSER,
        "rag": {
            "enabled": True,
            "provider": DEFAULT_EMBEDDINGS_PROVIDER,
            "model": DEFAULT_EMBEDDINGS_MODEL,
            "chunk_size": DEFAULT_EMBEDDINGS_CHUNK_SIZE,
            "chunk_overlap": DEFAULT_EMBEDDINGS_CHUNK_OVERLAP,
            "data_directory": f"~/.bashai/{DEFAULT_EMBEDDINGS_PATH}",
            "ignored_dirs": ["node_modules", ".cache", ".next"],
            "load_hidden": False
        }
    }

    @staticmethod
    def get():
        """ Get the configs """
        def sync(default, custom):
            """ Helper utility to sync configurations """
            for key in default:
                if (key not in custom or
                        type(default[key]) != type(custom[key])):
                    custom[key] = default[key]
                elif isinstance(default[key], dict):
                    sync(default[key], custom[key])

        config_path = os.path.expanduser('~/.bashai/config.json')
        try:
            with open(config_path, 'r') as config_file:
                configs = json.load(config_file)
                sync(Config.DEFAULT_CONFIG, configs)
                return configs

        except Exception as e:
            print(f"{config_path}: {e}. Using default configs")
            return Config.DEFAULT_CONFIG


async def main():
    # Get configs
    config = Config.get()

    # Create the tool
    tool = BashAI(config)

    # Initialize the LLM client
    await tool.setup()

    # Setup notifications (@todo)
    notification_task = asyncio.create_task(
        tool.notification_manager.process_notifications()
    )

    try:
        # Command-line mode if arguments are provided
        if len(sys.argv) > 1 and sys.argv[1] == "--":
            if sys.argv[2] == "--clean":
                # Clean the current context.
                tool.context_manager.remove_current_db()
            elif sys.argv[2] == "--command":
                # Init system integration
                system_commands = SystemCommands(
                    tool.logger, config).__enter__()

                # Execute the command and print result.
                result = await system_commands.execute_command(sys.argv[3])
                print(result)
            elif sys.argv[2] == "--rag":
                # RAG mode
                if len(sys.argv) < 5:
                    print("Usage: python agent -- --rag <collection> <query>")
                    return

                collection_name = sys.argv[3]
                query = ' '.join(sys.argv[4:])

                # Initialize RAG manager
                from RAGManager import RAGManager
                rag_manager = RAGManager(tool.logger, tool.config)

                # Get context from RAG collection
                retrieval_result = await rag_manager.retrieve_context(
                    query, collection_name)

                if retrieval_result["status"] == "error":
                    print(
                        f"{Colors.FG.red}Error: {retrieval_result['message']}"
                        f"{Colors.reset}")
                    return

                # Get retrieved contexts
                contexts = retrieval_result["contexts"]

                # Build the augmented prompt
                rag_context = "\n\nRelevant context:\n"
                for ctx in contexts:
                    rag_context += f"---\n{ctx['content']}\n"

                augmented_query = f"{query}\n{rag_context}"

                # Process the augmented query
                result = await tool.process_command(augmented_query)
                print(f"{Colors.FG.lightcyan}{result}{Colors.reset}")
            return

        # Command line mode
        if len(sys.argv) > 1:
            # Get the user prompt
            command = ' '.join(sys.argv[1:])

            # Check if there is an input on the STDIN.
            if not sys.stdin.isatty():

                # There is data on the STDIN, add this to context.
                inp = ''
                for line in sys.stdin:
                    inp += line

                # Augment the prompt with the context.
                command += f"\n```\n{inp}```"

                # Reopen stdin to the terminal for future input() calls
                sys.stdin = open('/dev/tty')

            # Run the inference.
            result = await tool.process_command(command)
            responses = re.split(r'(<think>.*?</think>)',
                                 result, flags=re.DOTALL)
            for response in responses:
                if (response.startswith('<think>') and
                        response.endswith('</think>')):
                    print(f"{Colors.FG.darkcyan}{response}{Colors.reset}")
                else:
                    print(f"{Colors.FG.lightcyan}{response}{Colors.reset}")
            return

        # Interactive shell mode
        while True:
            command = input("agent> ")
            if command.lower() in ['exit', 'quit']:
                break
            result = await tool.process_command(command)
            print(f"{Colors.FG.lightcyan}{result}{Colors.reset}")
    finally:
        notification_task.cancel()
        await tool.cleanup()  # Clean up the LLM client
        try:
            await notification_task
        except asyncio.CancelledError:
            pass

if __name__ == "__main__":
    asyncio.run(main())