feat: Add cli for interacting with Axiom agent

cli.py

@@ -0,0 +1,131 @@
+import asyncio
+from typing import List, Dict, Optional
+
+from src.axiom.agent import AxiomAgent
+from src.axiom.config import settings, load_mcp_servers_from_config 
+from agents.mcp import MCPServer 
+import logging
+
+# --- Configure Logging ---
+httpx_logger = logging.getLogger("httpx")
+httpx_logger.setLevel(logging.WARNING) 
+
+# Use Rich for pretty output
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.rule import Rule
+from rich.text import Text
+
+console = Console()
+
+async def get_user_input(prompt_text: str) -> str:
+    """Gets user input asynchronously with a styled prompt."""
+    # Render the prompt text using Rich
+    prompt_markup = Text.from_markup(prompt_text, style="bold blue")
+    return await asyncio.to_thread(console.input, prompt_markup)
+
+
+async def main():
+    console.print(Rule("[bold blue] Welcome to Axiom CLI [/bold blue]", style="blue"))
+    console.print("[dim]Type 'quit' or 'exit' to end the chat.[/dim]")
+    console.print("") # Blank line
+
+    loaded_mcp_servers: List[MCPServer] = []
+    started_mcp_servers: List[MCPServer] = []
+    agent: Optional[AxiomAgent] = None
+    chat_history: List[Dict[str, str]] = []
+
+    try:
+        # 1. Load MCP Servers
+        console.print("[bold]Loading MCP configurations...[/bold]")
+        try:
+            loaded_mcp_servers = load_mcp_servers_from_config()
+            if loaded_mcp_servers:
+                 console.print(f"[green]Loaded {len(loaded_mcp_servers)} server(s) config.[/green] Attempting to start...")
+            else:
+                 console.print("[yellow]No MCP server configurations found or loaded.[/yellow]")
+
+        except Exception as e:
+            console.print(f"[bold red]Error loading MCP config:[/bold red] [red]{e}[/red]")
+
+        # 2. Start MCP Servers
+        if loaded_mcp_servers:
+            for server in loaded_mcp_servers:
+                try:
+                    await server.__aenter__()
+                    started_mcp_servers.append(server)
+                    console.print(f"  [green]Started:[/green] {server.name}")
+                except Exception as e:
+                    console.print(f"  [yellow]Failed to start:[/yellow] {server.name} - [yellow]{e}[/yellow]")
+
+        if loaded_mcp_servers and not started_mcp_servers:
+             console.print("[yellow]Warning: All configured MCP servers failed to start. Agent will operate without MCP tools.[/yellow]")
+        elif started_mcp_servers:
+             console.print(f"[green]Successfully started {len(started_mcp_servers)} MCP server(s).[/green]")
+
+        # 3. Initialize the Agent
+        console.print("[bold]Initializing Agent...[/bold]")
+        try:
+            agent = AxiomAgent(mcp_servers=started_mcp_servers)
+            console.print(f"[bold green]{settings.AGENT_NAME} is ready![/bold green]")
+        except Exception as e:
+            console.print(f"[bold red]Fatal Error: Could not initialize Agent:[/bold red] [red]{e}[/red]")
+            raise
+
+        console.print(Rule(style="blue"))
+
+        # 4. Main chat loop
+        while True:
+            # Get user input with styled prompt
+            user_input = await get_user_input("You: ")
+
+            if user_input.lower() in ['quit', 'exit']:
+                break
+
+            chat_history.append({"role": "user", "content": user_input})
+
+            console.print("[bold green]Axiom:[/bold green]") # Print Agent prefix before response
+
+            full_response = ""
+            # Use Rich Status for the thinking indicator while streaming
+            with console.status("[bold green]Thinking...[/bold green]", spinner="dots", speed=0.5) as status:
+                 try:
+                    response_generator = agent.stream_agent(chat_history)
+                    async for token in response_generator:
+                        full_response += token
+                    status.stop()
+
+                 except Exception as e:
+                     status.stop()
+                     console.print(f"\n[bold red]Error during response:[/bold red] [red]{e}[/red]", style="red")
+                     full_response = f"[Error: {e}]"
+
+
+            # Render the full response using Markdown
+            if full_response:
+                 console.print(Markdown(full_response)) # Renders code blocks, lists, etc.
+            else:
+                 console.print("[dim](No response generated)[/dim]")
+
+            chat_history.append({"role": "assistant", "content": full_response})
+            console.print(Rule(style="blue")) # Print a rule after each turn
+
+    except Exception as e:
+        # Catch any unexpected errors from Agent init or loop setup
+        console.print(f"\n[bold red]An unhandled error occurred:[/bold red] [red]{e}[/red]", style="red", highlight=True)
+
+    finally:
+        # 5. Cleanup: Stop MCP Servers
+        if started_mcp_servers:
+            console.print(Rule("[dim]Stopping MCP servers[/dim]", style="blue"))
+            for server in started_mcp_servers:
+                try:
+                    await server.__aexit__(None, None, None)
+                    console.print(f"  [dim]Stopped:[/dim] {server.name}")
+                except Exception as e:
+                    console.print(f"  [red]Error stopping:[/red] {server.name} - [red]{e}[/red]")
+
+        console.print(Rule("[bold blue] Chat ended. Goodbye! [/bold blue]", style="blue"))
+
+if __name__ == "__main__":
+    asyncio.run(main())

