app.py

from openai import OpenAI
import google.generativeai as genai

import os
import requests
import json
import gradio as gr
import time
import re
#export GRADIO_DEBUG=1

GENAI_API = "gemini" # or "openai"

def search_inspire(query, size=10):
    """
    Search INSPIRE HEP database using fulltext search

    Args:
        query (str): Search query
        size (int): Number of results to return
    """
    base_url = "https://inspirehep.net/api/literature"
    params = {
        "q": query,
        "size": size,
        "format": "json"
    }

    response = requests.get(base_url, params=params)
    return response.json()

def format_reference(metadata):
  output = f"{', '.join(author.get('full_name', '') for author in metadata.get('authors', []))} "
  output += f"({metadata.get('publication_info', [{}])[0].get('year', 'N/A')}). "
  output += f"*{metadata.get('titles', [{}])[0].get('title', 'N/A')}*. "
  output += f"DOI: {metadata.get('dois', [{}])[0].get('value', 'N/A') if metadata.get('dois') else 'N/A'}. "
  output += f"[INSPIRE record {metadata['control_number']}](https://inspirehep.net/literature/{metadata['control_number']})"
  output += "\n\n"
  return output

def format_results(results):
    """Print formatted search results"""
    output = ""
    for i, hit in enumerate(results['hits']['hits']):
        metadata = hit['metadata']
        output += f"**[{i}]** "
        output += format_reference(metadata)
    return output

def results_context(results):
  """ Prepare a context from the results for the LLM """
  context = ""
  for i, hit in enumerate(results['hits']['hits']):
    metadata = hit['metadata']
    context += f"Result [{i}]\n\n"
    context += f"Title: {metadata.get('titles', [{}])[0].get('title', 'N/A')}\n\n"
    context += f"Abstract: {metadata.get('abstracts', [{}])[0].get('value', 'N/A')}\n\n"
  return context

def user_prompt(query, context):
  """ Generate a prompt for the LLM """
  prompt = f"""
  QUERY: {query}

  CONTEXT:

  {context}

  ANSWER:

  """
  return prompt

def llm_expand_query(query):
  """ Expands a query to variations of fulltext searches """


  prompt = f"""
    Expand this query into a the query format used for a search
    over the INSPIRE HEP database. Propose alternatives of the query to
    maximize the recall and join those variantes using OR operators. 
    Just provide the expanded query, without explanations.

    Example of query:
    how far are black holes?

    Expanded query:
    "how far are black holes" OR "distance from black holes" OR
    "distances to black holes" OR "measurement of distance to black
    holes"  OR "remoteness of black holes" OR "distance to black
    holes"  OR "how far are singularities" OR "distance to
    singularities" OR "distances to event horizon" OR "distance
    from Schwarzschild radius" OR "black hole distance"

    Query: {query}

    Expanded query:
  """
  
  if GENAI_API == "openai":
    response = client.chat.completions.create(
      model="gpt-4o-mini",
      messages=[
        {
          "role": "user",
          "content": [
            {
              "type": "text",
              "text": prompt
            }
          ]
        }
      ],
      response_format={
        "type": "text"
      },
      temperature=0,
      max_tokens=2048,
      top_p=1,
      frequency_penalty=0,
      presence_penalty=0
    )

    return response.choices[0].message.content
  else:
     response = genai.GenerativeModel("gemini-1.5-flash").generate_content(prompt)
     return response.text

