118 lines
3.3 KiB
Python
118 lines
3.3 KiB
Python
import argparse
|
|
import asyncio
|
|
import os
|
|
|
|
from src.pipeline.orchestrator import PipelineOrchestrator
|
|
from src.rag.manager import RAGManager
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="D&D Helpers CLI")
|
|
|
|
# RAG Ingestion Arguments
|
|
parser.add_argument(
|
|
"--ingest-pdf",
|
|
type=str,
|
|
help="Path to a PDF file to ingest into the RAG system",
|
|
)
|
|
parser.add_argument(
|
|
"--ingest-file",
|
|
type=str,
|
|
help="Path to a markdown file to ingest into the RAG system",
|
|
)
|
|
parser.add_argument(
|
|
"--ingest-dir",
|
|
type=str,
|
|
help="Path to a directory of markdown files to ingest into the RAG system",
|
|
)
|
|
|
|
# LLM Configuration Arguments
|
|
parser.add_argument(
|
|
"--llm-backend",
|
|
type=str,
|
|
choices=["openai", "ollama", "vllm"],
|
|
default=os.environ.get("LLM_BACKEND", "openai"),
|
|
help="LLM backend to use",
|
|
)
|
|
parser.add_argument(
|
|
"--llm-model",
|
|
type=str,
|
|
default=os.environ.get("LLM_MODEL", "gpt-4o"),
|
|
help="The model to use for processing",
|
|
)
|
|
parser.add_argument(
|
|
"--llm-api-key",
|
|
type=str,
|
|
default=os.environ.get("OPENAI_API_KEY"),
|
|
help="API key for the LLM backend",
|
|
)
|
|
parser.add_argument(
|
|
"--llm-base-url",
|
|
type=str,
|
|
default=os.environ.get("OPENAI_BASE_URL"),
|
|
help="Base URL for the LLM backend",
|
|
)
|
|
|
|
# STT Configuration Arguments
|
|
parser.add_argument(
|
|
"--whisper-model",
|
|
type=str,
|
|
default=os.environ.get("WHISPER_MODEL", "turbo"),
|
|
help="The Whisper model to use for STT",
|
|
)
|
|
|
|
# Pipeline Execution Argument
|
|
parser.add_argument(
|
|
"--run-pipeline",
|
|
action="store_true",
|
|
help="Start the main orchestration pipeline (TUI + STT + LLM)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
llm_config = {
|
|
"backend": args.llm_backend,
|
|
"model": args.llm_model,
|
|
"api_key": args.llm_api_key,
|
|
"base_url": args.llm_base_url,
|
|
}
|
|
|
|
# Remove None values to allow defaults to take over if not provided
|
|
llm_config = {k: v for k, v in llm_config.items() if v is not None}
|
|
|
|
if args.run_pipeline:
|
|
async def run_pipeline():
|
|
loop = asyncio.get_event_loop()
|
|
orchestrator = PipelineOrchestrator(loop, llm_config=llm_config, whisper_model=args.whisper_model)
|
|
try:
|
|
await orchestrator.run()
|
|
except KeyboardInterrupt:
|
|
orchestrator.stop()
|
|
|
|
asyncio.run(run_pipeline())
|
|
return
|
|
|
|
rag_manager = RAGManager(llm_config=llm_config)
|
|
|
|
if args.ingest_pdf:
|
|
print(f"Ingesting PDF: {args.ingest_pdf}...")
|
|
rag_manager.ingest_pdf(args.ingest_pdf)
|
|
print("PDF ingestion complete.")
|
|
|
|
if args.ingest_file:
|
|
print(f"Ingesting File: {args.ingest_file}...")
|
|
rag_manager.ingest_file(args.ingest_file)
|
|
print("File ingestion complete.")
|
|
|
|
if args.ingest_dir:
|
|
print(f"Ingesting Directory: {args.ingest_dir}...")
|
|
rag_manager.ingest_directory(args.ingest_dir)
|
|
print("Directory ingestion complete.")
|
|
|
|
if not any([args.ingest_pdf, args.ingest_file, args.ingest_dir, args.run_pipeline]):
|
|
print("Hello from dnd-helpers! Use --help to see available commands.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|