Update app.py

app.py

@@ -3,8 +3,10 @@ import datetime
 import requests
 import pytz
 import yaml
+import re
 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 !
@@ -18,6 +20,42 @@ def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return
     """
     return "What magic will you build ?"
 
+@tool
+def get_bibtex_citation(title: str) -> str:
+    """Retrieves the citation of a particular paper in the specified style
+    Args:
+        title: Title of the work to cite in the bibtex format"""
+    try:
+        search_url = f"https://api.crossref.org/works?query.title={title.replace(' ', '+')}&rows=1"
+        response = requests.get(search_url)
+        data = response.json()
+
+        if data['message']['total-results'] == 0:
+            return f"No results found for '{title}'"
+
+        paper = data['message']['items'][0]
+        doi = paper.get('DOI')
+
+        if not doi:
+            return f"No DOI found for paper: {paper.get('title', [''])[0]}"
+        
+        bibtex_url = f"https://www.doi2bib.org/bib/{doi}"
+        bibtex_response = requests.get(bibtex_url)
+        
+        if bibtex_response.status_code != 200:
+            return f"Error accessing doi2bib.org: Status code {bibtex_response.status_code}"
+        
+        html = bibtex_response.text
+        bibtex_match = re.search(r'<pre class="bibtex-code text-left">.*?<code>(.*?)</code>', html, re.DOTALL)
+
+        if bibtex_match:
+            bibtex = bibtex_match.group(1)
+            return bibtex.strip()
+        
+    except Exception as e:
+        return f"Error generating citation: {str(e)}"
+
+        
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """A tool that fetches the current local time in a specified timezone.
@@ -55,7 +93,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer], ## add your tools here (don't remove final answer)
+    tools=[final_answer, get_bibtex_citation], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,