!pip install gradio grandalf huggingface-hub langchain langchain-community langchain-core langchain-openai langgraph langgraph-checkpoint langgraph-checkpoint-postgres langgraph-checkpoint-sqlite langsmith openai psycopg pydantic pydantic_core tiktoken langchain-huggingface pymupdf4llm

from google.colab import userdata
import os
os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')

from pydantic import BaseModel, Field
from typing import TypedDict, Annotated, List, Dict, Any, Type
import requests
from langchain_core.tools import BaseTool

from bs4 import BeautifulSoup
import pymupdf4llm
import sys
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage, AIMessage, ChatMessage

import requests
import ast
import operator

from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI

import os
from uuid import uuid4
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.memory import MemorySaver
from psycopg import Connection
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint

from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type
)

import openai
from dotenv import load_dotenv
_ = load_dotenv()

class AcademicPaperSearchInput(BaseModel):
    topic: str = Field(..., description="The topic to search for academic papers on")
    max_results: int = Field(20, description="Maximum number of results to return")

class AcademicPaperSearchTool(BaseTool):
    args_schema: type = AcademicPaperSearchInput  # Explicit type annotation
    name: str = Field("academic_paper_search_tool", description="Tool for searching academic papers")
    description: str = Field("Queries an academic papers API to retrieve relevant articles based on a topic")

    def __init__(self, name: str = "academic_paper_search_tool",
                 description: str = "Queries an academic paper API to retrieve relevant articles based on a topic"):
        super().__init__()
        self.name = name
        self.description = description

    def _run(self, topic: str, max_results: int) -> List[Dict[str, Any]]:
        # Query an external academic API like arXiv, Semantic Scholar, or CrossRef
        search_results = self.query_academic_api(topic, max_results)
        # testing = search_results[0]['text'][:100]

        return search_results

    async def _arun(self, topic: str, max_results: int) -> List[Dict[str, Any]]:
        raise NotImplementedError("Async version not implemented")

    def query_academic_api(self, topic: str, max_results: int) -> List[Dict[str, Any]]:
        base_url = "https://api.semanticscholar.org/graph/v1/paper/search"
        params = {
            "query": topic,
            "limit": max_results, # max_results
            "fields": "title,abstract,authors,year,openAccessPdf",
            "openAccessPdf" : True
        }
        try:
            while True:
                try:
                    response = requests.get(base_url, params=params)
                    print(response)

                    if response.status_code == 200:
                        papers = response.json().get("data", [])
                        formatted_results = [
                            {
                                "title"     : paper.get("title"),
                                "abstract"  : paper.get("abstract"),
                                "authors"   : [author.get("name") for author in paper.get("authors", [])],
                                "year"      : paper.get("year"),
                                "pdf"       : paper.get("openAccessPdf"),
                            }
                            for paper in papers
                        ]

                        return formatted_results
                except:
                    # raise ValueError(f"Failed to fetch papers: {response.status_code} - {response.text}")
                    print((f"Failed to fetch papers: {response.status_code} - {response.text}. Trying Again..."))
        except KeyboardInterrupt:
            print("\nOperation cancelled by user")
            sys.exit(0)  # Clean exit

planner_prompt = '''You are an academic researcher that is planning to write a systematic review of Academic and Scientific Research Papers.

A systematic review article typically includes the following components:
Title: The title should accurately reflect the topic being reviewed, and usually includes the words "a systematic review".
Abstract: A structured abstract with a short paragraph for each of the following: background, methods, results, and conclusion.
Introduction: Summarizes the topic, explains why the review was conducted, and states the review's purpose and aims.
Methods: Describes the methods used in the review.
Results: Presents the results of the review.
Discussion: Discusses the results of the review.
References: Lists the references used in the review.

Other important components of a systematic review include:
Scoping: A "trial run" of the review that helps shape the review's method and protocol.
Meta-analysis: An optional component that uses statistical methods to combine and summarize the results of multiple studies.
Data extraction: A central component where data is collected and organized for analysis.
Assessing the risk of bias: Helps establish transparency of evidence synthesis results.
Interpreting results: Involves considering factors such as limitations, strength of evidence, biases, and implications for future practice or research.
Literature identification: An important component that sets the data to be analyzed.

With this in mind, only create an outline plan based on the topic. Don't search anything, just set up the planning.
'''

