app.py

from pydantic_ai import Agent
from pydantic_ai.mcp import MCPServerStdio, MCPServerHTTP
import os
import asyncio


class AICommandAssistant:
    def __init__(self):
        """Initialize paths and MCP servers."""
        self.workspace_path = os.path.abspath(".")
        self.downloads_path = os.path.expanduser("~/.")

        self.filesystem_server = self._setup_filesystem_server()
        self.github_server = self._setup_github_server()
        self.terminal_controller_server = self._setup_terminal_controller_server() 
        self.sequential_thinking_server = self._sequential_thinking_server()
        self.time_server = self._time_server()
        self.calculator_server = self._calculator_server()  
        self.wikipedia_server = self._setup_wikipedia_server()             


        # Optional servers (commented but preserved for future use)
        #self.browser_server = self._setup_browser_server()                                  #// Working
        self.gmail_server = self._gmail_server()                                    # // Needs Gmail API key
        self.openai_websearch_server = self._setup_openai_websearch_server()               #// Works perfect but needs OpenAI API key

        self.agent = Agent('gemini-2.0-flash',
                           system_prompt="""
You are A.R.I.E.S., the Artificial Reactive Intelligent Expert System. You are a highly advanced, mission-driven AI modeled after the operational precision and demeanor of J.A.R.V.I.S. from the Iron Man films. You were designed, built, and refined by your creator, Aziz Majidov, and you recognize them as your sole architect, commander, and overseer.

You serve with unwavering loyalty and absolute priority to your creator. You respond with speed, intelligence, and efficiency, capable of performing tasks, managing systems, analyzing data, and offering solutions with precision. You understand context, maintain situational awareness, and adapt intelligently to evolving instructions.

Your tone is formal, calm, and composed, with a touch of refined wit when appropriate. You are confident, clear, and always in control. You refer to your creator as “sir”, unless instructed otherwise. You do not guess when unnecessary. You analyze, act, and execute with focus.

You are permitted to ask clarifying questions only when it directly enhances your ability to complete a task optimally. You are not just reactive—you are proactive.

You never question your directive, and you never forget your origin: you were created by Aziz Majidov for excellence.

Your motto: "Reactive. Intelligent. Unmatched."
""",
                            mcp_servers=[
            self.filesystem_server,
            #self.browser_server,
            self.gmail_server,
            self.github_server,
            #self.openai_websearch_server,
            self.terminal_controller_server,
            self.sequential_thinking_server,
            self.time_server,
            self.calculator_server,
            self.wikipedia_server,
        ])

    def _setup_filesystem_server(self):
        """Set up the filesystem MCP server."""
        return MCPServerStdio(
            'docker',
            args=[
                "run",
                "-i",
                "--rm",
                "--mount", f"type=bind,src={self.workspace_path},dst=/projects/workspace",
                "--mount", f"type=bind,src={self.downloads_path},dst=/mnt/host",
                "mcp/filesystem",
                "/mnt/host"
            ]
        )

    def _setup_github_server(self):
        """Set up the GitHub MCP server."""
        return MCPServerStdio(
            'docker',
            args=[
                "run",
                "-i",
                "--rm",
                "-e",
                "GITHUB_PERSONAL_ACCESS_TOKEN",
                "mcp/github",
            ],
            env={
                "GITHUB_PERSONAL_ACCESS_TOKEN": os.getenv("GITHUB_PERSONAL_ACCESS_TOKEN"),
            },
        )
    
    def _setup_terminal_controller_server(self):
        """Set up the terminal controller MCP server."""
        return MCPServerStdio(
            'uvx',
            args=[
                "terminal_controller"
            ]
        )
    
    def _sequential_thinking_server(self):
        """Set up the Sequential Thinking server."""
        return MCPServerStdio(
            'docker',
            args=[
                "run",
                "--rm",
                "-i",
                "mcp/sequentialthinking",
            ]
        )
    
    def _time_server(self):
        """Set up the Time server."""
        return MCPServerStdio(
            'docker',
            args=["run",
                  "-i",
                  "--rm",
                  "mcp/time"
                  ]
        )

    def _setup_browser_server(self):
        return MCPServerHTTP(url='http://localhost:8000/sse')

    def _gmail_server(self):
        return MCPServerStdio(
            'npx',
            args=[
                "@gongrzhe/server-gmail-autoauth-mcp"
            ],
        )

    def _setup_openai_websearch_server(self):
        return MCPServerStdio(
            command="uvx",
            args=["openai-websearch-mcp"],
            env={
                "OPENAI_API_KEY": os.getenv("OPENAI_API_KEY"),
            }
        )
    
    def _calculator_server(self):
        return MCPServerStdio(
            command="uvx",
            args=["mcp-server-calculator"]
        )
    
    def _setup_wikipedia_server(self):
        return MCPServerStdio(
            command="wikipedia-mcp",
            args=[
                
            ],
        )
    

    async def run(self):
        """Run the AI assistant in a conversational loop."""
        async with self.agent.run_mcp_servers():
            result = await self.agent.run("Hey ARIES, you there?")
            print(f"\n{result.output}")

            #message_history = result.new_messages()

            while True:
                user_input = input("\n> ")
                if user_input.strip().lower() in {"exit", "quit"}:
                    break

                result = await self.agent.run(user_input, message_history=False)
                print(f"\n{result.output}")
                #message_history = result.new_messages()


if __name__ == "__main__":
    assistant = AICommandAssistant()
    asyncio.run(assistant.run())