|
|
from smolagents import CodeAgent, HfApiModel, tool as smol_tool, GradioUI |
|
|
from smolagents.tools import Tool |
|
|
import datetime |
|
|
import pytz |
|
|
import yaml |
|
|
import tempfile |
|
|
import gradio as gr |
|
|
import os |
|
|
import re |
|
|
|
|
|
|
|
|
from tools.final_answer import FinalAnswerTool |
|
|
from tools.visit_webpage import VisitWebpageTool |
|
|
from tools.create_presentation import create_presentation |
|
|
from tools.web_search import DuckDuckGoSearchTool |
|
|
from tools.write_code import write_code_to_file |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@smol_tool |
|
|
def create_document(text: str, format: str = "docx") -> str: |
|
|
""" |
|
|
Creates a document file (DOCX or PDF) containing the provided text. |
|
|
Args: |
|
|
text: The full textual content that will be written into the document. |
|
|
format: The output file format. Must be either "docx" or "pdf". |
|
|
""" |
|
|
try: |
|
|
import docx |
|
|
from docx.shared import Pt |
|
|
temp_dir = tempfile.mkdtemp() |
|
|
doc = docx.Document() |
|
|
doc.add_heading("Generated Document", 0) |
|
|
|
|
|
for paragraph in text.split("\n"): |
|
|
if paragraph.strip(): |
|
|
doc.add_paragraph(paragraph) |
|
|
|
|
|
file_path = os.path.join(temp_dir, f"generated_document.{format.lower()}") |
|
|
doc.save(file_path) |
|
|
return file_path |
|
|
except Exception as e: |
|
|
return f"Error creating document: {str(e)}" |
|
|
|
|
|
@smol_tool |
|
|
def create_presentation_preview(title: str, slides: list[str]) -> list[str]: |
|
|
""" |
|
|
Create a presentation and return slide previews as images. |
|
|
Args: |
|
|
title: Presentation title |
|
|
slides: List of slide contents |
|
|
""" |
|
|
from PIL import Image, ImageDraw |
|
|
temp_dir = tempfile.mkdtemp() |
|
|
image_paths = [] |
|
|
for i, slide_text in enumerate(slides, start=1): |
|
|
img = Image.new("RGB", (800, 450), "white") |
|
|
draw = ImageDraw.Draw(img) |
|
|
draw.text((40, 40), f"{title}\n\nSlide {i}:\n{slide_text[:100]}", fill="black") |
|
|
path = os.path.join(temp_dir, f"slide_{i}.png") |
|
|
img.save(path) |
|
|
image_paths.append(path) |
|
|
return image_paths |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
final_answer = FinalAnswerTool() |
|
|
visit_webpage = VisitWebpageTool() |
|
|
web_search = DuckDuckGoSearchTool() |
|
|
|
|
|
model = HfApiModel( |
|
|
max_tokens=2096, |
|
|
temperature=0.5, |
|
|
model_id="Qwen/Qwen2.5-72B-Instruct", |
|
|
) |
|
|
|
|
|
|
|
|
with open("prompts.yaml", "r") as stream: |
|
|
prompt_templates = yaml.safe_load(stream) |
|
|
|
|
|
agent = CodeAgent( |
|
|
model=model, |
|
|
tools=[ |
|
|
final_answer, |
|
|
visit_webpage, |
|
|
web_search, |
|
|
create_document, |
|
|
create_presentation, |
|
|
create_presentation_preview, |
|
|
write_code_to_file, |
|
|
], |
|
|
max_steps=10, |
|
|
verbosity_level=1, |
|
|
prompt_templates=prompt_templates, |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class CustomGradioUI(GradioUI): |
|
|
def build_interface(self): |
|
|
with gr.Blocks(theme=gr.themes.Soft()) as interface: |
|
|
gr.Markdown("# 🚀 Enterprise AI Agent") |
|
|
gr.Markdown("Search the web, create docs, and generate PowerPoints instantly.") |
|
|
|
|
|
chatbot = gr.Chatbot(height=500, show_label=False) |
|
|
with gr.Row(): |
|
|
msg = gr.Textbox(placeholder="E.g., 'Research AI trends and make a 5-slide PPT'", scale=9) |
|
|
submit = gr.Button("Send", variant="primary", scale=1) |
|
|
|
|
|
with gr.Row(): |
|
|
download_btn = gr.Button("📂 Download Generated File", visible=False) |
|
|
file_output = gr.File(label="Ready for Download", visible=False) |
|
|
|
|
|
self._latest_file_path = None |
|
|
|
|
|
def respond(message, chat_history): |
|
|
|
|
|
agent_response = self.agent.run(message) |
|
|
chat_history.append((message, str(agent_response))) |
|
|
|
|
|
|
|
|
self._latest_file_path = None |
|
|
show_download = False |
|
|
match = re.search(r"(/tmp/[\w\d_-]+/generated_[\w-]+\.(docx|pdf|pptx))", str(agent_response)) |
|
|
|
|
|
if match: |
|
|
self._latest_file_path = match.group(1) |
|
|
show_download = True |
|
|
|
|
|
return chat_history, gr.update(visible=show_download), gr.update(visible=False) |
|
|
|
|
|
def prepare_download(): |
|
|
if self._latest_file_path and os.path.exists(self._latest_file_path): |
|
|
return gr.update(value=self._latest_file_path, visible=True) |
|
|
return gr.update(visible=False) |
|
|
|
|
|
msg.submit(respond, [msg, chatbot], [chatbot, download_btn, file_output]) |
|
|
submit.click(respond, [msg, chatbot], [chatbot, download_btn, file_output]) |
|
|
download_btn.click(prepare_download, [], [file_output]) |
|
|
|
|
|
return interface |
|
|
|
|
|
if __name__ == "__main__": |
|
|
CustomGradioUI(agent).launch() |