pyproject.toml

@@ -7,7 +7,9 @@ requires-python = ">=3.11"
 dependencies = [
     "chainlit>=2.5.5",
     "openai-agents>=0.0.11",
+    "prompt-toolkit>=3.0.51",
     "python-dotenv>=1.1.0",
+    "rich>=14.0.0",
 ]
 
 [tool.hatch.build.targets.wheel]

src/axiom/agent.py

@@ -93,6 +93,7 @@ class AxiomAgent:
             result = Runner.run_streamed(
                 starting_agent=self.agent,
                 input=chat_history,
+                max_turns=20,
                 run_config=config
             )
             async for event in result.stream_events():

src/axiom/config.py

@@ -80,8 +80,6 @@ def load_mcp_servers_from_config(config_path: Path = settings.MCP_CONFIG_PATH) -
         logger.error(f"MCP configuration file not found: {config_path}")
         raise FileNotFoundError(f"MCP configuration file not found: {config_path}")
 
-    logger.info(f"Loading MCP servers from: {config_path}")
-    
     # Allow json.JSONDecodeError to propagate if file is invalid JSON
     with open(config_path, 'r', encoding='utf-8') as f:
         data = json.load(f)
@@ -104,7 +102,6 @@ def load_mcp_servers_from_config(config_path: Path = settings.MCP_CONFIG_PATH) -
                     }
                 )
             servers.append(server_instance)
-            logger.info(f"Prepared MCP Server: {name}")
 
         except (ValidationError, Exception) as e:
             logger.warning(f"Skipping MCP server '{name}' due to configuration error: {e}")

uv.lock

@@ -169,14 +169,18 @@ source = { editable = "." }
 dependencies = [
     { name = "chainlit" },
     { name = "openai-agents" },
+    { name = "prompt-toolkit" },
     { name = "python-dotenv" },
+    { name = "rich" },
 ]
 
 [package.metadata]
 requires-dist = [
     { name = "chainlit", specifier = ">=2.5.5" },
     { name = "openai-agents", specifier = ">=0.0.11" },
+    { name = "prompt-toolkit", specifier = ">=3.0.51" },
     { name = "python-dotenv", specifier = ">=1.1.0" },
+    { name = "rich", specifier = ">=14.0.0" },
 ]
 
 [[package]]
@@ -707,6 +711,18 @@ dependencies = [
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/7e/c1/7bd34ad0ae6cfd99512f8a40b28b9624c3b1f4e1d40c9038eabc2f870b15/literalai-0.1.201.tar.gz", hash = "sha256:29e4ccadd9d68bfea319a7f0b4fc32611b081990d9195f98e5e97a14d24d3713", size = 67832 }
 
+[[package]]
+name = "markdown-it-py"
+version = "3.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "mdurl" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/38/71/3b932df36c1a044d397a1f92d1cf91ee0a503d91e470cbd670aa66b07ed0/markdown-it-py-3.0.0.tar.gz", hash = "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb", size = 74596 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/42/d7/1ec15b46af6af88f19b8e5ffea08fa375d433c998b8a7639e76935c14f1f/markdown_it_py-3.0.0-py3-none-any.whl", hash = "sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1", size = 87528 },
+]
+
 [[package]]
 name = "markupsafe"
 version = "3.0.2"
@@ -786,6 +802,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/10/30/20a7f33b0b884a9d14dd3aa94ff1ac9da1479fe2ad66dd9e2736075d2506/mcp-1.6.0-py3-none-any.whl", hash = "sha256:7bd24c6ea042dbec44c754f100984d186620d8b841ec30f1b19eda9b93a634d0", size = 76077 },
 ]
 
+[[package]]
+name = "mdurl"
+version = "0.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979 },
+]
+
 [[package]]
 name = "monotonic"
 version = "1.6"
