Update app.py

app.py

@@ -9,12 +9,18 @@ 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 = [
@@ -27,8 +33,8 @@ functions_list = [
                 "type": "object",
                 "properties": {
                     "query": {
-                        "type": "string",
-                        "description": "Search query (e.g., 'deep learning', 'quantum computing')"
+                        "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",
@@ -63,9 +69,27 @@ 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:
+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,
@@ -76,8 +100,12 @@ def search_arxiv_papers(query: str, max_results: int = 5, sort_by: str = 'releva
         }
         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 = {
@@ -90,12 +118,16 @@ def search_arxiv_papers(query: str, max_results: int = 5, sort_by: str = 'releva
                 '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',
@@ -104,6 +136,7 @@ def search_arxiv_papers(query: str, max_results: int = 5, sort_by: str = 'releva
 
 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,
 }
 
@@ -116,6 +149,8 @@ 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:
@@ -128,15 +163,18 @@ 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")
@@ -155,6 +193,7 @@ 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:
@@ -168,17 +207,21 @@ 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:
@@ -196,10 +239,13 @@ 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)
@@ -207,8 +253,10 @@ 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 = {
@@ -227,41 +275,58 @@ 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 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],
-    }
+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((message, assistant_answer))
-    return chat_history
+    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)
 
-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)"),
-    ],
-)
+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...",
+    )
 
-if __name__ == "__main__":
+    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()