app.py

import os
import datetime
import json
import gradio as gr
from openai import OpenAI
import urllib.request
import feedparser
import time
from typing import Dict, List, Optional

ENDPOINT_URL = "https://api.hyperbolic.xyz/v1"

OAI_API_KEY = os.getenv('HYPERBOLIC_XYZ_KEY')

VERBOSE_SHELL = True

todays_date_string = datetime.date.today().strftime("%d %B %Y")


NAME_OF_SERVICE = "arXiv Paper Search"
DESCRIPTION_OF_SERVICE = (
    "a service that searches and retrieves academic papers from arXiv based on various criteria"
)
PAPER_SEARCH_FUNCTION_NAME = "search_arxiv_papers"

functions_list = [
    {
        "type": "function",
        "function": {
            "name": PAPER_SEARCH_FUNCTION_NAME,
            "description": DESCRIPTION_OF_SERVICE,
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string", # function names for AI agents should be chosen carefully to avoid confusion
                        "description": "Search query (e.g., 'deep learning', 'quantum computing')" # descriptions help the AI agent's LLM backend understand the function
                    },
                    "max_results": {
                        "type": "integer",
                        "description": "Maximum number of results to return (default: 5)",
                        "optional": True
                    },
                    "sort_by": {
                        "type": "string",
                        "description": "Sort criteria (e.g., 'relevance', 'lastUpdatedDate', 'submittedDate')",
                        "optional": True
                    }
                },
                "required": ["query"]
            }
        }
    }
]

system_prompt = """Cutting Knowledge Date: December 2023
Today Date: """ + todays_date_string + """

You are a helpful assistant with tool calling capabilities.

You can search for academic papers on arXiv. When given a research topic or paper query, you should call the search_arxiv_papers function to find relevant papers.
If you choose to use one of the following functions, respond with a JSON for a function call with its proper arguments that best answers the given prompt.

Your tool request should be in the format {{"name": function name, "parameters": dictionary of argument name and its value}}. Do not use variables. Just a two-key dictionary, starting with the function name, followed by a dictionary of parameters.

{{functions}}

After receiving the results back from a function (formatted as {{"name": function name, "return": returned data after running function}}) formulate your response to the user. If the information needed is not found in the returned data, either attempt a new function call, or inform the user that you cannot answer based on your available knowledge. The user cannot see the function results. You have to interpret the data and provide a response based on it.

If the user request does not necessitate a function call, simply respond to the user's query directly."""

def search_arxiv_papers(
    query: str,
    max_results: int = 5,
    sort_by: str = 'relevance'
) -> Dict:
    """
    Search for papers on arXiv using their API.
    
    Args:
        query: Search query string
        max_results: Maximum number of results to return (default: 5)
        sort_by: Sorting criteria (default: 'relevance')
        
    Returns:
        Dictionary containing search results and metadata
    """
    try:
        # Construct the search query
        search_query = f'all:{query}'
            
        # Construct the API URL
        base_url = 'http://export.arxiv.org/api/query?'
        params = {
            'search_query': search_query,
            'start': 0,
            'max_results': max_results,
            'sortBy': sort_by,
            'sortOrder': 'descending'
        }
        query_string = '&'.join([f'{k}={urllib.parse.quote(str(v))}' for k, v in params.items()])
        url = base_url + query_string
        
        # Make the API request
        response = urllib.request.urlopen(url)
        feed = feedparser.parse(response.read().decode('utf-8'))
        
        # Process the results
        papers = []
        for entry in feed.entries:
            paper = {
                'id': entry.id.split('/abs/')[-1],
                'title': entry.title,
                'authors': [author.name for author in entry.authors],
                'summary': entry.summary,
                'published': entry.published,
                'link': entry.link,
                'primary_category': entry.tags[0]['term']
            }
            papers.append(paper)
            
        # Add a delay to respect API rate limits
        time.sleep(3)
        
        return {
            'status': 'success',
            'total_results': len(papers),
            'papers': papers
        }
        
    except Exception as e:
        return {
            'status': 'error',
            'message': str(e)
        }

functions_dict = {f["function"]["name"]: f for f in functions_list}
FUNCTION_BACKENDS = {
    #WALLET_CHECK_FUNCTION_NAME: check_wallet_balance,
    PAPER_SEARCH_FUNCTION_NAME: search_arxiv_papers,
}

EOT_STRING = "<|eot_id|>"
FUNCTION_EOT_STRING = "<|eom_id|>"
ROLE_HEADER = "<|start_header_id|>{role}<|end_header_id|>"

class LLM:
    def __init__(self, max_model_len: int = 4096):
        self.api_key = OAI_API_KEY
        self.max_model_len = max_model_len
        self.client = OpenAI(base_url=ENDPOINT_URL, api_key=self.api_key)
        #models_list = self.client.models.list()
        #self.model_name = models_list.data[0].id
        self.model_name = "meta-llama/Llama-3.3-70B-Instruct"

    def generate(self, prompt: str, sampling_params: dict) -> dict:
        completion_params = {
            "model": self.model_name,
            "prompt": prompt,
            "max_tokens": sampling_params.get("max_tokens", 2048),
            "temperature": sampling_params.get("temperature", 0.8),
            "top_p": sampling_params.get("top_p", 0.95),
            "n": sampling_params.get("n", 1),
            "stream": False,
        }
        
        if "stop" in sampling_params:
            completion_params["stop"] = sampling_params["stop"]
        if "presence_penalty" in sampling_params:
            completion_params["presence_penalty"] = sampling_params["presence_penalty"]
        if "frequency_penalty" in sampling_params:
            completion_params["frequency_penalty"] = sampling_params["frequency_penalty"]
            
        return self.client.completions.create(**completion_params)