def llm_generate_answer(prompt):
  """ Generate a response from the LLM """
  
  system_desc = """You are part of a Retrieval Augmented Generation system
              (RAG) and are asked with a query and a context of results. Generate an
              answer substantiated by the results provided and citing them using
              their index when used to provide an answer text. Do not put two or more
              references together (ex: use [1][2] instead of [1, 2] or [1][2][3] instead of [1, 2, 3]). Do not generate an answer
              that cannot be entailed from cited abstract, so all paragraphs should cite a
              search result. End the answer with the query and a brief answer as
              summary of the previous discussed results. Do not consider results
              that are not related to the query and, if no specific answer can be
              provided, assert that in the brief answer."""
  
  if GENAI_API == "openai":
     
    response = client.chat.completions.create(
      model="gpt-4o-mini",
      messages=[
        {
          "role": "system",
          "content": [
            {
              "type": "text",
              "text": system_desc
            }
          ]
        },
        {
          "role": "user",
          "content": [
            {
              "type": "text",
              "text": prompt
            }
          ]
        }
      ],
      response_format={
        "type": "text"
      },
      temperature=0,
      max_tokens=2048,
      top_p=1,
      frequency_penalty=0,
      presence_penalty=0
    )

    return response.choices[0].message.content
  
  else:
     response = genai.GenerativeModel("gemini-1.5-flash").generate_content(system_desc + "\n\n" + prompt)
     return response.text

def clean_refs(answer, results):
  """ Clean the references from the answer """

  # Find references
  unique_ordered = []
  for match in re.finditer(r'\[(\d+)\]', answer):
    ref_num = int(match.group(1))
    if ref_num not in unique_ordered:
        unique_ordered.append(ref_num)

  # Filter references
  new_i = 1
  new_results = ""
  for i, hit in enumerate(results['hits']['hits']):
    if i not in unique_ordered:
      continue
    metadata = hit['metadata']
    new_results += f"**[{new_i}]** "
    new_results += format_reference(metadata)
    new_i += 1

  new_i = 1
  for i in unique_ordered:
    answer = answer.replace(f"[{i}]", f" **[__NEW_REF_ID_{new_i}]**")
    new_i += 1
  answer = answer.replace("__NEW_REF_ID_", "")

  return answer, new_results

def search(query, progress=gr.Progress()):
    time.sleep(1)
    progress(0, desc="Expanding query...")
    expanded_query = llm_expand_query(query)
    progress(0.25, desc="Searching INSPIRE HEP...")
    results = search_inspire(expanded_query)
    progress(0.50, desc="Generating answer...")
    context = results_context(results)
    prompt = user_prompt(query, context)
    answer = llm_generate_answer(prompt)
    new_answer, references = clean_refs(answer, results)
    progress(1, desc="Done!")

    #json_str = json.dumps(results['hits']['hits'][0]['metadata'], indent=4)
    return "**Answer**:\n\n" + new_answer +"\n\n**References**:\n\n" + references #+ "\n\n <pre>\n" + json_str + "</pre>"

# ----------- MAIN ------------------------------------------------------------

if GENAI_API == "openai":
  client = OpenAI()
else:
   genai.configure(api_key=os.getenv('GEMINI_API_KEY'))

with gr.Blocks() as demo:
    gr.Markdown("# Feynbot on INSPIRE HEP Search")
    gr.Markdown("""Specialized academic search tool that combines traditional 
                database searching with AI-powered query expansion and result 
                synthesis, focused on High Energy Physics research papers.""")
    with gr.Row():
        with gr.Column():
            query = gr.Textbox(label="Search Query")
            search_btn = gr.Button("Search")
            examples = gr.Examples([["Which one is closest star?"], ["In which particles does the Higgs Boson decay to?"]], query)
            with gr.Row():
              gr.HTML("<a href='https://sinai.ujaen.es'><img src='https://sinai.ujaen.es/sites/default/files/SINAI%20-%20logo%20tx%20azul%20%5Baf%5D.png' width='200'></img></a>")
              gr.HTML("<a href='https://www.ujaen.es'><img src='https://diariodigital.ujaen.es/sites/default/files/general/logo-uja.svg' width='180'></img></a>")
        with gr.Column():
           results = gr.Markdown("Answer will appear here...", label="Search Results", )
        search_btn.click(fn=search, inputs=query, outputs=results, api_name="search", show_progress=True)
    


demo.launch()
#print(search("how far are black holes?"))