run_agent.py

"""This script runs the BigQuery optimization agent in a conversational mode."""
from bq_agent_app.agent import root_agent
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.genai import types
import uuid
import logging
from google.adk.artifacts import InMemoryArtifactService
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

APP_NAME = "bq_agent_app"
USER_ID = "biquery_user_101"

async def run_conversation(prompt: str):
  """Runs a conversation with the BigQuery optimization agent.

  Args:
      prompt: The user's prompt to the agent.

  Returns:
      The final response from the agent.
  """
  logger.info(f"Starting conversation with prompt: {prompt}")
  session_service = InMemorySessionService()
  artifact_service = InMemoryArtifactService()
  session_id = f"{APP_NAME}-{uuid.uuid4().hex[:8]}"
 
  runner = Runner(
      agent=root_agent,
      app_name=APP_NAME,
      session_service=session_service,
      artifact_service=artifact_service
  )
 
  await session_service.create_session(
       app_name=APP_NAME,
       user_id=USER_ID,
       session_id=session_id
  )
 
  try:
      # Run the agent and process the events as they are generated
      async for event in runner.run_async(
          user_id=USER_ID,
          session_id=session_id,
          new_message=types.Content(role='user', parts=[types.Part(text=prompt)])
      ):
          if event.is_final_response():
              if event.content and event.content.parts:
                  logger.info("Finished conversation.")
                  return event.content.parts[0].text
  except Exception as e:
      logger.error(f"Error in run_conversation: {e}", exc_info=True)
      return f"An error occurred: {e}"
    
  logger.info("Finished conversation.")
  return "Unable to retrieve final response."