Spaces:
Sleeping
Sleeping
Add tools and system prompt to base_agent, fix code exec tool, and add
Browse files- .gitignore +1 -0
- agents/base_agent.py +23 -1
- config/models.py +0 -1
- test.py +10 -0
- tools/code_execution.py +2 -2
.gitignore
CHANGED
|
@@ -1,2 +1,3 @@
|
|
| 1 |
.venv
|
| 2 |
.env
|
|
|
|
|
|
| 1 |
.venv
|
| 2 |
.env
|
| 3 |
+
__pycache__/
|
agents/base_agent.py
CHANGED
|
@@ -1,6 +1,28 @@
|
|
| 1 |
from langgraph.prebuilt import create_react_agent
|
| 2 |
from config.models import get_default_model
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
|
|
|
|
|
|
|
|
|
|
| 4 |
model = get_default_model()
|
| 5 |
|
| 6 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
from langgraph.prebuilt import create_react_agent
|
| 2 |
from config.models import get_default_model
|
| 3 |
+
from tools import code_execution
|
| 4 |
+
from tools.arxiv_search import arxiv_tool
|
| 5 |
+
from tools.wikipedia_search import wikipedia_search_tool
|
| 6 |
+
from tools.calculator import calculator_tool
|
| 7 |
+
from tools.code_execution import CodeExecutionTool
|
| 8 |
+
from tools.google_search import GoogleSearchTool
|
| 9 |
+
from tools.wikipedia_search import wikipedia_search_tool
|
| 10 |
+
import os
|
| 11 |
+
from dotenv import load_dotenv
|
| 12 |
|
| 13 |
+
load_dotenv()
|
| 14 |
+
|
| 15 |
+
SYSTEM_PROMPT = "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."
|
| 16 |
model = get_default_model()
|
| 17 |
|
| 18 |
+
code_execution_tool = CodeExecutionTool(api_key=os.environ['GOOGLE_API_KEY'])
|
| 19 |
+
google_search_tool = GoogleSearchTool(api_key=os.environ['GOOGLE_API_KEY'])
|
| 20 |
+
tools = [
|
| 21 |
+
arxiv_tool,
|
| 22 |
+
wikipedia_search_tool,
|
| 23 |
+
calculator_tool,
|
| 24 |
+
code_execution_tool,
|
| 25 |
+
google_search_tool,
|
| 26 |
+
]
|
| 27 |
+
|
| 28 |
+
agent_executor = create_react_agent(model, tools, prompt=SYSTEM_PROMPT)
|
config/models.py
CHANGED
|
@@ -1,5 +1,4 @@
|
|
| 1 |
from langchain.chat_models import init_chat_model
|
| 2 |
-
import os
|
| 3 |
|
| 4 |
def get_default_model():
|
| 5 |
"""Get the default chat model"""
|
|
|
|
| 1 |
from langchain.chat_models import init_chat_model
|
|
|
|
| 2 |
|
| 3 |
def get_default_model():
|
| 4 |
"""Get the default chat model"""
|
test.py
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from agents.base_agent import agent_executor
|
| 2 |
+
from langchain_core.messages import HumanMessage
|
| 3 |
+
|
| 4 |
+
# question = "In Series 9, Episode 11 of Doctor Who, the Doctor is trapped inside an ever-shifting maze. What is this location called in the official script for the episode? Give the setting exactly as it appears in the first scene heading."
|
| 5 |
+
# question = "In Emily Midkiff's June 2014 article in a journal named for the one of Hreidmar's sons that guarded his house, what word was quoted from two different authors in distaste for the nature of dragon depictions?"
|
| 6 |
+
question = "If Eliud Kipchoge could maintain his record-making marathon pace indefinitely, how many thousand hours would it take him to run the distance between the Earth and the Moon its closest approach? Please use the minimum perigee value on the Wikipedia page for the Moon when carrying out your calculation. Round your result to the nearest 1000 hours and do not use any comma separators if necessary."
|
| 7 |
+
messages = [HumanMessage(content=question)]
|
| 8 |
+
response = agent_executor.invoke({"messages": messages})
|
| 9 |
+
for message in response["messages"]:
|
| 10 |
+
message.pretty_print()
|
tools/code_execution.py
CHANGED
|
@@ -6,13 +6,13 @@ from pydantic import BaseModel, Field
|
|
| 6 |
from google.genai.types import Tool, ToolCodeExecution, GenerateContentConfig, Content, Part
|
| 7 |
from google import genai
|
| 8 |
|
| 9 |
-
class
|
| 10 |
task: str = Field(description="The task to be executed")
|
| 11 |
|
| 12 |
class CodeExecutionTool(BaseTool):
|
| 13 |
name: str = "code_execution"
|
| 14 |
description: str = "Useful for generate and execute python code for a given task"
|
| 15 |
-
args_schema: Optional[ArgsSchema] =
|
| 16 |
client: Any = None
|
| 17 |
model_id: str = "gemini-2.0-flash"
|
| 18 |
return_direct: bool = False
|
|
|
|
| 6 |
from google.genai.types import Tool, ToolCodeExecution, GenerateContentConfig, Content, Part
|
| 7 |
from google import genai
|
| 8 |
|
| 9 |
+
class CodeExecutionInput(BaseModel):
|
| 10 |
task: str = Field(description="The task to be executed")
|
| 11 |
|
| 12 |
class CodeExecutionTool(BaseTool):
|
| 13 |
name: str = "code_execution"
|
| 14 |
description: str = "Useful for generate and execute python code for a given task"
|
| 15 |
+
args_schema: Optional[ArgsSchema] = CodeExecutionInput
|
| 16 |
client: Any = None
|
| 17 |
model_id: str = "gemini-2.0-flash"
|
| 18 |
return_direct: bool = False
|