enhance-research-agent

Enhanced research capabilities and improved UI

app.py

@@ -1,69 +1,142 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
 import datetime
 import requests
 import pytz
 import yaml
-from tools.final_answer import FinalAnswerTool
+from transformers import PerplexityForTextGeneration
 
+# Import additional tools
+from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
 
-# Below is an example of a tool that does nothing. Amaze us with your creativity !
+# Initialize Perplexity for research capabilities
+perplexity = PerplexityForTextGeneration()
+
+@tool
+def get_ai_research_papers(query: str) -> str:
+    """A tool that fetches relevant AI research papers using Perplexity API.
+    Args:
+        query: The search query for AI research papers.
+    """
+    try:
+        response = perplexity.search(
+            query=query,
+            max_results=8,
+            filters={"domain": "AI"}
+        )
+        if response:
+            papers = "\n".join([f"📄 {paper.title}" for paper in response])
+            return f"Found these relevant AI papers about {query}:\n{papers}"
+        return f"No relevant AI papers found for your query: {query}"
+    except Exception as e:
+        return f"Error fetching research papers: {str(e)}"
+
 @tool
-def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
-    #Keep this format for the description / args / args description but feel free to modify the tool
-    """A tool that does nothing yet 
+def summarize_paper(paper_title: str) -> str:
+    """A tool that summarizes an AI research paper.
     Args:
-        arg1: the first argument
-        arg2: the second argument
+        paper_title: Title of the paper to summarize.
     """
-    return "What magic will you build ?"
+    try:
+        try:
+            response = perplexity.search(
+                query=paper_title,
+                max_results=1
+            )
+            if response:
+                paper = response[0]
+                summary = paper.summary
+                return f"Summary of '{paper_title}':\n{summary}"
+        except IndexError:
+            return f"Could not find a paper titled '{paper_title}"
+    except Exception as e:
+        return f"Error summarizing paper: {str(e)}"
 
 @tool
-def get_current_time_in_timezone(timezone: str) -> str:
-    """A tool that fetches the current local time in a specified timezone.
+def get_citation(paper_title: str) -> str:
+    """A tool that generates a citation for an AI research paper.
     Args:
-        timezone: A string representing a valid timezone (e.g., 'America/New_York').
+        paper_title: Title of the paper to cite.
     """
     try:
-        # Create timezone object
-        tz = pytz.timezone(timezone)
-        # Get current time in that timezone
-        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
-        return f"The current local time in {timezone} is: {local_time}"
+        try:
+            response = perplexity.search(
+                query=paper_title,
+                max_results=1
+            )
+            if response:
+                paper = response[0]
+                citation = f"""{paper.title}
+Author(s): {paper.authors}
+Published in: {paper.venue}
+Year: {paper.year}
+DOI: {paper.doi if paper.doi else 'Not available}"""
+                return citation
+        except IndexError:
+            return f"Could not find a paper titled '{paper_title}"
     except Exception as e:
-        return f"Error fetching time for timezone '{timezone}': {str(e)}"
+        return f"Error generating citation: {str(e)}"
 
+@tool
+def explain_concept(concept: str) -> str:
+    """A tool that explains an AI-related concept in simple terms.
+    Args:
+        concept: The concept to explain.
+    """
+    try:
+        try:
+            response = perplexity.search(
+                query=f"explain {concept}",
+                max_results=1
+            )
+            if response:
+                explanation = response[0].summary
+                return f"Explanation of {concept}:\n{explanation}"
+        except IndexError:
+            return f"Could not find an explanation for '{concept}"
+    except Exception as e:
+        return f"Error explaining concept: {str(e)}"
 
 final_answer = FinalAnswerTool()
 
-# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
-# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 
-
+# Initialize the model
 model = HfApiModel(
-max_tokens=2096,
-temperature=0.5,
-model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
-custom_role_conversions=None,
+    max_tokens=2096,
+    temperature=0.5,
+    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
 )
 
-
-# Import tool from Hub
+# Load additional tools
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
-    
+
+# Create the enhanced research agent
 agent = CodeAgent(
     model=model,
-    tools=[final_answer], ## add your tools here (don't remove final answer)
+    tools=[
+        final_answer,
+        get_ai_research_papers,
+        summarize_paper,
+        get_citation,
+        explain_concept,
+        image_generation_tool,
+    ],
     max_steps=6,
     verbosity_level=1,
     grammar=None,
     planning_interval=None,
-    name=None,
-    description=None,
+    name="AI Research Assistant",
+    description="An AI-powered research assistant for exploring AI-related papers and concepts.",
     prompt_templates=prompt_templates
 )
 
-
-GradioUI(agent).launch()
+# Launch the enhanced UI
+GradioUI(
+    agent,
+    theme="peach",
+    title="AI Research Assistant",
+    description="Your go-to assistant for finding and understanding AI research papers.",
+    allow_flagging="never"
+).launch()