Add CodeExecutionTool for generating and running Python code

tools/code_execution.py

@@ -0,0 +1,66 @@
+from typing import Any, Optional
+from langchain_core.callbacks import CallbackManagerForToolRun
+from langchain_core.tools import BaseTool
+from langchain_core.tools.base import ArgsSchema
+from pydantic import BaseModel, Field
+from google.genai.types import Tool, ToolCodeExecution, GenerateContentConfig, Content, Part
+from google import genai
+
+class CodeExectionInput(BaseModel):
+    task: str = Field(description="The task to be executed")
+
+class CodeExecutionTool(BaseTool):
+    name: str = "code_execution"
+    description: str = "Useful for generate and execute python code for a given task"
+    args_schema: Optional[ArgsSchema] = CodeExectionInput
+    client: Any = None
+    model_id: str = "gemini-2.0-flash"
+    return_direct: bool = False
+    code_execution_tool: Any = None
+
+    def __init__(self, api_key: Optional[str] = None, **kwargs):
+        super().__init__(**kwargs)
+        self.client = genai.Client(api_key=api_key)
+        self.code_execution_tool = Tool(code_execution=ToolCodeExecution())
+
+    def _run(self, task: str, run_manager: Optional[CallbackManagerForToolRun]=None) -> str:
+        """Generate and execute python code for a given task"""
+        try:
+            response = self.client.models.generate_content(
+                model=self.model_id,
+                contents=f"Task: {task} "
+                "Plan on how to solve the task. Generate and run code to solve the task.",
+                config=GenerateContentConfig(
+                    tools=[self.code_execution_tool],
+                    response_modalities=["TEXT"]
+                )
+            )
+
+            result = ""
+            for part in response.candidates[0].content.parts:
+                if part.executable_code is not None:
+                    result += part.executable_code.code
+                if part.code_execution_result is not None:
+                    result += "\n\n" + part.code_execution_result.output
+
+            return result
+        except Exception as e:
+            return f"Error performing code execution: {str(e)}"
+
+    async def _arun(self, query: str, run_manager: Optional[CallbackManagerForToolRun]=None) -> str:
+        """Run the code execution tool asynchronously."""
+        return self._run(query, run_manager=run_manager.get_sync())
+
+# Example usage
+if __name__ == "__main__":
+    from dotenv import load_dotenv
+
+    # Load API key from environment variables
+    load_dotenv()
+
+    # Create the Google Search runnable
+    code_execution = CodeExecutionTool()
+
+    # Run a search query
+    result = code_execution.invoke({"task": "What is the sum of the first 50 prime numbers?"})
+    print(result)