@@ -1567,6 +1592,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/54/e2/c158366e621562ef224f132e75c1d1c1fce6b078a19f7d8060451a12d4b9/posthog-3.25.0-py2.py3-none-any.whl", hash = "sha256:85db78c13d1ecb11aed06fad53759c4e8fb3633442c2f3d0336bc0ce8a585d30", size = 89115 },
 ]
 
+[[package]]
+name = "prompt-toolkit"
+version = "3.0.51"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "wcwidth" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/bb/6e/9d084c929dfe9e3bfe0c6a47e31f78a25c54627d64a66e884a8bf5474f1c/prompt_toolkit-3.0.51.tar.gz", hash = "sha256:931a162e3b27fc90c86f1b48bb1fb2c528c2761475e57c9c06de13311c7b54ed", size = 428940 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ce/4f/5249960887b1fbe561d9ff265496d170b55a735b76724f10ef19f9e40716/prompt_toolkit-3.0.51-py3-none-any.whl", hash = "sha256:52742911fde84e2d423e2f9a4cf1de7d7ac4e51958f648d9540e0fb8db077b07", size = 387810 },
+]
+
 [[package]]
 name = "propcache"
 version = "0.3.1"
@@ -1748,6 +1785,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b6/5f/d6d641b490fd3ec2c4c13b4244d68deea3a1b970a97be64f34fb5504ff72/pydantic_settings-2.9.1-py3-none-any.whl", hash = "sha256:59b4f431b1defb26fe620c71a7d3968a710d719f5f4cdbbdb7926edeb770f6ef", size = 44356 },
 ]
 
+[[package]]
+name = "pygments"
+version = "2.19.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/7c/2d/c3338d48ea6cc0feb8446d8e6937e1408088a72a39937982cc6111d17f84/pygments-2.19.1.tar.gz", hash = "sha256:61c16d2a8576dc0649d9f39e089b5f02bcd27fba10d8fb4dcc28173f7a45151f", size = 4968581 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8a/0b/9fcc47d19c48b59121088dd6da2488a49d5f72dacf8262e2790a1d2c7d15/pygments-2.19.1-py3-none-any.whl", hash = "sha256:9ea1544ad55cecf4b8242fab6dd35a93bbce657034b0611ee383099054ab6d8c", size = 1225293 },
+]
+
 [[package]]
 name = "pyjwt"
 version = "2.10.1"
@@ -1915,6 +1961,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f9/9b/335f9764261e915ed497fcdeb11df5dfd6f7bf257d4a6a2a686d80da4d54/requests-2.32.3-py3-none-any.whl", hash = "sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6", size = 64928 },
 ]
 
+[[package]]
+name = "rich"
+version = "14.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "markdown-it-py" },
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a1/53/830aa4c3066a8ab0ae9a9955976fb770fe9c6102117c8ec4ab3ea62d89e8/rich-14.0.0.tar.gz", hash = "sha256:82f1bc23a6a21ebca4ae0c45af9bdbc492ed20231dcb63f297d6d1021a9d5725", size = 224078 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/0d/9b/63f4c7ebc259242c89b3acafdb37b41d1185c07ff0011164674e9076b491/rich-14.0.0-py3-none-any.whl", hash = "sha256:1c9491e1951aac09caffd42f448ee3d04e58923ffe14993f6e83068dc395d7e0", size = 243229 },
+]
+
 [[package]]
 name = "simple-websocket"
 version = "1.1.0"
@@ -2252,6 +2311,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/3f/82/45dddf4f5bf8b73ba27382cebb2bb3c0ee922c7ef77d936b86276aa39dca/watchfiles-0.20.0-cp37-abi3-win_arm64.whl", hash = "sha256:b17d4176c49d207865630da5b59a91779468dd3e08692fe943064da260de2c7c", size = 265344 },
 ]
 
+[[package]]
+name = "wcwidth"
+version = "0.2.13"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/6c/63/53559446a878410fc5a5974feb13d31d78d752eb18aeba59c7fef1af7598/wcwidth-0.2.13.tar.gz", hash = "sha256:72ea0c06399eb286d978fdedb6923a9eb47e1c486ce63e9b4e64fc18303972b5", size = 101301 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fd/84/fd2ba7aafacbad3c4201d395674fc6348826569da3c0937e75505ead3528/wcwidth-0.2.13-py2.py3-none-any.whl", hash = "sha256:3da69048e4540d84af32131829ff948f1e022c1c6bdb8d6102117aac784f6859", size = 34166 },
+]
+
 [[package]]
 name = "wrapt"
 version = "1.17.2"