add option between tool calling and code agent

smolagent_chat.py

@@ -12,7 +12,7 @@ import re
 from typing import List, Dict, Any
 import gradio as gr
 from gradio import ChatMessage
-from smolagents import MCPClient, ToolCallingAgent, OpenAIServerModel, AzureOpenAIModel, InferenceClientModel, stream_to_gradio
+from smolagents import MCPClient, ToolCallingAgent, CodeAgent, OpenAIServerModel, AzureOpenAIModel, InferenceClientModel, stream_to_gradio
 
 
 class Colors:
@@ -63,6 +63,7 @@ class KnowledgeGraphChatAgent:
         self.mcp_server_url = mcp_server_url or os.getenv("MCP_SERVER_URL", "http://localhost:4000/mcp")
         self.model = None
         self.agent = None
+        self.agent_type = None
         self.mcp_client = None
         self.tools = None
         self.conversation_history = []
@@ -126,6 +127,7 @@ class KnowledgeGraphChatAgent:
                 model_kwargs = {
                     "model_id": model_name,
                     "token": api_key,
+                    "bill_to": "epita"
                 }
                 if provider:
                     model_kwargs["provider"] = provider
@@ -151,7 +153,7 @@ class KnowledgeGraphChatAgent:
             print_error(f"Failed to initialize model: {e}")
             raise
     
-    def _initialize_agent(self, max_steps: int = None):
+    def _initialize_agent(self, agent_type: str = "tool_calling", max_steps: int = None):
         """Initialize the agent using the configured model and pre-loaded MCP tools."""
         if not self.model:
             raise ValueError("Model must be initialized before creating agent!")
@@ -160,16 +162,30 @@ class KnowledgeGraphChatAgent:
         
         try:
             max_steps = max_steps or int(os.getenv("MAX_STEPS", 5))
+            self.agent_type = agent_type
             
-            self.agent = ToolCallingAgent(
-                tools=self.tools,
-                model=self.model,
-                name="KnowledgeGraphAgent",
-                max_steps=max_steps,
-                add_base_tools=False,
-                instructions=CUSTOM_INSTRUCTIONS
-            )
-            print_success("Agent initialized successfully!")
+            print_info(f"Initializing {agent_type} agent...")
+            
+            if agent_type == "code":
+                self.agent = CodeAgent(
+                    tools=self.tools,
+                    model=self.model,
+                    name="KnowledgeGraphCodeAgent",
+                    max_steps=max_steps,
+                    add_base_tools=False,
+                    additional_authorized_imports=["numpy", "pandas", "re", "json"],
+                    instructions=CUSTOM_INSTRUCTIONS
+                )
+            else:  # tool_calling
+                self.agent = ToolCallingAgent(
+                    tools=self.tools,
+                    model=self.model,
+                    name="KnowledgeGraphAgent",
+                    max_steps=max_steps,
+                    add_base_tools=False,
+                    instructions=CUSTOM_INSTRUCTIONS
+                )
+            print_success(f"{agent_type.title()} agent initialized successfully!")
         except Exception as e:
             print_error(f"Failed to initialize agent: {e}")
             raise
@@ -299,6 +315,25 @@ def create_gradio_interface(agent: KnowledgeGraphChatAgent):
                         info="Choose between HuggingFace Inference, OpenAI, or Azure OpenAI"
                     )
                 
+                # Agent type selection
+                with gr.Row():
+                    agent_type = gr.Radio(
+                        choices=["tool_calling", "code"],
+                        value="tool_calling",
+                        label="Agent Type",
+                        info="Choose the agent architecture"
+                    )
+                
+                gr.Markdown("""
+                **Agent Type Guide:**
+
+- **Tool Calling Agent** (Recommended): Uses native function-calling capabilities. Works well for both medium and large models (GPT-4, Claude, Qwen 70B+). More reliable and efficient because the output format is constrained.
+
+- **Code Agent**: Generates Python code to solve tasks. Best used with larger, more capable models that have strong code-generation abilities. Smaller models often struggle with correctness, debugging, and multi-step reasoning.
+
+⚠️ **Note**: Smaller models (<30B parameters) generally struggle with both complex tool calling and code generation, but code generation tends to be more error-prone. For best results, use high-capability models like GPT-4, Claude 3.5, or Qwen 70B+ for Code Agent workflows, and prefer Tool Calling Agents whenever possible.
+""")
+                
                 # Model name field (shown for all types)
                 with gr.Row() as model_name_row:
                     model_name = gr.Textbox(
@@ -439,7 +474,7 @@ def create_gradio_interface(agent: KnowledgeGraphChatAgent):
             """)
         
         # Handle agent initialization
-        def initialize_agent(mtype, mname, akey, burl, azure_akey, azure_ep, aversion, hf_tok, hf_prov, msteps):
+        def initialize_agent(atype, mtype, mname, akey, burl, azure_akey, azure_ep, aversion, hf_tok, hf_prov, msteps):
             try:
                 if mtype == "azure":
                     agent._initialize_model(
@@ -463,9 +498,9 @@ def create_gradio_interface(agent: KnowledgeGraphChatAgent):
                         base_url=burl,
                         model_name=mname
                     )
-                agent._initialize_agent(max_steps=int(msteps))
+                agent._initialize_agent(agent_type=atype, max_steps=int(msteps))
                 return (
-                    gr.update(value="**Status:** ✅ Agent Ready!"),
+                    gr.update(value=f"**Status:** ✅ Agent Ready! (Using {atype.replace('_', ' ').title()} Agent)"),
                     gr.update(visible=False),  # Hide init section
                     gr.update(visible=True)    # Show chat section
                 )
@@ -479,7 +514,7 @@ def create_gradio_interface(agent: KnowledgeGraphChatAgent):
         
         init_btn.click(
             fn=initialize_agent,
-            inputs=[model_type, model_name, api_key, base_url, azure_api_key, azure_endpoint, api_version, hf_token, hf_provider, max_steps],
+            inputs=[agent_type, model_type, model_name, api_key, base_url, azure_api_key, azure_endpoint, api_version, hf_token, hf_provider, max_steps],
             outputs=[init_status, init_section, chat_section]
         )
         
@@ -547,3 +582,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+