suggestoor/main.py

from libs.ollama_client import ollama_chat, ollama_client
from libs.get_all_mans import get_all_mans
import os
from datetime import datetime, timezone
import sys
from colorama import init, Fore, Style

# Initialize colorama for cross-platform colored output
init()


def check_vector_store():
    """Check if vector store exists and is up to date"""
    if not os.path.exists("vector_store"):
        print(f"{Fore.YELLOW}Vector store not found{Style.RESET_ALL}")
        return False

    # Check if any man page is newer than vector store
    vector_store_time = datetime.fromtimestamp(
        os.path.getmtime("vector_store"), timezone.utc
    )
    mans_dir = "mans"

    if not os.path.exists(mans_dir):
        print(f"{Fore.YELLOW}Man pages directory not found{Style.RESET_ALL}")
        return False

    for man_file in os.listdir(mans_dir):
        if man_file.endswith(".man"):
            man_time = datetime.fromtimestamp(
                os.path.getmtime(os.path.join(mans_dir, man_file)), timezone.utc
            )
            if man_time > vector_store_time:
                print(
                    f"{Fore.YELLOW}Vector store outdated (newer man pages found){Style.RESET_ALL}"
                )
                return False

    return True


def initialize_knowledge_base():
    """Initialize or update the knowledge base with progress indicators"""
    if check_vector_store():
        print(f"{Fore.GREEN}✓ Vector store is up to date{Style.RESET_ALL}")
        return

    print(f"{Fore.CYAN}Initializing knowledge base...{Style.RESET_ALL}")
    try:
        db = get_all_mans()
        print(f"{Fore.GREEN}✓ Successfully created vector store{Style.RESET_ALL}")
    except Exception as e:
        print(f"{Fore.RED}Error creating vector store: {e}{Style.RESET_ALL}")
        sys.exit(1)


def main():
    print(f"{Fore.CYAN}Welcome to suggestoor!{Style.RESET_ALL}")

    # Initialize knowledge base
    initialize_knowledge_base()

    print(f"\n{Fore.GREEN}Ready to help! Type your questions below.{Style.RESET_ALL}")
    print(f"{Fore.YELLOW}(Press Ctrl+C to exit){Style.RESET_ALL}\n")

    while True:
        try:
            prompt = input(f"{Fore.CYAN}Which tool do you need? {Style.RESET_ALL}")
            client = ollama_client()
            print(f"{Fore.YELLOW}Thinking...{Style.RESET_ALL}")
            response = ollama_chat(client, prompt)
            print(f"\n{Fore.GREEN}Answer:{Style.RESET_ALL}")
            print(response.message.content)
            print()  # Empty line for readability
        except KeyboardInterrupt:
            print(f"\n{Fore.YELLOW}Goodbye!{Style.RESET_ALL}")
            break
        except Exception as e:
            print(f"{Fore.RED}Error: {e}{Style.RESET_ALL}")


if __name__ == "__main__":
    main()