def form_chat_prompt(message_history, functions=functions_dict.keys()):
    """Builds the chat prompt for the LLM."""
    functions_string = "\n\n".join([json.dumps(functions_dict[f], indent=4) for f in functions])
    full_prompt = (
        ROLE_HEADER.format(role="system")
        + "\n\n"
        + system_prompt.format(functions=functions_string)
        + EOT_STRING
    )
    for message in message_history:
        full_prompt += (
            ROLE_HEADER.format(role=message["role"])
            + "\n\n"
            + message["content"]
            + EOT_STRING
        )
    full_prompt += ROLE_HEADER.format(role="assistant")
    return full_prompt

def check_assistant_response_for_tool_calls(response):
    """Check if the LLM response contains a function call."""
    response = response.split(FUNCTION_EOT_STRING)[0].split(EOT_STRING)[0]
    for tool_name in functions_dict.keys():
        if f"\"{tool_name}\"" in response and "{" in response:
            response = "{" + "{".join(response.split("{")[1:])
            for _ in range(10):
                response = "}".join(response.split("}")[:-1]) + "}"
                try:
                    return json.loads(response)
                except json.JSONDecodeError:
                    continue
    return None

def process_tool_request(tool_request_data):
    """Process tool requests from the LLM."""
    tool_name = tool_request_data["name"]
    tool_parameters = tool_request_data["parameters"]

    if tool_name == PAPER_SEARCH_FUNCTION_NAME:
        query = tool_parameters["query"]
        max_results = tool_parameters.get("max_results", 5)
        sort_by = tool_parameters.get("sort_by", "relevance")
        search_results = FUNCTION_BACKENDS[tool_name](query, max_results, sort_by)
        return {"name": PAPER_SEARCH_FUNCTION_NAME, "results": search_results}

    return None

def restore_message_history(full_history):
    """Restore the complete message history including tool interactions."""
    restored = []
    for message in full_history:
        if message["role"] == "assistant" and "metadata" in message:
            tool_interactions = message["metadata"].get("tool_interactions", [])
            if tool_interactions:
                for tool_msg in tool_interactions:
                    restored.append(tool_msg)
                final_msg = message.copy()
                del final_msg["metadata"]["tool_interactions"]
                restored.append(final_msg)
            else:
                restored.append(message)
        else:
            restored.append(message)
    return restored

def iterate_chat(llm, sampling_params, full_history):
    """Handle conversation turns with tool calling."""
    tool_interactions = []

    for _ in range(10):
        prompt = form_chat_prompt(restore_message_history(full_history) + tool_interactions)
        output = llm.generate(prompt, sampling_params)
        
        if VERBOSE_SHELL:
            print(f"Input prompt: {prompt}")
            print("-" * 50)
            print(f"Model response: {output.choices[0].text}")
            print("=" * 50)
        if not output or not output.choices:
            raise ValueError("Invalid completion response")
            
        assistant_response = output.choices[0].text.strip()
        assistant_response = assistant_response.split(FUNCTION_EOT_STRING)[0].split(EOT_STRING)[0]

        tool_request_data = check_assistant_response_for_tool_calls(assistant_response)
        if not tool_request_data:
            final_message = {
                "role": "assistant",
                "content": assistant_response,
                "metadata": {
                    "tool_interactions": tool_interactions
                }
            }
            full_history.append(final_message)
            return full_history
        else:
            assistant_message = {
                "role": "assistant",
                "content": json.dumps(tool_request_data),
            }
            tool_interactions.append(assistant_message)
            tool_return_data = process_tool_request(tool_request_data)

            tool_message = {
                "role": "function",
                "content": json.dumps(tool_return_data)
            }
            tool_interactions.append(tool_message)

    return full_history

def user_conversation(user_message, chat_history, full_history):
    """Handle user input and maintain conversation state."""
    if full_history is None:
        full_history = []

    full_history.append({"role": "user", "content": user_message})
    updated_history = iterate_chat(llm, sampling_params, full_history)
    assistant_answer = updated_history[-1]["content"]
    chat_history.append((user_message, assistant_answer))

    return "", chat_history, updated_history

sampling_params = {
    "temperature": 0.8,
    "top_p": 0.95,
    "max_tokens": 512,
    "stop_token_ids": [128001,128008,128009,128006],
}

# Initialize LLM
llm = LLM(max_model_len=32000)

with gr.Blocks() as demo:
    gr.Markdown(f"<h2>{NAME_OF_SERVICE}</h2>")
    chat_state = gr.State([])
    chatbot = gr.Chatbot(label="Chat with the arXiv Paper Search Assistant")
    user_input = gr.Textbox(
        lines=1, 
        placeholder="Type your message here...",
    )

    user_input.submit(
        fn=user_conversation,
        inputs=[user_input, chatbot, chat_state],
        outputs=[user_input, chatbot, chat_state],
        queue=False
    )

    send_button = gr.Button("Send")
    send_button.click(
        fn=user_conversation,
        inputs=[user_input, chatbot, chat_state],
        outputs=[user_input, chatbot, chat_state],
        queue=False
    )
    demo.launch()