feat: create Chainlit app for using agent.

chainlit_ui.py

@@ -0,0 +1,54 @@
+import chainlit as cl
+
+from agents.mcp import MCPServerStdio
+from src.axiom.agent import AxiomAgent
+
+@cl.on_chat_start 
+async def on_chat_start():
+
+    server1 = MCPServerStdio(
+        name="Documentation MCP",
+        params={
+            "command": "npx",
+            "args": ['-y', '@upstash/context7-mcp@latest']
+        },
+    )
+    server2 = MCPServerStdio(
+        name="Sequential Thinking MCP",
+        params={
+            "command": "npx",
+            "args": ['-y', '@modelcontextprotocol/server-sequential-thinking']
+        },
+    )
+    try:
+         # Manually enter the async context
+        await server1.__aenter__()
+        await server2.__aenter__() 
+    except Exception as e:
+        cl.Message(content=f"Failed to start MCP Server: {e}").send()
+        return
+
+    agent = AxiomAgent(mcp_servers=[server1, server2])
+    cl.user_session.set("axiom_agent", agent)
+    cl.user_session.set("chat_history", [])
+
+@cl.on_message
+async def on_message(message: cl.Message):
+    agent = cl.user_session.get("axiom_agent")
+    chat_history = cl.user_session.get("chat_history")
+    # Add user message to history
+    chat_history.append({"role": "user", "content": message.content})
+    cl.user_session.set("chat_history", chat_history)
+
+    response_generator = agent.stream_agent(chat_history)
+
+    full_response = ""
+    msg = cl.Message(content="") # Initialize an empty message for streaming
+
+    async for token in response_generator:
+        full_response += token
+        await msg.stream_token(token) # Stream tokens to the UI
+    await msg.send() # Send the final message
+
+    chat_history.append({"role": "assistant", "content": full_response})
+    cl.user_session.set("chat_history", chat_history) # Update chat history

src/axiom/agent.py

@@ -22,14 +22,16 @@ class AxiomAgent:
         model: Optional[str] = None,
         tools: Optional[list[Tool]] = None,
         mcp_servers: Optional[list[MCPServer]] = None,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
     ):
-        self._api_key = settings.GOOGLE_API_KEY
-        self.base_url = settings.BASE_URL
+        self._api_key = api_key or settings.GOOGLE_API_KEY
+        self.base_url = base_url or settings.BASE_URL
         self.model_name = model or settings.DEFAULT_MODEL
         
         self._client: AsyncOpenAI = AsyncOpenAI(
             api_key=self._api_key,
-            base_url=self._base_url,
+            base_url=self.base_url,
         )
 
         self.agent = Agent(