updated google sdk calls

app.py

@@ -18,6 +18,87 @@ GENAI_API = os.getenv("GENAI_API", "gemini")
 
 # LLM_MODEL_NAME must be set in the environment
 
+def _extract_text_from_message(message):
+  """Extract plain text from a message entry used in this codebase.
+
+  Messages in this project often look like:
+    {"role": "user", "content": [{"type": "text", "text": "..."}]}
+
+  This helper normalizes that shape to a single string.
+  """
+  content = message.get("content")
+  if isinstance(content, list) and len(content) > 0:
+    first = content[0]
+    if isinstance(first, dict) and "text" in first:
+      return first.get("text", "")
+    return str(first)
+  if isinstance(content, dict) and "text" in content:
+    return content.get("text", "")
+  if isinstance(content, str):
+    return content
+  return str(content)
+
+
+def create_chat_response(messages, model, temperature=0, max_tokens=2048):
+  """Unified helper to produce a text response from either OpenAI or
+  Google's GenAI backends.
+
+  Returns a plain string with the assistant reply.
+  """
+  # OpenAI-style client: keep calling the same API
+  if GENAI_API == "openai":
+    response = client.chat.completions.create(
+      model=model,
+      messages=messages,
+      temperature=temperature,
+      max_tokens=max_tokens,
+      top_p=1,
+      frequency_penalty=0,
+      presence_penalty=0,
+    )
+    # Expect OpenAI-style response
+    try:
+      return response.choices[0].message.content
+    except Exception:
+      # Fallback: stringify
+      return str(response)
+
+  # Google GenAI path: convert messages to a single prompt and call
+  # the available model API (best-effort mapping).
+  prompt = "\n\n".join(f"{m.get('role','')}: {_extract_text_from_message(m)}" for m in messages)
+
+  # Try common modern GenAI SDK pattern: client.models.generate_content
+  try:
+    if hasattr(client, "models") and hasattr(client.models, "generate_content"):
+      # Use names similar to examples: contents and optional params
+      try:
+        resp = client.models.generate_content(model=model, contents=prompt, temperature=temperature, max_output_tokens=max_tokens)
+      except TypeError:
+        # Some versions may not accept those named args; try minimal call
+        resp = client.models.generate_content(model=model, contents=prompt)
+
+      # Response object often has `.text` or `.content`
+      text = getattr(resp, "text", None) or getattr(resp, "content", None)
+      if text is None:
+        return str(resp)
+      return text
+
+    # Older `google.generativeai` (legacy) had different surface; try a
+    # generous fallback: look for a top-level `generate` or `generate_text`.
+    if hasattr(client, "generate"):
+      resp = client.generate(prompt)
+      return getattr(resp, "text", str(resp))
+    if hasattr(client, "generate_text"):
+      resp = client.generate_text(prompt)
+      return getattr(resp, "text", str(resp))
+
+  except Exception as e:
+    # Surface the error with context to help debugging.
+    raise RuntimeError(f"GenAI model call failed: {e}")
+
+  raise RuntimeError("No suitable GenAI method found on `client`; please install/initialize supported SDK or set GENAI_API=openai")
+
+
 def search_inspire(query, size=10):
     """
     Search INSPIRE HEP database using fulltext search
@@ -80,16 +161,13 @@ def user_prompt(query, context):
 
 def llm_expand_query(query):
   """ Expands a query to variations of fulltext searches """
-
-  response = client.chat.completions.create(
-    model=MODEL_NAME,
-    messages=[
-      {
-        "role": "user",
-        "content": [
-          {
-            "type": "text",
-            "text": f"""
+  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
@@ -110,36 +188,24 @@ def llm_expand_query(query):
             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
+  return create_chat_response(messages=messages, model=MODEL_NAME, temperature=0, max_tokens=2048)
 
 
 def llm_generate_answer(prompt):
   """ Generate a response from the LLM """
-
-  response = client.chat.completions.create(
-    model=MODEL_NAME,
-    messages=[
-      {
-        "role": "system",
-        "content": [
-          {
-            "type": "text",
-            "text": """You are part of a Retrieval Augmented Generation system
+  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
@@ -149,30 +215,21 @@ def llm_generate_answer(prompt):
             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."""
-          }
-        ]
-      },
-      {
-        "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
+    {
+      "role": "user",
+      "content": [
+        {
+          "type": "text",
+          "text": prompt
+        }
+      ]
+    }
+  ]
+
+  return create_chat_response(messages=messages, model=MODEL_NAME, temperature=0, max_tokens=2048)
   
 
 def clean_refs(answer, results):