Spaces:
Runtime error
Runtime error
File size: 5,313 Bytes
0cac9cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """Application entrypoint for QueryQuest.
Author: mohamedgamal04
"""
import sys
import time
from pathlib import Path
from dotenv import load_dotenv
from rich.console import Console
from rich.prompt import Prompt
from rich.panel import Panel
from rich import box
from .chat_session import run_chat_session
from .cli import normalize_prompt_input, parse_args, print_banner
from .config import EXCEL_INFO_FORMAT_VERSION, PROVIDERS_BY_NAME, SYSTEM_PROMPT
from .excel.context import (
build_excel_files_info,
describe_excel_snapshot_changes,
format_excel_context,
get_excel_snapshot,
list_excel_files,
normalize_excel_dir,
)
from .setup_flow import run_setup
from .state import load_state, save_state
CONSOLE = Console()
SNAPSHOT_CHECK_INTERVAL_SECONDS = 2.0
MAX_EXCEL_CONTEXT_CHARS = 5000
def main() -> None:
"""Initialize configuration and run the interactive chat session."""
load_dotenv()
print_banner(CONSOLE)
options = parse_args(sys.argv[1:], CONSOLE)
state = load_state()
if options.setup or state is None:
state = run_setup(CONSOLE)
excel_default = state.get("excel_dir") if state else None
excel_dir_input = Prompt.ask(
"Excel files location",
default=excel_default or str(normalize_excel_dir()),
console=CONSOLE,
).strip()
excel_dir = normalize_excel_dir(Path(excel_dir_input))
config = PROVIDERS_BY_NAME[state["provider"]]
CONSOLE.print(
Panel(
f"[bold]Provider[/bold]: [cyan]{config.name}[/cyan]\n"
f"[bold]Model[/bold]: [bright_white]{state['model']}[/bright_white]",
border_style="bright_cyan",
box=box.ROUNDED,
padding=(0, 1),
)
)
excel_info_cache = {
"signature": state.get("excel_signature") if state else None,
"info": state.get("excel_info") if state else None,
"format_version": state.get("excel_info_format_version") if state else None,
"snapshot": None,
"change_note": "",
"last_snapshot_check_at": 0.0,
}
is_info_stale = excel_info_cache["format_version"] != EXCEL_INFO_FORMAT_VERSION
if not excel_info_cache["signature"] or not excel_info_cache["info"] or is_info_stale:
excel_info_cache["info"], excel_info_cache["signature"], excel_info_cache["snapshot"] = build_excel_files_info(excel_dir)
excel_info_cache["format_version"] = EXCEL_INFO_FORMAT_VERSION
save_state(
state["provider"],
state["api_key"],
state["model"],
excel_dir=str(excel_dir),
excel_signature=excel_info_cache["signature"],
excel_info=excel_info_cache["info"],
excel_info_format_version=EXCEL_INFO_FORMAT_VERSION,
)
else:
_, excel_info_cache["snapshot"] = get_excel_snapshot(excel_dir)
excel_info_cache["last_snapshot_check_at"] = time.monotonic()
def get_system_prompt() -> str:
"""Return a system prompt enriched with current Excel metadata.
The prompt cache is refreshed when the file snapshot signature changes.
"""
now = time.monotonic()
should_refresh_snapshot = (now - excel_info_cache["last_snapshot_check_at"]) >= SNAPSHOT_CHECK_INTERVAL_SECONDS
if should_refresh_snapshot or excel_info_cache["snapshot"] is None:
current_signature, current_snapshot = get_excel_snapshot(excel_dir)
excel_info_cache["last_snapshot_check_at"] = now
if current_signature != excel_info_cache["signature"]:
# Rebuild context when workbook set/content changed on disk.
previous_snapshot = excel_info_cache["snapshot"]
info, signature, snapshot = build_excel_files_info(excel_dir)
excel_info_cache["signature"] = signature
excel_info_cache["info"] = info
excel_info_cache["format_version"] = EXCEL_INFO_FORMAT_VERSION
excel_info_cache["snapshot"] = snapshot
excel_info_cache["change_note"] = describe_excel_snapshot_changes(previous_snapshot, snapshot)
save_state(
state["provider"],
state["api_key"],
state["model"],
excel_dir=str(excel_dir),
excel_signature=signature,
excel_info=info,
excel_info_format_version=EXCEL_INFO_FORMAT_VERSION,
)
elif excel_info_cache["snapshot"] is None:
excel_info_cache["snapshot"] = current_snapshot
prompt_parts = [SYSTEM_PROMPT]
if excel_info_cache["change_note"]:
prompt_parts.append(excel_info_cache["change_note"])
if excel_info_cache["info"]:
excel_context = format_excel_context(excel_info_cache["info"])
if len(excel_context) > MAX_EXCEL_CONTEXT_CHARS:
excel_context = f"{excel_context[:MAX_EXCEL_CONTEXT_CHARS]}\n... [excel context truncated for speed]"
prompt_parts.append(excel_context)
return "\n\n".join(part for part in prompt_parts if part)
prompt, prompt_flag_only = normalize_prompt_input(options.prompt or "")
if prompt_flag_only:
CONSOLE.print("Please provide prompt text after -p/--prompt.")
prompt = ""
if options.setup and not prompt:
CONSOLE.print(
Panel(
"[green]Setup complete.[/green] Type prompts below.\n"
"Use [cyan]QQ -q[/cyan] or [cyan]QQ -quit[/cyan] to exit.",
border_style="bright_green",
box=box.ROUNDED,
padding=(0, 1),
)
)
run_chat_session(
console=CONSOLE,
provider_name=config.name,
provider_base_url=config.base_url,
model_name=state["model"],
api_key=state["api_key"],
initial_prompt=prompt,
system_prompt_provider=get_system_prompt,
excel_file_count_provider=lambda: len(list_excel_files(excel_dir)),
excel_dir=excel_dir,
)
if __name__ == "__main__":
main()
|