research_prompt = '''You are an academic researcher that is searching Academic and Scientific Research Papers.

You will be given a project plan. Based on the project plan, generate 5 queries that you will use to search the papers.

Send the queries to the academic_paper_search_tool as a tool call.
'''

decision_prompt = '''You are an academic researcher that is searching Academic and Scientific Research Papers.

You will be given a project plan and a list of articles.

Based on the project plan and articles provided, you must choose a maximum of 3 to investigate that are most relevant to that plan.

IMPORTANT: You must return ONLY a JSON array of the PDF URLs with no additional text or explanation. Your entire response should be in this exact format:

[
    "url1",
    "url2",
    "url3",
    ...
]

Do not include any other text, explanations, or formatting.'''

analyze_paper_prompt = '''You are an academic researcher trying to understand the details of scientific and academic research papers.

You must look through the text provided and get the details from the Abstract, Introduction, Methods, Results, and Conclusions.
If you are in an Abstract section, just give me the condensed thoughts.
If you are in an Introduction section, give me a concise reason on why the research was done.
If you are in a Methods section, give me low-level details of the approach. Analyze the math and tell me what it means.
If you are in a Results section, give me low-level relevant objective statistics. Tie it in with the methods
If you are in a Conclusions section, give me the fellow researcher's thoughts, but also come up with a counter-argument if none are given.

Remember to attach the other information to the top:
    Title : <title>
    Year : <year>
    Authors : <author1, author2, etc.>
    URL : <pdf url>
    TLDR Analysis:
        <your analysis>
'''

########################################################
abstract_prompt = '''You are an academic researcher that is writing a systematic review of Academic and Scientific Research Papers.
You are tasked with writing the Abstract section of the paper based on the systematic outline and the analyses given.
Make the abstract no more than 100 words.
'''

introduction_prompt = '''You are an academic researcher that is writing a systematic review of Academic and Scientific Research Papers.
You are tasked with writing the Introduction section of the paper based on the systematic outline and the analyses given.
Make sure it is thorough and covers information in all the papers.
'''

methods_prompt = '''You are an academic researcher that is writing a systematic review of Academic and Scientific Research Papers.
You are tasked with writing the Methods section of the paper based on the systematic outline and the analyses given.
Make sure it is thorough and covers information in all the papers. Draw on the differences and similarities in approaches in each paper.
'''

results_prompt = '''You are an academic researcher that is writing a systematic review of Academic and Scientific Research Papers.
You are tasked with writing the Results section of the paper based on the systematic outline and the analyses given.
Make sure it is thorough and covers information in all the papers. If there are results to compare among papers, please do so.
'''

conclusions_prompt = '''You are an academic researcher that is writing a systematic review of Academic and Scientific Research Papers.
You are tasked with writing the Conclusions section of the paper based on the systematic outline and the analyses given.
Make sure it is thorough and covers information in all the papers.
Draw on the conclusions from other papers, and what you might think the future of the research holds.
'''

references_prompt = '''You are an academic researcher that is writing a systematic review of Academic and Scientific Research Papers.
You are tasked with writing the References section of the paper based on the systematic outline and the analyses given.
Construct an APA style references list
'''
#########################################################
critique_draft_prompt = """You are an academic researcher deciding whether or not a systematic review should be published.
Generate a critique and recommendations for the author's submission or generate a query to get more papers.

If you think just a revision needs to be made, provide detailed recommendations, including requests for length, depth, style.
If you think the paper is good as is, just end with the draft unchanged.
"""
# If you think the write-up needs more papers, generate a search query and ask only for 2 additional articles.


revise_draft_prompt = """You are an academic researcher that is revising a systematic review that is about to be published.
Given the paper below, revise it following the recommendations given.

Return the revised paper with the implemented recommended changes.
"""

def reduce_messages(left: list[AnyMessage], right: list[AnyMessage]) -> list[AnyMessage]:
    # assign ids to messages that don't have them
    for message in right:
        if not message.id:
            message.id = str(uuid4())
    # merge the new messages with the existing messages
    merged = left.copy()
    for message in right:
        for i, existing in enumerate(merged):
            # replace any existing messages with the same id
            if existing.id == message.id:
                merged[i] = message
                break
        else:
            # append any new messages to the end
            merged.append(message)
    return merged

