refactor: streamline agent and model wrapping with improved error handling

Author: Dwij1704

agentops/instrumentation/agentic/smolagents/instrumentor.py

@@ -83,49 +83,18 @@ def _custom_wrap(self, **kwargs):
         This is called after normal wrapping, but we use it for all wrapping
         since we don't have normal wrapped methods.
         """
-        # Core agent operations - wrap with error handling
-        try:
-            # Import the module first to ensure it's loaded
-            import smolagents.agents
-
-            # Check if classes exist before wrapping
-            if hasattr(smolagents.agents, "CodeAgent"):
-                wrap_function_wrapper("smolagents.agents", "CodeAgent.run", self._agent_run_wrapper(self._tracer))
-                logger.debug("Successfully wrapped CodeAgent.run")
-            else:
-                logger.debug("CodeAgent not found in smolagents.agents")
-
-            if hasattr(smolagents.agents, "ToolCallingAgent"):
-                wrap_function_wrapper(
-                    "smolagents.agents", "ToolCallingAgent.run", self._agent_run_wrapper(self._tracer)
-                )
-                wrap_function_wrapper(
-                    "smolagents.agents",
-                    "ToolCallingAgent.execute_tool_call",
-                    self._tool_execution_wrapper(self._tracer),
-                )
-                logger.debug("Successfully wrapped ToolCallingAgent methods")
-            else:
-                logger.debug("ToolCallingAgent not found in smolagents.agents")
-        except (ImportError, AttributeError) as e:
-            logger.debug(f"Failed to wrap agent operations: {e}")
+        # Core agent operations
+        wrap_function_wrapper("smolagents.agents", "CodeAgent.run", self._agent_run_wrapper(self._tracer))
+        wrap_function_wrapper("smolagents.agents", "ToolCallingAgent.run", self._agent_run_wrapper(self._tracer))
+
+        # Tool calling operations
+        wrap_function_wrapper(
+            "smolagents.agents", "ToolCallingAgent.execute_tool_call", self._tool_execution_wrapper(self._tracer)
+        )
 
         # Model operations with proper model name extraction
-        try:
-            # Import the module first to ensure it's loaded
-            import smolagents.models
-
-            # Check if class exists before wrapping
-            if hasattr(smolagents.models, "LiteLLMModel"):
-                wrap_function_wrapper("smolagents.models", "LiteLLMModel.generate", self._llm_wrapper(self._tracer))
-                wrap_function_wrapper(
-                    "smolagents.models", "LiteLLMModel.generate_stream", self._llm_wrapper(self._tracer)
-                )
-                logger.debug("Successfully wrapped LiteLLMModel methods")
-            else:
-                logger.debug("LiteLLMModel not found in smolagents.models")
-        except (ImportError, AttributeError) as e:
-            logger.debug(f"Failed to wrap model operations: {e}")
+        wrap_function_wrapper("smolagents.models", "LiteLLMModel.generate", self._llm_wrapper(self._tracer))
+        wrap_function_wrapper("smolagents.models", "LiteLLMModel.generate_stream", self._llm_wrapper(self._tracer))
 
         logger.info("SmoLAgents instrumentation enabled")
 
@@ -289,48 +258,29 @@ def _custom_unwrap(self, **kwargs):
         # Unwrap all instrumented methods
         from opentelemetry.instrumentation.utils import unwrap
 
-        # Try to import modules before attempting unwrap
         try:
-            import smolagents.agents
-
-            if hasattr(smolagents.agents, "CodeAgent"):
-                try:
-                    unwrap("smolagents.agents", "CodeAgent.run")
-                    logger.debug("Successfully unwrapped CodeAgent.run")
-                except Exception as e:
-                    logger.debug(f"Failed to unwrap CodeAgent.run: {e}")
-
-            if hasattr(smolagents.agents, "ToolCallingAgent"):
-                try:
-                    unwrap("smolagents.agents", "ToolCallingAgent.run")
-                    logger.debug("Successfully unwrapped ToolCallingAgent.run")
-                except Exception as e:
-                    logger.debug(f"Failed to unwrap ToolCallingAgent.run: {e}")
+            unwrap("smolagents.agents", "CodeAgent.run")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap CodeAgent.run: {e}")
 
-                try:
-                    unwrap("smolagents.agents", "ToolCallingAgent.execute_tool_call")
-                    logger.debug("Successfully unwrapped ToolCallingAgent.execute_tool_call")
-                except Exception as e:
-                    logger.debug(f"Failed to unwrap ToolCallingAgent.execute_tool_call: {e}")
-        except ImportError as e:
-            logger.debug(f"smolagents.agents module not available for unwrapping: {e}")
+        try:
+            unwrap("smolagents.agents", "ToolCallingAgent.run")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap ToolCallingAgent.run: {e}")
 
         try:
-            import smolagents.models
+            unwrap("smolagents.agents", "ToolCallingAgent.execute_tool_call")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap ToolCallingAgent.execute_tool_call: {e}")
 
-            if hasattr(smolagents.models, "LiteLLMModel"):
-                try:
-                    unwrap("smolagents.models", "LiteLLMModel.generate")
-                    logger.debug("Successfully unwrapped LiteLLMModel.generate")
-                except Exception as e:
-                    logger.debug(f"Failed to unwrap LiteLLMModel.generate: {e}")
+        try:
+            unwrap("smolagents.models", "LiteLLMModel.generate")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap LiteLLMModel.generate: {e}")
 
-                try:
-                    unwrap("smolagents.models", "LiteLLMModel.generate_stream")
-                    logger.debug("Successfully unwrapped LiteLLMModel.generate_stream")
-                except Exception as e:
-                    logger.debug(f"Failed to unwrap LiteLLMModel.generate_stream: {e}")
-        except ImportError as e:
-            logger.debug(f"smolagents.models module not available for unwrapping: {e}")
+        try:
+            unwrap("smolagents.models", "LiteLLMModel.generate_stream")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap LiteLLMModel.generate_stream: {e}")
 
         logger.info("SmoLAgents instrumentation disabled")