RamV12's picture
Create agent.py
2f9cbee verified
raw
history blame
1.7 kB
# agent.py
import os
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from llama_index.core.agent.workflow import AgentWorkflow, ReActAgent
from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
from llama_index.core.tools import FunctionTool
class BasicAgent:
def __init__(self):
print("Initializing BasicAgent...")
# Define tools
def add(a: int, b: int) -> int:
return a + b
def multiply(a: int, b: int) -> int:
return a * b
add_tool = FunctionTool.from_defaults(fn=add)
multiply_tool = FunctionTool.from_defaults(fn=multiply)
# Load web search tool
tool_spec = DuckDuckGoSearchToolSpec()
search_tools = tool_spec.to_tool_list()
all_tools = [add_tool, multiply_tool] + search_tools
# Load LLM
hf_token = os.getenv("HF_TOKEN")
self.llm = HuggingFaceInferenceAPI(
model_name="Qwen/Qwen2.5-Coder-32B-Instruct",
token=hf_token,
)
# Create the agent
self.agent = ReActAgent(
name="universal_agent",
description="An assistant that can search the web and do math.",
system_prompt="You are a helpful assistant with access to tools.",
tools=all_tools,
llm=self.llm,
)
# Setup workflow
self.workflow = AgentWorkflow(
agents=[self.agent],
root_agent="universal_agent",
)
async def __call__(self, question: str) -> str:
print(f"[BasicAgent] Received question: {question}")
response = await self.workflow.run(user_msg=question)
return str(response)