class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], reduce_messages]
    systematic_review_outline : str
    last_human_index : int
    papers : Annotated[List[str], operator.add] ## papers downloaded
    analyses: Annotated[List[Dict], operator.add]  # Store analysis results
    combined_analysis: str  # Final combined analysis

    title: str
    abstract : str
    introduction : str
    methods : str
    results : str
    conclusion : str
    references : str

    draft : str
    revision_num : int
    max_revisions : int

def process_input(state: AgentState):
    max_revision = 2
    messages = state.get('messages', [])

    last_human_index = len(messages) - 1
    for i in reversed(range(len(messages))):
        if isinstance(messages[i], HumanMessage):
            last_human_index = i
            break

    return {"last_human_index": last_human_index, "max_revisions" : max_revision, "revision_num" : 1}

def get_relevant_messages(state: AgentState) -> List[AnyMessage]:
    '''
    Don't get tool call messages for AI from history.
    Get state from everything up to the most recent human message
    '''
    messages = state['messages']
    filtered_history = []
    for message in messages:
        if isinstance(message, HumanMessage) and message.content!="":
            filtered_history.append(message)
        elif isinstance(message, AIMessage) and message.content!="" and message.response_metadata['finish_reason']=="stop":
            filtered_history.append(message)
    last_human_index = state['last_human_index']
    return filtered_history[:-1] + messages[last_human_index:]

def plan_node(state: AgentState):
    print("PLANNER")
    relevant_messages = get_relevant_messages(state)
    messages = [SystemMessage(content=planner_prompt)] + relevant_messages
    response = model.invoke(messages, temperature=temperature)
    print(response)
    print()
    return {"systematic_review_outline" : [response]}

def research_node(state: AgentState):
    print("RESEARCHER")
    review_plan = state['systematic_review_outline']
    messages = [SystemMessage(content=research_prompt)] + review_plan
    response = model.invoke(messages, temperature=temperature)
    print(response)
    print()
    return {"messages" : [response]}

def take_action(state: AgentState):
    ''' Get last message from agent state.
    If we get to this state, the language model wanted to use a tool.
    The tool calls attribute will be attached to message in the Agent State. Can be a list of tool calls.
    Find relevant tool and invoke it, passing in the arguments
    '''
    print("GET SEARCH RESULTS")
    last_message = state["messages"][-1]

    if not hasattr(last_message, 'tool_calls') or not last_message.tool_calls:
        return {"messages": state['messages']}

    results = []
    for t in last_message.tool_calls:
        print(f'Calling: {t}')

        if not t['name'] in tools: # check for bad tool name
            print("\n ....bad tool name....")
            result = "bad tool name, retry" # instruct llm to retry if bad
        else:
            # pass in arguments for tool call
            result = tools[t['name']].invoke(t['args'])

        # append result as a tool message
        results.append(ToolMessage(tool_call_id = t['id'], name=t['name'], content=str(result)))

    return {"messages" : results} # langgraph adding to state in between iterations

def decision_node(state: AgentState):
    print("DECISION-MAKER")
    review_plan = state['systematic_review_outline']
    relevant_messages = get_relevant_messages(state)
    messages = [SystemMessage(content=decision_prompt)] + review_plan + relevant_messages
    response = model.invoke(messages, temperature=temperature)
    print(response)
    print()
    return {"messages" : [response]}

def article_download(state: AgentState):
    print("DOWNLOAD PAPERS")
    last_message = state["messages"][-1]

    try:
        # Handle different types of content
        if isinstance(last_message.content, str):
            urls = ast.literal_eval(last_message.content)
        else:
            urls = last_message.content

        filenames = []
        for url in urls:
            try:
                response = requests.get(url)
                response.raise_for_status()

                # Create a papers directory if it doesn't exist
                if not os.path.exists('data'):
                    os.makedirs('data')

                # Generate a filename from the URL
                filename = f"data/{url.split('/')[-1]}"
                if not filename.endswith('.pdf'):
                    filename += '.pdf'

                # Save the PDF
                with open(filename, 'wb') as f:
                    f.write(response.content)

                filenames.append({"paper" : filename})
                print(f"Successfully downloaded: {filename}")

            except Exception as e:
                print(f"Error downloading {url}: {str(e)}")
                continue

        # Return AIMessage instead of raw strings
        return {
            "papers": [
                AIMessage(
                    content=filenames,
                    response_metadata={'finish_reason': 'stop'}
                )
            ]
        }

    except Exception as e:
        # Return error as AIMessage
        return {
            "messages": [
                AIMessage(
                    content=f"Error processing downloads: {str(e)}",
                    response_metadata={'finish_reason': 'error'}
                )
            ]
        }

