Update app.py

app.py

@@ -9,18 +9,12 @@ 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"
-)
+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 = [
@@ -33,8 +27,8 @@ functions_list = [
                 "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
+                        "type": "string",
+                        "description": "Search query (e.g., 'deep learning', 'quantum computing')"
                     },
                     "max_results": {
                         "type": "integer",
@@ -69,27 +63,9 @@ After receiving the results back from a function (formatted as {{"name": functio
 
 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
-    """
+def search_arxiv_papers(query: str, max_results: int = 5, sort_by: str = 'relevance') -> Dict:
     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,
@@ -100,12 +76,8 @@ def search_arxiv_papers(
         }
         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 = {
@@ -118,16 +90,12 @@ def search_arxiv_papers(
                 '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',
@@ -136,7 +104,6 @@ def search_arxiv_papers(
 
 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,
 }
 
@@ -149,8 +116,6 @@ class LLM:
         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:
@@ -163,18 +128,15 @@ class LLM:
             "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")
@@ -193,7 +155,6 @@ def form_chat_prompt(message_history, functions=functions_dict.keys()):
     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:
@@ -207,21 +168,17 @@ def check_assistant_response_for_tool_calls(response):
     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:
@@ -239,13 +196,10 @@ def restore_message_history(full_history):
     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)
@@ -253,10 +207,8 @@ def iterate_chat(llm, sampling_params, full_history):
             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 = {
@@ -275,58 +227,41 @@ def iterate_chat(llm, sampling_params, full_history):
             }
             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})
+def respond(message, chat_history, system_message, max_tokens, temperature, top_p):
+    if chat_history is None:
+        chat_history = []
+    full_history = chat_history.copy()
+    full_history.append({"role": "user", "content": message})
+    sampling_params = {
+        "temperature": temperature,
+        "top_p": top_p,
+        "max_tokens": max_tokens,
+        "stop_token_ids": [128001, 128008, 128009, 128006],
+    }
     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],
-}
+    chat_history.append((message, assistant_answer))
+    return chat_history
 
 # Initialize LLM
 llm = LLM(max_model_len=8096)
 
-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
-    )
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+    ],
+)
 
-    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
-    )
+if __name__ == "__main__":
     demo.launch()