1  #!/usr/bin/env python3
  2  """Simple script to query the vector database with a question."""
  3  
  4  import argparse
  5  import logging
  6  import sys
  7  
  8  from src import Config
  9  from src.cli.constants import (
 10      EXIT_CODE_ERROR,
 11      SEPARATOR_CHAR,
 12      SEPARATOR_LENGTH,
 13  )
 14  from src.components import execute_query, initialize_rag_components
 15  from src.logger import Logger
 16  from src.pipeline import PipelineStatus, QueryContext
 17  
 18  
 19  def _log_query_info(
 20      query: str,
 21      user_role: str | None,
 22      role_mapping: dict[str, list[str]],
 23      logger: logging.Logger
 24  ) -> None:
 25      """Log query information and role mapping."""
 26      logger.info(SEPARATOR_CHAR * SEPARATOR_LENGTH)
 27      logger.info("Querying Vector Database")
 28      logger.info(SEPARATOR_CHAR * SEPARATOR_LENGTH)
 29      logger.info(f"Query: {query}")
 30      if user_role:
 31          logger.info(f"User role: {user_role}")
 32      if role_mapping:
 33          logger.info(f"Role mapping loaded: {len(role_mapping)} roles")
 34          if user_role and user_role in role_mapping:
 35              logger.info(f"Role '{user_role}' maps to tags: {role_mapping[user_role]}")
 36      logger.info("")
 37  
 38  
 39  def _display_results(
 40      context: QueryContext,
 41      logger: logging.Logger,
 42  ) -> None:
 43      """Display query results, citations, and retrieved documents."""
 44      logger.info("\n" + SEPARATOR_CHAR * SEPARATOR_LENGTH)
 45      logger.info("Final Answer")
 46      logger.info(SEPARATOR_CHAR * SEPARATOR_LENGTH)
 47      logger.info(context.llm_response or "(no response)")
 48  
 49      if context.citations:
 50          logger.info("\nSources:")
 51          for c in context.citations:
 52              cid = c.get("id")
 53              source = c.get("source")
 54              score = c.get("score")
 55              logger.info(f"  [{cid}] {source}  score={score}")
 56  
 57  def main():
 58      """Query the vector database with a question from command line arguments."""
 59      parser = argparse.ArgumentParser(
 60          description="Query vector database with role-aware access control from config"
 61      )
 62      parser.add_argument(
 63          "query",
 64          nargs="+",
 65          help="Search query string",
 66      )
 67      args = parser.parse_args()
 68  
 69      config = Config.get_config()
 70  
 71      Logger.setup(config)
 72      logger = logging.getLogger(__name__)
 73  
 74      query = " ".join(args.query)
 75  
 76      user_role = config.access_control.default_user_role
 77      role_mapping = config.access_control.get_role_mapping()
 78  
 79      _log_query_info(query, user_role, role_mapping, logger)
 80  
 81      try:
 82          components = initialize_rag_components()
 83  
 84          context = execute_query(
 85              components,
 86              query,
 87              user_role=user_role,
 88              role_mapping=role_mapping,
 89          )
 90  
 91          if context.status == PipelineStatus.FAILED:
 92              raise RuntimeError(f"Pipeline failed: {context.error}")
 93  
 94          _display_results(context, logger)
 95  
 96      except Exception as e:
 97          logger.error(f"✗ Error: {e}", exc_info=True)
 98          sys.exit(EXIT_CODE_ERROR)
 99  
100  if __name__ == "__main__":
101      main()