dawid-lorek's picture
Update agent.py
92b0d1a verified
raw
history blame
1.35 kB
# agent.py
from llama_index.llms.openai import OpenAI
from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.core.tools import FunctionTool
# LangChain-based tools
from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun
from langchain_experimental.tools.python.tool import PythonREPLTool
from langchain_community.document_loaders import YoutubeLoader
# Tool wrappers
def search_duckduckgo(query: str) -> str:
return DuckDuckGoSearchRun().run(query)
def search_wikipedia(query: str) -> str:
return WikipediaQueryRun(api_wrapper=None).run(query)
def run_python(code: str) -> str:
return PythonREPLTool().run(code)
def get_youtube_transcript(url: str) -> str:
loader = YoutubeLoader.from_youtube_url(url, add_video_info=False)
docs = loader.load()
return " ".join(d.page_content for d in docs)
TOOLS = [
FunctionTool.from_defaults(search_duckduckgo),
FunctionTool.from_defaults(search_wikipedia),
FunctionTool.from_defaults(run_python),
FunctionTool.from_defaults(get_youtube_transcript),
]
llm = OpenAI(model="gpt-4")
agent = AgentWorkflow.from_tools_or_functions(
TOOLS,
llm=llm,
system_prompt="You are a helpful AI agent using external tools."
)
import asyncio
async def answer_question(question: str) -> str:
return await agent.arun(question)