Update app.py

app.py

@@ -5,118 +5,147 @@ import streamlit as st
 from dotenv import load_dotenv
 from duckduckgo_search import DDGS
 import subprocess
+import re
 
-# Load API
+# Load environment
 load_dotenv()
 API_KEY = os.getenv("OPENROUTER_API_KEY")
 
-def call_llm(prompt):
-    url = "https://openrouter.ai/api/v1/chat/completions"
-    headers = {"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
-    data = {
-        "model": "google/gemma-3-1b-it:free",
-        "messages": [{"role": "user", "content": prompt}]
+# ========== UTILITY FUNCTIONS ==========
+def sanitize_filename(title):
+    return re.sub(r'[^\w\s-]', '', title).replace(" ", "_")[:50]
+
+def escape_latex(text):
+    replacements = {
+        '&': r'\&', '%': r'\%', '$': r'\$', '#': r'\#',
+        '_': r'\_', '{': r'\{', '}': r'\}', '~': r'\textasciitilde{}',
+        '^': r'\^{}', '\\': r'\textbackslash{}',
     }
-    res = requests.post(url, headers=headers, json=data)
-    if res.status_code == 200:
-        return res.json()['choices'][0]['message']['content']
-    else:
-        return f"❌ Error: {res.status_code} - {res.text}"
-
-def get_arxiv_papers(topic):
-    search = arxiv.Search(query=topic, max_results=5, sort_by=arxiv.SortCriterion.SubmittedDate)
-    papers = []
-    for result in search.results():
-        paper = {
-            "title": result.title,
-            "summary": result.summary,
-            "url": result.pdf_url
+    for original, replacement in replacements.items():
+        text = text.replace(original, replacement)
+    return text
+
+# ========== LLM INTERACTION ==========
+def call_llm(prompt):
+    try:
+        url = "https://openrouter.ai/api/v1/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {API_KEY}",
+            "Content-Type": "application/json"
+        }
+        data = {
+            "model": "google/gemma-3-1b-it:free",
+            "messages": [{"role": "user", "content": prompt}]
         }
-        papers.append(paper)
-    return papers
+        response = requests.post(url, headers=headers, json=data)
+        response.raise_for_status()
+        return response.json()['choices'][0]['message']['content']
+    except Exception as e:
+        return f"❌ LLM Error: {str(e)}"
+
+# ========== FETCH RESEARCH ==========
+def fetch_arxiv_papers(topic):
+    results = arxiv.Search(query=topic, max_results=5, sort_by=arxiv.SortCriterion.SubmittedDate)
+    return [{
+        "title": result.title,
+        "summary": result.summary,
+        "url": result.pdf_url
+    } for result in results.results()]
 
 def fetch_image_url(query):
-    with DDGS() as ddgs:
-        results = list(ddgs.images(query, max_results=1))
-        if results:
-            return results[0]['image']
-        return None
-
-def create_latex(paper_title, full_content, image_url, file_name="research_paper.tex"):
-    latex = f"""
+    try:
+        with DDGS() as ddgs:
+            results = list(ddgs.images(query, max_results=1))
+            if results:
+                return results[0]['image']
+    except:
+        pass
+    return None
+
+# ========== PDF EXPORT ==========
+def create_latex(title, content, image_url):
+    title_safe = sanitize_filename(title)
+    content = escape_latex(content)
+    tex = f"""
 \\documentclass[12pt]{{article}}
-\\usepackage{{graphicx}}
 \\usepackage[margin=1in]{{geometry}}
+\\usepackage{{graphicx}}
 \\usepackage{{hyperref}}
-\\title{{\\textbf{{{paper_title}}}}}
+\\title{{\\textbf{{{escape_latex(title)}}}}}
 \\date{{\\today}}
 \\begin{{document}}
 \\maketitle
-
-\\begin{{figure}}[h]
-    \\centering
-    \\includegraphics[width=0.6\\textwidth]{{image.jpg}}
-    \\caption{{Relevant diagram for {paper_title}}}
-\\end{{figure}}
-
-{full_content.replace('\\n', '\\\\ \\n')}
-
-\\end{{document}}
 """
-    with open(file_name, "w") as f:
-        f.write(latex)
 
-    # Save image
     if image_url:
         img_data = requests.get(image_url).content
         with open("image.jpg", "wb") as img:
             img.write(img_data)
+        tex += """
+\\begin{figure}[h]
+    \\centering
+    \\includegraphics[width=0.7\\textwidth]{image.jpg}
+    \\caption{Auto-fetched Diagram}
+\\end{figure}
+"""
+
+    tex += f"{content}\n\\end{{document}}"
+
+    tex_file = f"{title_safe}.tex"
+    with open(tex_file, "w", encoding="utf-8") as f:
+        f.write(tex)
 
-    subprocess.run(["pdflatex", file_name], stdout=subprocess.DEVNULL)
-    return "research_paper.pdf"
+    subprocess.run(["pdflatex", tex_file], stdout=subprocess.DEVNULL)
+    return f"{title_safe}.pdf"
 
-# Streamlit UI
-st.set_page_config("MTech AI Research Generator", layout="wide")
-st.title("📚 Deep Research Paper Generator (AI + arXiv + PDF Export)")
+# ========== STREAMLIT UI ==========
+st.set_page_config("AI Research Assistant", layout="wide")
+st.title("🧪 AI-Powered Research Assistant")
 
 topic = st.text_input("🔍 Enter your research topic:")
 
 if topic:
-    with st.spinner("Fetching latest papers from arXiv..."):
-        papers = get_arxiv_papers(topic)
+    with st.spinner("🔎 Fetching relevant arXiv papers..."):
+        papers = fetch_arxiv_papers(topic)
         summaries = "\n\n".join([f"Title: {p['title']}\nSummary: {p['summary']}" for p in papers])
-        st.success("✅ Papers Retrieved!")
-        st.subheader("📄 Top 5 arXiv Papers")
-        for p in papers:
-            st.markdown(f"**{p['title']}**  \n[PDF Link]({p['url']})  \n> {p['summary'][:300]}...\n")
-
-    with st.spinner("Analyzing literature & generating novel paper..."):
-        gap = call_llm(f"Identify current research gaps based on these papers:\n\n{summaries}")
-        idea = call_llm(f"Propose a novel research direction based on this:\n\n{gap}")
-        paper = call_llm(f"""Write a detailed academic paper on "{topic}" using the idea:
+
+    st.subheader("📄 Recent arXiv Papers")
+    for p in papers:
+        st.markdown(f"**{p['title']}**\n[🔗 PDF Link]({p['url']})\n> {p['summary'][:300]}...")
+
+    with st.spinner("🧠 Analyzing gaps and proposing ideas..."):
+        gaps = call_llm(f"You're a top AI researcher. Read the summaries and find research gaps:\n\n{summaries}")
+        idea = call_llm(f"Based on the gaps below, propose a novel research idea:\n\n{gaps}")
+        paper = call_llm(f"""Write a detailed academic paper titled: "{topic}".
+Incorporate the following novel idea:
 {idea}
+
 Structure:
-1. Abstract
-2. Introduction
-3. Related Work
-4. Methodology
-5. Experiments
-6. Results
-7. Conclusion
-8. References""")
-
-    st.subheader("🧠 Novel Idea")
+- Abstract
+- Introduction
+- Related Work
+- Methodology
+- Experiments
+- Results & Discussion
+- Conclusion
+- References
+
+Write clearly in LaTeX-ready format.""")
+
+    st.subheader("💡 Novel Research Idea")
     st.markdown(idea)
 
-    st.subheader("📝 Full Paper")
-    st.markdown(paper)
+    st.subheader("📃 Full Generated Paper")
+    st.text_area("Academic Paper", paper, height=600)
 
-    st.subheader("🖼️ Diagram")
-    image_url = fetch_image_url(topic)
-    if image_url:
-        st.image(image_url, caption="Auto-selected diagram")
+    with st.spinner("🖼️ Fetching diagram..."):
+        image_url = fetch_image_url(topic)
+        if image_url:
+            st.image(image_url, caption="Auto-fetched relevant diagram")
 
     if st.button("📥 Export to PDF"):
-        pdf = create_latex(topic, paper, image_url)
-        with open(pdf, "rb") as f:
-            st.download_button("Download PDF", f, file_name=f"{topic.replace(' ', '_')}.pdf")
+        with st.spinner("📄 Generating PDF..."):
+            pdf_file = create_latex(topic, paper, image_url)
+            with open(pdf_file, "rb") as f:
+                st.download_button("Download Paper as PDF", f, file_name=pdf_file)
+