RAG usando Vertex AI Search e Gemini (PT2)

A2_search.py

# Copyright 2024 Google LLC.
# SPDX-License-Identifier: Apache-2.0

# Search request
def search(query):

    client = discoveryengine.SearchServiceClient()

    serving_config = client.serving_config_path(
        project=PROJECT_ID,
        location=SEARCH_ENGINE_LOCATION,
        data_store=SEARCH_ENGINE_ID,
        serving_config=SEARCH_ENGINE_SERVING_CONFIG_ID,
    )

    request = discoveryengine.SearchRequest(
        serving_config=serving_config,
        query=query, page_size=5
    )

    response = client.search(request)
    return response.results
    
    # RAG using Vertex AI and Gemini
def qa(question):
  # Search from Vertex AI Agent Builder
  search_result = search(question)

  # Get documents
  context = "\n".join([ json.dumps(result.document.derived_struct_data["extractive_answers"][0]["content"]) for result in search_result])
  sources = "\n".join([ result.document.derived_struct_data["link"] + " - Página: " +\
                       json.dumps(result.document.derived_struct_data["extractive_answers"][0]["pageNumber"])  for result in search_result])

  # Modify the prompt according to your use case
  prompt = '''
Você é uma especialista em contratos de plano de saúde que responde perguntas gerais de acordo com o contexto.
Contexto: {}\n
Pergunta: {}\n
Resposta:'''

  # Generate answer
  responses = model.generate_content(
      prompt.format(context, question),
      generation_config={
        "max_output_tokens": 2048,
        "temperature": 0.5
      },
      safety_settings={
        generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
        generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
        generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
        generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
      },
      stream=True,
  )
  for response in responses:
    print(response.text, end="")

  print(sources)
  
# Make your question
qa("YOUR QUESTION")