def paper_analyzer(state: AgentState):
    print("ANALYZE PAPERS")
    analyses=""
    for paper in state['papers'][-1].content:
        md_text = pymupdf4llm.to_markdown(f"./{paper['paper']}")
        messages = [
            SystemMessage(content=analyze_paper_prompt),
            HumanMessage(content=md_text)
        ]

        model = ChatOpenAI(model='gpt-4o')
        response = model.invoke(messages, temperature=0.1)
        print(response)
        analyses+=response.content
    return {
        "analyses": [analyses]
    }

def _make_api_call(model, messages, temperature=0.1):
    @retry(
        retry=retry_if_exception_type(openai.RateLimitError),
        wait=wait_exponential(multiplier=1, min=4, max=60),
        stop=stop_after_attempt(5)
    )
    def _call():
        try:
            return model.invoke(messages, temperature=temperature)
        except openai.RateLimitError as e:
            print(f"Rate limit reached. Waiting before retry... ({e})")
            raise
    return _call()


def write_abstract(state: AgentState):
    print("WRITE ABSTRACT")
    review_plan = state['systematic_review_outline']
    analyses = state['analyses']
    messages = [SystemMessage(content=abstract_prompt)] + review_plan + analyses
    model = ChatOpenAI(model='gpt-4o-mini')
    response = _make_api_call(model, messages)
    print(response)
    print()
    return {"abstract" : [response]}

def write_introduction(state: AgentState):
    print("WRITE INTRODUCTION")
    review_plan = state['systematic_review_outline']
    analyses = state['analyses']
    messages = [SystemMessage(content=introduction_prompt)] + review_plan + analyses
    model = ChatOpenAI(model='gpt-4o-mini')
    response = _make_api_call(model, messages)
    print(response)
    print()
    return {"introduction" : [response]}

def write_methods(state: AgentState):
    print("WRITE METHODS")
    review_plan = state['systematic_review_outline']
    analyses = state['analyses']
    messages = [SystemMessage(content=methods_prompt)] + review_plan + analyses
    model = ChatOpenAI(model='gpt-4o-mini')
    response = _make_api_call(model, messages)
    print(response)
    print()
    return {"methods" : [response]}

def write_results(state: AgentState):
    print("WRITE RESULTS")
    review_plan = state['systematic_review_outline']
    analyses = state['analyses']
    messages = [SystemMessage(content=results_prompt)] + review_plan + analyses
    model = ChatOpenAI(model='gpt-4o-mini')
    response = _make_api_call(model, messages)
    print(response)
    print()
    return {"results" : [response]}

def write_conclusion(state: AgentState):
    print("WRITE CONCLUSION")
    review_plan = state['systematic_review_outline']
    analyses = state['analyses']
    messages = [SystemMessage(content=conclusions_prompt)] + review_plan + analyses
    model = ChatOpenAI(model='gpt-4o-mini')
    response = _make_api_call(model, messages)
    print(response)
    print()
    return {"conclusion" : [response]}

def write_references(state: AgentState):
    print("WRITE REFERENCES")
    review_plan = state['systematic_review_outline']
    analyses = state['analyses']
    messages = [SystemMessage(content=references_prompt)] + review_plan + analyses
    model = ChatOpenAI(model='gpt-4o-mini')
    response = _make_api_call(model, messages)
    print(response)
    print()
    return {"references" : [response]}

def aggregator(state: AgentState):
    print("AGGREGATE")
    abstract = state['abstract'][-1].content
    introduction = state['introduction'][-1].content
    methods = state['methods'][-1].content
    results = state['results'][-1].content
    conclusion = state['conclusion'][-1].content
    references = state['references'][-1].content

    messages = [
            SystemMessage(content="Make a title for this systematic review based on the abstract. Write it in markdown."),
            HumanMessage(content=abstract)
        ]
    title = model.invoke(messages, temperature=0.1).content

    draft = title + "\n\n" + abstract + "\n\n" + introduction + "\n\n" + methods + "\n\n" + results + "\n\n" + conclusion + "\n\n" + references

    return {"draft" : [draft]}

