File size: 717 Bytes
879b924 |
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 |
from transformers import ReactCodeAgent, HfApiEngine
from prompts import SQUAD_REACT_CODE_SYSTEM_PROMPT
from tools.squad_tools import SquadRetrieverTool, SquadQueryTool
from tools.text_to_image import TextToImageTool
def get_agent():
# model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
model_name = "http://localhost:1234/v1"
llm_engine = HfApiEngine(model_name)
TASK_SOLVING_TOOLBOX = [
SquadRetrieverTool(),
SquadQueryTool(),
TextToImageTool(),
]
# Initialize the agent with both tools
agent = ReactCodeAgent(
tools=TASK_SOLVING_TOOLBOX,
llm_engine=llm_engine,
system_prompt=SQUAD_REACT_CODE_SYSTEM_PROMPT,
)
return agent
|