app.py

import os
import textwrap
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
import pubmed_parser
import requests

VERBOSE_SHELL = True
ENDPOINT_URL = "https://api.hyperbolic.xyz/v1"
OAI_API_KEY = os.environ['HYPERBOLIC_XYZ_API_KEY']
PERPLEXITY_API_KEY = os.environ["PERPLEXITY_API_KEY"]
MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"

def lgs(log_string):
    if VERBOSE_SHELL:
        print(log_string)

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

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

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

def system_prompt_format(function_descriptions,function_jsons):
    return """Cutting Knowledge Date: December 2023
Today Date: """ + todays_date_string + """

You are a helpful assistant with tool calling capabilities.

""" + "\n".join(function_descriptions) + """
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 exact 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.

""" + "\n".join([json.dumps(d,indent=2) for d in function_jsons]) + """

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 build_sys_prompt(tool_objects):
    function_descriptions = [t.system_prompt_paragraph for t in tool_objects]
    function_jsons = [t.json_definition_of_function for t in tool_objects]
    return system_prompt_format(function_descriptions,function_jsons)

class ToolBase:
    def __init__(self,
            programmatic_name: str,
            natural_name: str,
            active_voice_description_of_capability: str,
            passive_voice_description_of_function: str,
            prescriptive_conditional: str,
            input_params: Dict[str, Dict],
            required_params: Optional[List[str]] = None,
        ):
        self.json_name = programmatic_name
        self.json_description = passive_voice_description_of_function
        self.json_definition_of_function = {
            "type": "function",
                "function": {
                    "name": self.json_name,
                    "description": self.json_description,
                    "parameters": {
                        "type": "object",
                        "properties": input_params,
                        "required": required_params,
                    }
                }
            }
        self.system_prompt_paragraph = active_voice_description_of_capability + " " + prescriptive_conditional
    def actual_function(self, **kwargs):
        raise NotImplementedError("Subclasses must implement this method.")

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 = 'https://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)

        time.sleep(1)

        return {
            'status': 'success',
            'total_results': len(papers),
            'papers': papers
        }

    except Exception as e:
        return {
            'status': 'error',
            'message': str(e)
        }

class ArxivSearchTool(ToolBase):
    def __init__(self):
        super().__init__(
            programmatic_name="search_arxiv_papers",
            natural_name="arXiv Paper Search",
            active_voice_description_of_capability="You can search for academic papers on arXiv.",
            passive_voice_description_of_function="a service that searches and retrieves academic papers from arXiv based on various criteria",
            prescriptive_conditional="When given a research topic or paper query, you should call the search_arxiv_papers function to find relevant papers.",
            input_params={
                "query": {
                    "type": "string",
                    "description": "Search query (e.g., 'deep learning', 'quantum computing')"
                },
                "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_params=["query"],
        )

    def actual_function(self, **kwargs):
        """
        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
        """
        return search_arxiv_papers(**kwargs)

arxiv_tool = ArxivSearchTool()

def get_snp_info(rsid):
    base_url = "https://api.ncbi.nlm.nih.gov/variation/v0/"
    result = {"rsid": rsid, "error": "No data found"}

    # Fetch RefSNP data
    snp_url = f"{base_url}refsnp/{rsid}"
    response = requests.get(snp_url)

    if response.status_code != 200:
        return {"error": f"Failed to retrieve data for rs{rsid}"}

    data = response.json()

    # Extract useful information
    result = {
        "create_date": data.get("create_date", "Unknown"),
        "last_update_date": data.get("last_update_date", "Unknown"),
        "genes": [],
        "hgvs": [],
        "spdi": [],
        "clinical_significance": [],
        "frequency_data": {},
    }

    # Extract gene associations
    primary_data = data.get("primary_snapshot_data", {})
    if "allele_annotations" in primary_data:
        for annotation in primary_data["allele_annotations"]:
            for gene in annotation.get("assembly_annotation", []):
                for gene_info in gene.get("genes", []):
                    result["genes"].append(gene_info.get("locus", "Unknown"))

    # Extract HGVS notation
    for placement in primary_data.get("placements_with_allele", []):
        for allele in placement.get("alleles", []):
            if "hgvs" in allele:
                result["hgvs"].append(allele["hgvs"])
            if "spdi" in allele.get("allele", {}):
                spdi_data = allele["allele"]["spdi"]
                spdi_notation = f"{spdi_data['seq_id']}:{spdi_data['position']}:{spdi_data['deleted_sequence']}:{spdi_data['inserted_sequence']}"
                result["spdi"].append(spdi_notation)

    # Extract clinical significance from ClinVar
    for annotation in primary_data.get("allele_annotations", []):
        for clinical in annotation.get("clinical", []):
            result["clinical_significance"].extend([str(s)[:600] for s in clinical.get("clinical_significances", [])])

    # Fetch ALFA frequency data
    freq_url = f"{base_url}refsnp/{rsid}/frequency"
    freq_response = requests.get(freq_url)

    if freq_response.status_code == 200:
        freq_data = freq_response.json().get("results", {})
        for key, value in freq_data.items():
            if "counts" in value:
                result["frequency_data"] = value["counts"]
                break
    citations = data.get("citations", [])[:6]
    lgs("citations: " + str(citations))
    result["citations"] = [pubmed_parser.parse_xml_web(c, sleep=0.5, save_xml=False,) for c in citations]
    lgs("full citations data: " + str(result["citations"]))
    return result

class NIHRefSNPTool(ToolBase):
    def __init__(self):
        super().__init__(
            programmatic_name="search_nih_refsnp",
            natural_name="NIH RefSNP Searcher",
            active_voice_description_of_capability=(
                "You can search for refSNP data on the NIH Variation API."
            ),
            passive_voice_description_of_function=(
                "a service that retrieves refSNP data from the NIH Variation API "
                "based on a provided SNP identifier"
            ),
            prescriptive_conditional=(
                "When given a refSNP identifier (e.g., 'rs79220014'), "
                "you should call the search_nih_refsnp function "
                "to find its associated data."
            ),
            input_params={
                "snp": {
                    "type": "string",
                    "description": "The refSNP identifier (e.g., 'rs79220014')"
                }
            },
            required_params=["snp"],
        )

    def actual_function(self, **kwargs):
        return get_snp_info(kwargs["snp"][2:])

nih_ref_snp_tool=NIHRefSNPTool()

def query_perplexity(query: str, api_key: str) -> Dict:
    """
    Query the Perplexity API for research information using default settings.
    
    Args:
        query: The research question or topic to query
        api_key: Your Perplexity API key
    
    Returns:
        Dictionary containing the response from Perplexity API
    """
    if not api_key:
        return {
            "status": "error",
            "message": "API key is required. Please provide your Perplexity API key."
        }
    
    url = "https://api.perplexity.ai/chat/completions"
    
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    # Using Perplexity's recommended defaults
    messages = [
        {
            "role": "system",
            "content": "You are a helpful AI research assistant. Provide accurate, detailed information with relevant citations."
        },
        {
            "role": "user",
            "content": query
        }
    ]
    
    payload = {
        "model": "sonar",
        "messages": messages,
        "temperature": 0.2,
        "top_p": 0.9,
    }
    
    try:
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()  # Raise exception for 4XX/5XX responses
        data = response.json()
        
        # Format the response for easier consumption
        result = {
            "status": "success",
            "content": data["choices"][0]["message"]["content"] if data.get("choices") else None,
            "citations": data.get("citations", []),
        }
            
        return result
    
    except requests.exceptions.RequestException as e:
        return {
            "status": "error",
            "message": f"API request failed: {str(e)}"
        }
    except Exception as e:
        return {
            "status": "error",
            "message": f"An unexpected error occurred: {str(e)}"
        }

class PerplexityQueryTool(ToolBase):
    def __init__(self):
        super().__init__(
            programmatic_name="query_perplexity",
            natural_name="Perplexity Research Assistant",
            active_voice_description_of_capability="You can search for up-to-date information using the Perplexity API.",
            passive_voice_description_of_function="a service that retrieves information from the web using the Perplexity AI model with proper citations.",
            prescriptive_conditional="When you need to find current information or answer research questions, you should use the query_perplexity function.",
            input_params={
                "query": {
                    "type": "string",
                    "description": "The research question or topic to query."
                },
                "api_key": {
                    "type": "string",
                    "description": "Your Perplexity API key."
                },
            },
            required_params=["query", "api_key"],
        )

    def actual_function(self, **kwargs):
        """
        Query the Perplexity API for research information.
        """
        return query_perplexity(**kwargs)

perplexity_tool = PerplexityQueryTool()

tool_objects_list = [arxiv_tool, nih_ref_snp_tool, perplexity_tool]
system_prompt = build_sys_prompt(tool_objects_list)
functions_dict = {t.json_name: t.actual_function for t in tool_objects_list}

print(system_prompt)

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 = MODEL_NAME

    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."""

    full_prompt = (
        ROLE_HEADER.format(role="system")
        + "\n\n"
        + system_prompt
        + 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"]
    tool_return = None
    
    if tool_name == arxiv_tool.json_name:
        query = tool_parameters["query"]
        max_results = tool_parameters.get("max_results", 5)
        sort_by = tool_parameters.get("sort_by", "relevance")
        search_results = arxiv_tool.actual_function(query=query, max_results=max_results, sort_by=sort_by)
        tool_return = {"name": arxiv_tool.json_name, "return": search_results}
        
    elif tool_name == nih_ref_snp_tool.json_name:
        snp = tool_parameters["snp"]
        search_results = nih_ref_snp_tool.actual_function(snp=snp)
        tool_return = {"name": nih_ref_snp_tool.json_name, "return": search_results}
        
    elif tool_name == perplexity_tool.json_name:
        query = tool_parameters["query"]
        # Using the environment variable instead of requiring it as a parameter
        search_results = perplexity_tool.actual_function(query=query, api_key=PERPLEXITY_API_KEY)
        tool_return = {"name": perplexity_tool.json_name, "return": search_results}
        
    else:
        raise ValueError(f"Unknown tool name: {tool_name}")
        
    lgs("TOOL: " + str(tool_return))
    return tool_return

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()
        lgs("ASSISTANT: " + assistant_response.replace("\n", "\\n"))
        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 = []
    lgs("USER: " + user_message.replace("\n", "\\n"))
    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

llm = LLM(max_model_len=32000)

lgs("STARTING NEW CHAT")
with gr.Blocks() as demo:
    gr.Markdown(f"<h2>Search/Arxiv/SNP Multi-tool Calling Bot</h2>")
    chat_state = gr.State([])
    chatbot = gr.Chatbot(label="Chat with the multi-tool bot")
    user_input = gr.Textbox(
        lines=1,
        placeholder="Type your message here...",
    )
    gr.Examples([
            [
                "What is the current weather in Åfjord?",
            ],
            [
                "List some papers about humor in LLMs",
            ],
            [
                "What does this SNP do?: rs429358",
            ]
        ],
        inputs=[user_input],
        label="Examples",
    )
    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()
    share_url = demo.share_url