updated to work with Gemini

app.py

@@ -1,4 +1,6 @@
 from openai import OpenAI
+import google.generativeai as genai
+
 import os
 import requests
 import json
@@ -7,6 +9,8 @@ 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
@@ -70,97 +74,112 @@ def user_prompt(query, context):
 def llm_expand_query(query):
   """ Expands a query to variations of fulltext searches """
 
-  response = client.chat.completions.create(
-    model="gpt-4o-mini",
-    messages=[
-      {
-        "role": "user",
-        "content": [
-          {
-            "type": "text",
-            "text": f"""
-            Expand this query into a the query format used for a fulltext search
-            over the INSPIRE HEP database. Propose alternatives of the query to
-            maximize the recall and join those variantes using OR operators and
-            prepend each variant with the ft prefix. Just provide the expanded
-            query, without explanations.
-
-            Example of query:
-            how far are black holes?
-
-            Expanded query:
-            ft "how far are black holes" OR ft "distance from black holes" OR ft
-            "distances to black holes" OR ft "measurement of distance to black
-            holes"  OR ft "remoteness of black holes"  OR ft "distance to black
-            holes"  OR ft "how far are singularities"  OR ft "distance to
-            singularities"  OR ft "distances to event horizon"  OR ft "distance
-            from Schwarzschild radius" OR ft "black hole distance"
-
-            Query: {query}
-
-            Expanded query:
-            """
-          }
-        ]
-      }
-    ],
-    response_format={
-      "type": "text"
-    },
-    temperature=0,
-    max_tokens=2048,
-    top_p=1,
-    frequency_penalty=0,
-    presence_penalty=0
-  )
-
-  return response.choices[0].message.content
+
+  prompt = f"""
+    Expand this query into a the query format used for a fulltext search
+    over the INSPIRE HEP database. Propose alternatives of the query to
+    maximize the recall and join those variantes using OR operators and
+    prepend each variant with the ft prefix. Just provide the expanded
+    query, without explanations.
+
+    Example of query:
+    how far are black holes?
+
+    Expanded query:
+    ft "how far are black holes" OR ft "distance from black holes" OR ft
+    "distances to black holes" OR ft "measurement of distance to black
+    holes"  OR ft "remoteness of black holes"  OR ft "distance to black
+    holes"  OR ft "how far are singularities"  OR ft "distance to
+    singularities"  OR ft "distances to event horizon"  OR ft "distance
+    from Schwarzschild radius" OR ft "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 """
-
-  response = client.chat.completions.create(
-    model="gpt-4o-mini",
-    messages=[
-      {
-        "role": "system",
-        "content": [
-          {
-            "type": "text",
-            "text": """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]. 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."""
-          }
-        ]
+  
+  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]. 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"
       },
-      {
-        "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
+      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 """
@@ -209,7 +228,10 @@ def search(query, progress=gr.Progress()):
 
 # ----------- MAIN ------------------------------------------------------------
 
-client = OpenAI()
+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")
@@ -221,10 +243,11 @@ with gr.Blocks() as demo:
             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)
+            gr.HTML("<img src='https://sinai.ujaen.es/sites/default/files/SINAI%20-%20logo%20tx%20azul%20%5Baf%5D.png' width='200'></img>")
         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)
-    gr.HTML("<img src='https://sinai.ujaen.es/sites/default/files/SINAI%20-%20logo%20tx%20azul%20%5Baf%5D.png' width='200'></img>")
+    
 
 
 demo.launch()

requirements.txt

@@ -1,4 +1,5 @@
 gradio
 openai
 requests
-httpx<0.28
+httpx<0.28
+google-generativeai