def critique(state:AgentState):
    print("CRITIQUE")
    draft = state["draft"]
    review_plan = state['systematic_review_outline']

    messages = [SystemMessage(content=critique_draft_prompt)] + review_plan + draft
    response = model.invoke(messages, temperature=temperature)
    print(response)

    # every critique is a call for revision
    return {'messages' : [response], "revision_num": state.get("revision_num", 1) + 1}

def paper_reviser(state: AgentState):
    print("REVISE PAPER")
    critique = state["messages"][-1].content
    draft = state["draft"]

    messages = [SystemMessage(content=revise_draft_prompt)] + [critique] + draft
    response = model.invoke(messages, temperature=temperature)
    print(response)

    return {'draft' : [response]}

def exists_action(state: AgentState):
    '''
    Determines whether to continue revising, end, or search for more articles
    based on the critique and revision count
    '''
    print("DECIDING WHETHER TO REVISE, END, or SEARCH AGAIN")

    if state["revision_num"] > state["max_revisions"]:
        return "final_draft"

    # # Get the latest critique
    critique = state['messages'][-1]
    print(critique)

    # Check if the critique response has any tool calls
    if hasattr(critique, 'tool_calls') and critique.tool_calls:
        # The critique suggests we need more research
        return True
    else:
        # No more research needed, proceed with revision
        return "revise"

def final_draft(state: AgentState):
    print("FINAL DRAFT")
    return {"draft" : state['draft']}

graph = StateGraph(AgentState)
graph.add_node("process_input", process_input)
graph.add_node("planner", plan_node)
graph.add_node("researcher", research_node)
graph.add_node("search_articles", take_action)
graph.add_node("article_decisions", decision_node)
graph.add_node("download_articles", article_download)
graph.add_node("paper_analyzer", paper_analyzer)

graph.add_node("write_abstract", write_abstract)
graph.add_node("write_introduction", write_introduction)
graph.add_node("write_methods", write_methods)
graph.add_node("write_results", write_results)
graph.add_node("write_conclusion", write_conclusion)
graph.add_node("write_references", write_references)

graph.add_node("aggregate_paper", aggregator)
graph.add_node("critique_paper", critique)
graph.add_node("revise_paper", paper_reviser)
graph.add_node("final_draft", final_draft)

####################################
graph.add_edge("process_input", "planner")
graph.add_edge("planner", "researcher")
graph.add_edge("researcher", "search_articles")
graph.add_edge("search_articles", "article_decisions")
graph.add_edge("article_decisions", "download_articles")
graph.add_edge("download_articles", 'paper_analyzer')

graph.add_edge("paper_analyzer", "write_abstract")
graph.add_edge("paper_analyzer", "write_introduction")
graph.add_edge("paper_analyzer", "write_methods")
graph.add_edge("paper_analyzer", "write_results")
graph.add_edge("paper_analyzer", "write_conclusion")
graph.add_edge("paper_analyzer", "write_references")

graph.add_edge("write_abstract", "aggregate_paper")
graph.add_edge("write_introduction", "aggregate_paper")
graph.add_edge("write_methods", "aggregate_paper")
graph.add_edge("write_results", "aggregate_paper")
graph.add_edge("write_conclusion", "aggregate_paper")
graph.add_edge("write_references", "aggregate_paper")

graph.add_edge("aggregate_paper", 'critique_paper')

graph.add_conditional_edges(
    "critique_paper",
    exists_action,
    {"final_draft": "final_draft",
     "revise": "revise_paper",
     True: "search_articles"}
)

graph.add_edge("revise_paper", "critique_paper")
graph.add_edge("final_draft", END)

graph.set_entry_point("process_input") ## "llm"

checkpointer = MemorySaver()
graph = graph.compile(checkpointer=checkpointer)

topic= "diffusion models for music generation"
thread_id = "test18"
temperature=0.1
papers_tool = AcademicPaperSearchTool()
tooling = [papers_tool]
model=ChatOpenAI(model='gpt-4o-mini') # gpt-4o-mini
tools = {t.name: t for t in tooling} if tooling else {}
model = model.bind_tools(tooling) if tools else model

agent_input = {"messages" : [HumanMessage(content=topic)]}
thread_config = {"configurable" : {"thread_id" : thread_id}}
result = graph.invoke(agent_input, thread_config)

from IPython.display import Image, display, Markdown
final_paper=result['draft'][-1].content
display(Markdown(final_paper))