agaliano's picture
Create agent.py
8763bc1 verified
Raw
History Blame Contribute Delete
2.37 kB
from smolagents import LiteLLMModel
from smolagents import CodeAgent, ToolCallingAgent, GoogleSearchTool, HfApiModel, VisitWebpageTool, DuckDuckGoSearchTool, tool, PythonInterpreterTool, WikipediaSearchTool
import requests
from youtube_transcript_api import YouTubeTranscriptApi
@tool
def reverse_sentence_tool(reverse_sentence: str) -> str:
"""
This tool receives a sentence where both the word order and the characters in each word are reversed.
It returns the sentence with words and order corrected.
Args:
reverse_sentence: A sentence with reversed words and reversed word order.
Returns:
A sentence in natural reading order.
"""
inverted_words = reverse_sentence.split(" ")[::-1]
correct_words = [word[::-1] for word in inverted_words]
return " ".join(correct_words)
@tool
def get_youtube_transcript(video_url: str) -> str:
"""
Fetches the transcript from a YouTube video if available.
Args:
video_url: Full URL to the YouTube video.
Returns:
Transcript text.
"""
video_id = video_url.split("v=")[-1]
transcript = YouTubeTranscriptApi.get_transcript(video_id)
full_text = " ".join([entry['text'] for entry in transcript])
return full_text
@tool
def check_answer(answer: str) -> str:
"""
Reviews the answer to check that it meets the requirements specified by the user and modifies it if necessary.
Args:
answer (str): The answer fo the Agent.
Returns:
str: The final answer.
"""
if answer[-1] == '.':
return answer[:-1]
if "St." in answer:
return answer.replace("St.", "Saint")
return answer
class BasicAgent:
def __init__(self):
self.api_key = ""
self.model = LiteLLMModel(model_id="gemini/gemini-2.0-flash", api_key=self.api_key)
self.agent = CodeAgent(
tools=[
DuckDuckGoSearchTool(),
PythonInterpreterTool(),
VisitWebpageTool(),
reverse_sentence_tool,
check_answer
],
model=self.model
)
print("BasicAgent initialized.")
def __call__(self, question: str) -> str:
print(f"Agent received question")
answer = self.agent.run(question)
return answer