Spaces:
Sleeping
Sleeping
File size: 2,638 Bytes
4107a1a 9be739d 4107a1a 9be739d 4107a1a 818d867 4107a1a 818d867 4107a1a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
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 CodeExecutionInput(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] = CodeExecutionInput
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, task: str, run_manager: Optional[CallbackManagerForToolRun]=None) -> str:
"""Run the code execution tool asynchronously."""
return self._run(task, 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)
|