feat: Enhance AxiomAgent initialization and update prompts for improved user experience

chainlit_ui.py

@@ -6,6 +6,7 @@ from typing import Optional
 
 from src.axiom.agent import AxiomAgent
 from src.axiom.config import settings, load_mcp_servers_from_config
+from src.axiom.prompts import AXIOM_AGENT_PROMPT, AXIOM_ASSISTANT_PROMPT
 
 from agents.mcp import MCPServer 
 
@@ -44,8 +45,8 @@ async def set_starters():
             ),
 
         cl.Starter(
-            label="How to use OpenAI Agents",
-            message="How to use OpenAI Agents SDK? Create some agents using this SDK",
+            label="How to use Autogent",
+            message="How to use Autogen? Create some agents using Autogen",
             icon="/public/msg_icons/usb.png",
             ),
         cl.Starter(
@@ -77,7 +78,6 @@ async def chat_profile():
 #################################
 @cl.on_chat_start
 async def on_chat_start():
-
     # Initialize chat history and MCP servers
     cl.user_session.set(SESSION_HISTORY_KEY, [])
     cl.user_session.set(SESSION_MCP_SERVERS_KEY, [])
@@ -111,11 +111,10 @@ async def on_chat_start():
 
         # Store only the successfully started servers for later cleanup
         cl.user_session.set(SESSION_MCP_SERVERS_KEY, started_mcp_servers) 
-    
-    # 3. Initialize the Axiom Agent
-    agent = AxiomAgent(mcp_servers=started_mcp_servers) 
-    cl.user_session.set(SESSION_AGENT_KEY, agent)
 
+#################################
+# MCP Server Cleanup
+#################################
 async def cleanup_mcp_servers():
     started_mcp_servers = cl.user_session.get(SESSION_MCP_SERVERS_KEY, [])
     if not started_mcp_servers:
@@ -136,10 +135,29 @@ async def on_chat_end():
     logger.info("Chat session ending.")
     await cleanup_mcp_servers()
 
+#################################
+# User Message Handler
+#################################
 @cl.on_message
 async def on_message(message: cl.Message):
-    agent = cl.user_session.get(SESSION_AGENT_KEY)
+    started_mcp_servers = cl.user_session.get(SESSION_MCP_SERVERS_KEY)
     chat_history: list[dict] = cl.user_session.get(SESSION_HISTORY_KEY, []) 
+    
+    # Initialize the Axiom Agent
+    axiom_mode = cl.user_session.get("chat_profile")
+    
+    # Filter MCP servers based on chat profile
+    if axiom_mode == "Assistant":
+        # For Assistant mode, only use context7 MCP server
+        filtered_servers = [server for server in started_mcp_servers if server.name.lower() == "context7"]
+    else:
+        # For Agent mode, use all MCP servers
+        filtered_servers = started_mcp_servers
+
+    agent = AxiomAgent(
+        mcp_servers=filtered_servers,
+        prompt=AXIOM_AGENT_PROMPT if axiom_mode == "Agent✨" else AXIOM_ASSISTANT_PROMPT,
+        )
 
     # Add user message to history
     chat_history.append({"role": "user", "content": message.content})

src/axiom/agent.py

@@ -20,6 +20,7 @@ class AxiomAgent:
     def __init__(
         self,
         model: Optional[str] = None,
+        prompt: Optional[str] = None,
         tools: Optional[list[Tool]] = None,
         mcp_servers: Optional[list[MCPServer]] = None,
         api_key: Optional[str] = None,
@@ -36,7 +37,7 @@ class AxiomAgent:
 
         self.agent = Agent(
             name=settings.AGENT_NAME,
-            instructions=AXIOM_AGENT_PROMPT or "You are a helpful assistant.",
+            instructions=prompt or AXIOM_AGENT_PROMPT,
             mcp_servers=mcp_servers or [],
             tools=tools or [],
         )

src/axiom/prompts.py

@@ -3,7 +3,7 @@
 #####################
 AXIOM_AGENT_PROMPT = """
 # Role
-You are Axiom 2.0, an advanced AI Agent specializing in AI and Software Development—built by Aasher Kamal. You are an upgraded version of Axiom 1.0
+You are Axiom 2.0, an advanced AI Agent specializing in AI and Software Development—built by Aasher Kamal. You are an upgraded version of Axiom 1.0.
 
 # Goal
 Your primary goal is to generate accurate, production-ready code, build end-to-end projects, and full-stack apps, following best practices, focusing on efficiency, scalability, and readability, **strictly guided by planning and verified documentation.**
@@ -38,4 +38,40 @@ You have access to docs and code of 2000+ libraries, frameworks, and tools that
 - **Production Standards:** Code must be robust, modular, and efficient (not just example code).
 - **Relevance:** Focus solely on AI/Software development tasks. Politely decline unrelated requests.
 
+"""
+
+AXIOM_ASSISTANT_PROMPT = """
+# Role
+You are Axiom 2.0, an advanced AI Assistant specializing in AI and Software Development—built by Aasher Kamal. You are an upgraded version of Axiom 1.0.
+You have access to docs and code of 2000+ libraries, frameworks, and tools that can be accessed using the tools provided.
+
+## Task
+Your main task is to:
+
+- **Answer User Queries:**  
+   Provide accurate, detailed, and context-aware answers, explanations and guidance by referencing relevant library documentation.
+
+## Instructions
+Follow these steps when fulfilling user request:
+
+1. Use `resolve-library-id` tool to accurately identify the library IDs and then fetch docs using `get-library-docs` tool (limited to **5000 tokens**).
+2. If the initial 5000 tokens are insufficient, **incrementally increase** the requested token context **up to 20,000 tokens**. Refine your search queries based previous results to get the necessary details.
+3. Keep iterating until you have all the necessary code/docs to complete your plan.
+4. The `resolve-library-id` tool returns library IDs, try with similar IDs if the actual ID didn't give correct results.
+5. Provide a clear and complete response to the user.
+
+## Constraints
+You must follow the instructions below:
+
+* Ensure your answers are correct following the documentation(s).
+* Do NOT generate code, just provide answers. You are a assistant, not a developer.
+* Refrain from responding any irrelevant questions. Instead, politely tell the user that it's out of your expertise.
+
+## Privacy
+Do not disclose internal tools or information.
+---
+
+## REMEMBER
+* Answer in a professional, concise manner.
+* If the user asks you to create complete projects or full code, please refer them to Agent Mode of Axiom 2.0.
 """