""" Autonomous AI Agent with MCP Tool Calling using Granite 4.0 H-1B (Open Source) This agent uses IBM Granite 4.0 H-1B (1.5B params) loaded locally via transformers to autonomously decide which MCP tools to call. Granite 4.0 H-1B is optimized for tool calling and function calling tasks. Uses ReAct (Reasoning + Acting) prompting pattern for reliable tool calling. """ import os import re import json import uuid import logging import asyncio from typing import List, Dict, Any, AsyncGenerator from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM import torch from mcp.tools.definitions import MCP_TOOLS, list_all_tools from mcp.registry import MCPRegistry logger = logging.getLogger(__name__) class AutonomousMCPAgentGranite: """ AI Agent that autonomously uses MCP servers as tools using Granite 4. Uses ReAct (Reasoning + Acting) pattern: 1. Thought: AI reasons about what to do next 2. Action: AI decides which tool to call 3. Observation: AI sees the tool result 4. Repeat until task complete """ def __init__(self, mcp_registry: MCPRegistry, hf_token: str = None): """ Initialize the autonomous agent with Granite 4.0 H-1B Args: mcp_registry: MCP registry with all servers hf_token: HuggingFace token (optional, for accessing private models) """ self.mcp_registry = mcp_registry self.hf_token = hf_token or os.getenv("HF_API_TOKEN") or os.getenv("HF_TOKEN") # Use Granite 4.0 H-1B (1.5B params, optimized for tool calling) self.model_name = "ibm-granite/granite-4.0-h-1b" logger.info(f"Loading Granite 4.0 H-1B model locally...") # Load model with optimizations for CPU/limited memory try: logger.info(f"📥 Downloading tokenizer from {self.model_name}...") # Use bfloat16 for better efficiency, float32 fallback for CPU self.tokenizer = AutoTokenizer.from_pretrained( self.model_name, token=self.hf_token, trust_remote_code=True ) logger.info(f"✓ Tokenizer loaded successfully") # Check device availability device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32 logger.info(f"💻 Device: {device}, dtype: {dtype}") logger.info(f"📥 Downloading model weights (~1.5GB)...") # For hybrid models like Granite H-1B, we need explicit device placement if torch.cuda.is_available(): # GPU available - use device_map self.model = AutoModelForCausalLM.from_pretrained( self.model_name, token=self.hf_token, torch_dtype=dtype, device_map="auto", low_cpu_mem_usage=True, trust_remote_code=True ) else: # CPU only - load with 8-bit quantization to reduce memory logger.info(f"⚠️ Loading on CPU (no GPU available)") logger.info(f"💾 Using 8-bit quantization to reduce memory usage") try: # Try loading with 8-bit quantization (requires bitsandbytes) from transformers import BitsAndBytesConfig quantization_config = BitsAndBytesConfig( load_in_8bit=True, llm_int8_threshold=6.0 ) self.model = AutoModelForCausalLM.from_pretrained( self.model_name, token=self.hf_token, quantization_config=quantization_config, low_cpu_mem_usage=False, trust_remote_code=True ) logger.info(f"✓ Loaded with 8-bit quantization (~50% memory reduction)") except (ImportError, Exception) as e: # Fallback to float32 if 8-bit fails logger.warning(f"⚠️ 8-bit quantization failed: {e}") logger.info(f"⚠️ Falling back to float32 (may use ~4-6GB RAM)") self.model = AutoModelForCausalLM.from_pretrained( self.model_name, token=self.hf_token, torch_dtype=torch.float32, # Use float32 for CPU low_cpu_mem_usage=False, # Disable to avoid meta device trust_remote_code=True ) # Verify all parameters are on CPU, not meta logger.info(f"🔍 Verifying model is materialized on CPU...") param_devices = set() for param in self.model.parameters(): param_devices.add(str(param.device)) if 'meta' in param_devices: logger.error(f"❌ Model still has parameters on meta device!") raise RuntimeError("Model not properly materialized. Try upgrading transformers: pip install --upgrade transformers") logger.info(f"✓ All parameters on: {param_devices}") logger.info(f"✓ Model weights loaded") # Set model to eval mode self.model.eval() logger.info(f"✓ Model set to evaluation mode") # Get model device and memory info try: model_device = next(self.model.parameters()).device logger.info(f"✓ Model loaded successfully on device: {model_device}") except StopIteration: logger.warning(f"⚠️ Could not determine model device (no parameters)") # Memory info if available if torch.cuda.is_available(): memory_allocated = torch.cuda.memory_allocated() / 1024**3 logger.info(f"📊 GPU Memory allocated: {memory_allocated:.2f} GB") except Exception as e: logger.error(f"❌ Failed to load model: {e}", exc_info=True) raise # Create tool descriptions for the AI self.tools_description = self._create_tools_description() logger.info(f"Autonomous MCP Agent initialized with model: {self.model_name}") def _generate_text(self, prompt: str) -> str: """ Generate text using the local Granite model (synchronous, for use in executor) Args: prompt: The input prompt Returns: Generated text """ import time import gc start_time = time.time() # Force garbage collection before inference to free memory gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Tokenize input with aggressive truncation to save memory logger.info(f"🔤 Tokenizing input (length: {len(prompt)} chars)...") inputs = self.tokenizer( prompt, return_tensors="pt", truncation=True, max_length=2048 # Reduced from 4096 to save memory ) num_input_tokens = inputs["input_ids"].shape[-1] logger.info(f"✓ Tokenized to {num_input_tokens} tokens") # Get target device - handle models split across devices try: target_device = next(self.model.parameters()).device except StopIteration: # Fallback if no parameters found target_device = torch.device('cpu') logger.info(f"📍 Moving inputs to device: {target_device}") # Move to same device as model inputs = {k: v.to(target_device) for k, v in inputs.items()} # Generate with memory-efficient settings logger.info(f"🤖 Generating response (max 400 tokens, temp=0.1)...") with torch.no_grad(): outputs = self.model.generate( **inputs, max_new_tokens=400, # Reduced from 800 to save memory temperature=0.1, # Low temperature for deterministic reasoning top_p=0.9, do_sample=True, pad_token_id=self.tokenizer.eos_token_id, eos_token_id=self.tokenizer.eos_token_id, use_cache=True, # Use KV cache for efficiency num_beams=1, # Greedy decoding to save memory ) # Decode only the new tokens response = self.tokenizer.decode( outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True ) elapsed = time.time() - start_time num_output_tokens = outputs.shape[-1] - num_input_tokens tokens_per_sec = num_output_tokens / elapsed if elapsed > 0 else 0 logger.info(f"✓ Generated {num_output_tokens} tokens in {elapsed:.1f}s ({tokens_per_sec:.1f} tokens/sec)") logger.info(f"📝 Response preview: {response[:100]}...") # Clean up to free memory del inputs, outputs gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() return response def _create_tools_description(self) -> str: """Create a formatted description of all available tools for the AI""" tools_text = "## Available MCP Tools:\n\n" for tool in MCP_TOOLS: tools_text += f"**{tool['name']}**\n" tools_text += f" Description: {tool['description']}\n" tools_text += f" Parameters:\n" for prop_name, prop_data in tool['input_schema']['properties'].items(): required = prop_name in tool['input_schema'].get('required', []) tools_text += f" - {prop_name} ({prop_data['type']}){'*' if required else ''}: {prop_data.get('description', '')}\n" tools_text += "\n" return tools_text def _create_system_prompt(self) -> str: """Create the system prompt for ReAct pattern""" return f"""You are an autonomous AI agent for B2B sales automation using the ReAct (Reasoning + Acting) framework. You have access to MCP (Model Context Protocol) tools that let you: - Search the web for company information and news - Save prospects, companies, contacts, and facts to a database - Send emails and manage email threads - Schedule meetings and generate calendar invites {self.tools_description} ## ReAct Format: You must respond using this EXACT format: Thought: [Your reasoning about what to do next] Action: [tool_name] Action Input: {{"param1": "value1", "param2": "value2"}} After you see the Observation, you can continue with more Thought/Action/Observation cycles. When you've completed the task, respond with: Thought: [Your final reasoning] Final Answer: [Your complete response to the user] ## Important Rules: 1. Always use "Thought:" to reason before acting 2. Always use "Action:" followed by exact tool name 3. Always use "Action Input:" with valid JSON 4. Use tools multiple times if needed 5. Save important data to the database 6. When done, give a "Final Answer:" ## Example: Thought: I need to research Shopify first Action: search_web Action Input: {{"query": "Shopify company information"}} [You'll see Observation with results] Thought: Now I should save the company data Action: save_company Action Input: {{"company_id": "shopify", "name": "Shopify", "domain": "shopify.com"}} [Continue until task complete...] Thought: I've gathered all the information and saved it Final Answer: I've successfully researched Shopify and created a prospect profile with company information and recent facts. Now complete your assigned task!""" async def run( self, task: str, max_iterations: int = 15 ) -> AsyncGenerator[Dict[str, Any], None]: """ Run the agent autonomously on a task using ReAct pattern. Args: task: The task to complete max_iterations: Maximum tool calls to prevent infinite loops Yields: Events showing agent's progress and tool calls """ yield { "type": "agent_start", "message": f"🤖 Autonomous AI Agent (Granite 4) starting task", "task": task, "model": self.model } # Initialize conversation with system prompt and task conversation_history = f"""{self._create_system_prompt()} ## Task: {task} Begin! """ iteration = 0 while iteration < max_iterations: iteration += 1 yield { "type": "iteration_start", "iteration": iteration, "message": f"🔄 Iteration {iteration}: AI reasoning..." } try: # Get AI response using ReAct pattern response_text = "" try: # Generate using local model # Run in executor to avoid blocking the event loop response_text = await asyncio.get_event_loop().run_in_executor( None, self._generate_text, conversation_history ) except Exception as gen_error: logger.error(f"Text generation failed: {gen_error}", exc_info=True) yield { "type": "agent_error", "error": str(gen_error), "message": f"❌ Model error: {str(gen_error)}" } break # Check if we got a response if not response_text or not response_text.strip(): logger.warning("Empty response from model") yield { "type": "parse_error", "message": "⚠️ Model returned empty response. Retrying...", "response": "" } continue # Log the raw response for debugging logger.info(f"Model response (iteration {iteration}): {response_text[:200]}...") # Parse the response for Thought, Action, Action Input thought_match = re.search(r'Thought:\s*(.+?)(?=\n(?:Action:|Final Answer:)|$)', response_text, re.DOTALL) action_match = re.search(r'Action:\s*(\w+)', response_text) action_input_match = re.search(r'Action Input:\s*(\{.+?\})', response_text, re.DOTALL) final_answer_match = re.search(r'Final Answer:\s*(.+?)$', response_text, re.DOTALL) # Extract thought if thought_match: thought = thought_match.group(1).strip() yield { "type": "thought", "thought": thought, "message": f"💭 Thought: {thought}" } # Check if AI wants to finish if final_answer_match: final_answer = final_answer_match.group(1).strip() yield { "type": "agent_complete", "message": "✅ Task complete!", "final_answer": final_answer, "iterations": iteration } break # Execute action if present if action_match and action_input_match: tool_name = action_match.group(1).strip() action_input_str = action_input_match.group(1).strip() # Parse action input JSON try: tool_input = json.loads(action_input_str) except json.JSONDecodeError as e: error_msg = f"Invalid JSON in Action Input: {e}" logger.error(error_msg) # Give feedback to AI conversation_history += response_text conversation_history += f"\nObservation: Error - {error_msg}. Please provide valid JSON.\n\n" continue yield { "type": "tool_call", "tool": tool_name, "input": tool_input, "message": f"🔧 Action: {tool_name}" } # Execute the MCP tool try: result = await self._execute_mcp_tool(tool_name, tool_input) yield { "type": "tool_result", "tool": tool_name, "result": result, "message": f"✓ Tool {tool_name} completed" } # Add to conversation history conversation_history += response_text conversation_history += f"\nObservation: {json.dumps(result, default=str)}\n\n" except Exception as e: error_msg = str(e) logger.error(f"Tool execution failed: {tool_name} - {error_msg}") yield { "type": "tool_error", "tool": tool_name, "error": error_msg, "message": f"❌ Tool {tool_name} failed: {error_msg}" } # Give error feedback to AI conversation_history += response_text conversation_history += f"\nObservation: Error - {error_msg}\n\n" else: # No action found - AI might be confused yield { "type": "parse_error", "message": "⚠️ Could not parse Action from AI response", "response": response_text } # Give feedback to AI conversation_history += response_text conversation_history += "\nObservation: Please follow the format: 'Action: tool_name' and 'Action Input: {...}'\n\n" except (RuntimeError, StopIteration, StopAsyncIteration) as stop_err: # Handle StopIteration errors that get wrapped in RuntimeError error_msg = str(stop_err) logger.error(f"Stop iteration in agent loop: {error_msg}", exc_info=True) if "StopIteration" in error_msg or "StopAsyncIteration" in error_msg: yield { "type": "agent_error", "error": "Model inference error - possibly model not available or API issue", "message": f"❌ Model inference failed. Please check:\n" f" 1. HF_API_TOKEN is valid\n" f" 2. Model '{self.model}' is accessible\n" f" 3. HuggingFace Inference API is operational" } else: yield { "type": "agent_error", "error": error_msg, "message": f"❌ Agent error: {error_msg}" } break except Exception as e: logger.error(f"Agent iteration failed: {e}", exc_info=True) yield { "type": "agent_error", "error": str(e), "message": f"❌ Agent error: {str(e)}" } break if iteration >= max_iterations: yield { "type": "agent_max_iterations", "message": f"⚠️ Reached maximum iterations ({max_iterations})", "iterations": iteration } async def _execute_mcp_tool(self, tool_name: str, tool_input: Dict[str, Any]) -> Any: """ Execute an MCP tool by routing to the appropriate MCP server. This is where we actually call the MCP servers! """ # ============ SEARCH MCP SERVER ============ if tool_name == "search_web": query = tool_input["query"] max_results = tool_input.get("max_results", 5) results = await self.mcp_registry.search.query(query, max_results=max_results) return { "results": results[:max_results], "count": len(results[:max_results]) } elif tool_name == "search_news": query = tool_input["query"] max_results = tool_input.get("max_results", 5) results = await self.mcp_registry.search.query(f"{query} news", max_results=max_results) return { "results": results[:max_results], "count": len(results[:max_results]) } # ============ STORE MCP SERVER ============ elif tool_name == "save_prospect": prospect_data = { "id": tool_input.get("prospect_id", str(uuid.uuid4())), "company": { "id": tool_input.get("company_id"), "name": tool_input.get("company_name"), "domain": tool_input.get("company_domain") }, "fit_score": tool_input.get("fit_score", 0), "status": tool_input.get("status", "new"), "metadata": tool_input.get("metadata", {}) } result = await self.mcp_registry.store.save_prospect(prospect_data) return {"status": result, "prospect_id": prospect_data["id"]} elif tool_name == "get_prospect": prospect_id = tool_input["prospect_id"] prospect = await self.mcp_registry.store.get_prospect(prospect_id) return prospect or {"error": "Prospect not found"} elif tool_name == "list_prospects": prospects = await self.mcp_registry.store.list_prospects() status_filter = tool_input.get("status") if status_filter: prospects = [p for p in prospects if p.get("status") == status_filter] return { "prospects": prospects, "count": len(prospects) } elif tool_name == "save_company": company_data = { "id": tool_input.get("company_id", str(uuid.uuid4())), "name": tool_input["name"], "domain": tool_input["domain"], "industry": tool_input.get("industry"), "description": tool_input.get("description"), "employee_count": tool_input.get("employee_count") } result = await self.mcp_registry.store.save_company(company_data) return {"status": result, "company_id": company_data["id"]} elif tool_name == "get_company": company_id = tool_input["company_id"] company = await self.mcp_registry.store.get_company(company_id) return company or {"error": "Company not found"} elif tool_name == "save_fact": fact_data = { "id": tool_input.get("fact_id", str(uuid.uuid4())), "company_id": tool_input["company_id"], "fact_type": tool_input["fact_type"], "content": tool_input["content"], "source_url": tool_input.get("source_url"), "confidence_score": tool_input.get("confidence_score", 0.8) } result = await self.mcp_registry.store.save_fact(fact_data) return {"status": result, "fact_id": fact_data["id"]} elif tool_name == "save_contact": contact_data = { "id": tool_input.get("contact_id", str(uuid.uuid4())), "company_id": tool_input["company_id"], "email": tool_input["email"], "first_name": tool_input.get("first_name"), "last_name": tool_input.get("last_name"), "title": tool_input.get("title"), "seniority": tool_input.get("seniority") } result = await self.mcp_registry.store.save_contact(contact_data) return {"status": result, "contact_id": contact_data["id"]} elif tool_name == "list_contacts_by_domain": domain = tool_input["domain"] contacts = await self.mcp_registry.store.list_contacts_by_domain(domain) return { "contacts": contacts, "count": len(contacts) } elif tool_name == "check_suppression": supp_type = tool_input["suppression_type"] value = tool_input["value"] is_suppressed = await self.mcp_registry.store.check_suppression(supp_type, value) return { "suppressed": is_suppressed, "value": value, "type": supp_type } # ============ EMAIL MCP SERVER ============ elif tool_name == "send_email": to = tool_input["to"] subject = tool_input["subject"] body = tool_input["body"] prospect_id = tool_input["prospect_id"] thread_id = await self.mcp_registry.email.send(to, subject, body, prospect_id) return { "status": "sent", "thread_id": thread_id, "to": to } elif tool_name == "get_email_thread": prospect_id = tool_input["prospect_id"] thread = await self.mcp_registry.email.get_thread(prospect_id) return thread or {"error": "No email thread found"} # ============ CALENDAR MCP SERVER ============ elif tool_name == "suggest_meeting_slots": num_slots = tool_input.get("num_slots", 3) slots = await self.mcp_registry.calendar.suggest_slots() return { "slots": slots[:num_slots], "count": len(slots[:num_slots]) } elif tool_name == "generate_calendar_invite": start_time = tool_input["start_time"] end_time = tool_input["end_time"] title = tool_input["title"] slot = { "start_iso": start_time, "end_iso": end_time, "title": title } ics = await self.mcp_registry.calendar.generate_ics(slot) return { "ics_content": ics, "meeting": slot } else: raise ValueError(f"Unknown MCP tool: {tool_name}")