Create app.py

app.py

@@ -0,0 +1,339 @@
+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=8096)
+
+def build_demo():
+    """Build the Gradio interface."""
+    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("Check Balance")
+        send_button.click(
+            fn=user_conversation,
+            inputs=[user_input, chatbot, chat_state],
+            outputs=[user_input, chatbot, chat_state],
+            queue=False
+        )
+
+    return demo
+
+if __name__ == "__main__":
+    demo = build_demo()
+    demo.launch()