from google.cloud import discoveryengine_v1beta as discoveryengine

# Based on official documentation: https://cloud.google.com/generative-ai-app-builder/docs/answer
def answer_query_sample(
    project_id: str,
    location: str,
    engine_id: str,
) -> discoveryengine.AnswerQueryResponse:
    client = discoveryengine.ConversationalSearchServiceClient()
    
    # The full resource name of the Search serving config
    serving_config = f"projects/{project_id}/locations/{location}/collections/default_collection/engines/{engine_id}/servingConfigs/default_serving_config"
    
    # Optional: Options for answer phase
    answer_generation_spec = discoveryengine.AnswerQueryRequest.AnswerGenerationSpec(
        ignore_adversarial_query=False,
        ignore_non_answer_seeking_query=False,
        ignore_low_relevant_content=False,
        model_spec=discoveryengine.AnswerQueryRequest.AnswerGenerationSpec.ModelSpec(
            model_version="gemini-2.0-flash-001/answer_gen/v1",
        ),
        prompt_spec=discoveryengine.AnswerQueryRequest.AnswerGenerationSpec.PromptSpec(
            preamble="Give a detailed answer.",
        ),
        include_citations=True,
        answer_language_code="en",
    )
    
    # Initialize request argument(s)
    request = discoveryengine.AnswerQueryRequest(
        serving_config=serving_config,
        query=discoveryengine.Query(text="What is Vertex AI Search?"),
        session=None,
        answer_generation_spec=answer_generation_spec,
        user_pseudo_id="user-pseudo-id",
    )
    
    # Make the request
    response = client.answer_query(request)
    return response

# Test execution
try:
    response = answer_query_sample(
        project_id="project-id",
        location="global",
        engine_id="data-store-id"
    )
    print("Success:", response)
except Exception as e:
    print("Error:", str(e))