From 71ecdb34689399e4a23b2be338d819a7a5cc85d8 Mon Sep 17 00:00:00 2001 From: charles Date: Sun, 31 May 2026 14:13:58 -0700 Subject: [PATCH] Add LLM configuration and pipeline execution --- main.py | 67 ++++++++++++++++++++++++++++++++++-- src/llm/processor.py | 15 ++++---- src/pipeline/orchestrator.py | 7 ++-- src/rag/manager.py | 5 +-- 4 files changed, 80 insertions(+), 14 deletions(-) diff --git a/main.py b/main.py index 071291d..c233993 100644 --- a/main.py +++ b/main.py @@ -1,10 +1,15 @@ 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, @@ -21,9 +26,65 @@ def main(): 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", + ) + + # Pipeline Execution Argument + parser.add_argument( + "--run-pipeline", + action="store_true", + help="Start the main orchestration pipeline (TUI + STT + LLM)", + ) + args = parser.parse_args() - rag_manager = RAGManager() + 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) + 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}...") @@ -40,8 +101,8 @@ def main(): rag_manager.ingest_directory(args.ingest_dir) print("Directory ingestion complete.") - if not any([args.ingest_pdf, args.ingest_file, args.ingest_dir]): - print("Hello from dnd-helpers!") + 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__": diff --git a/src/llm/processor.py b/src/llm/processor.py index 04c329c..336e6b2 100644 --- a/src/llm/processor.py +++ b/src/llm/processor.py @@ -23,6 +23,7 @@ class LLMProcessor: api_key: Optional[str] = None, base_url: Optional[str] = None, model: Optional[str] = None, + backend: Optional[str] = None, ): """ Initializes the LLMProcessor. @@ -30,14 +31,16 @@ class LLMProcessor: :param api_key: OpenAI API key. If None, it looks for OPENAI_API_KEY in environment variables. :param base_url: OpenAI-compatible base URL (e.g., for vLLM). :param model: The model to use for processing. If None, it looks for LLM_MODEL in environment variables. + :param backend: The LLM backend to use (openai, ollama, or vllm). """ - backend = os.environ.get("LLM_BACKEND", "openai").lower() + # Use provided backend or fallback to environment variable + backend_env = backend or os.environ.get("LLM_BACKEND", "openai").lower() - if backend == "ollama": + if backend_env == "ollama": # Ollama's OpenAI-compatible API final_base_url = base_url or "http://localhost:11434/v1" final_api_key = api_key or "ollama" - elif backend == "vllm": + elif backend_env == "vllm": # Remote vLLM server final_base_url = base_url or os.environ.get("OPENAI_BASE_URL") final_api_key = api_key or os.environ.get("OPENAI_API_KEY") @@ -45,19 +48,19 @@ class LLMProcessor: final_base_url = base_url or os.environ.get("OPENAI_BASE_URL") final_api_key = api_key or os.environ.get("OPENAI_API_KEY") - logger.info(f"Using LLM backend: {backend}") + logger.info(f"Using LLM backend: {backend_env}") try: self.client = OpenAI( api_key=final_api_key, base_url=final_base_url, ) # Simple connectivity check for local backends - if backend == "ollama": + if backend_env == "ollama": # We can't easily check connectivity without making a call, # but we can ensure the client is initialized. pass except Exception as e: - logger.error(f"Error initializing LLM client for backend {backend}: {e}") + logger.error(f"Error initializing LLM client for backend {backend_env}: {e}") raise self.model = model or os.environ.get("LLM_MODEL", "gpt-4o") diff --git a/src/pipeline/orchestrator.py b/src/pipeline/orchestrator.py index 6023935..cb0817a 100644 --- a/src/pipeline/orchestrator.py +++ b/src/pipeline/orchestrator.py @@ -39,14 +39,15 @@ logger = logging.getLogger(__name__) class PipelineOrchestrator: - def __init__(self, loop: asyncio.AbstractEventLoop): + def __init__(self, loop: asyncio.AbstractEventLoop, llm_config: Optional[dict] = None): self.loop = loop + self.llm_config = llm_config or {} # Modules self.listener = AudioListener(loop=self.loop) self.transcriber = Transcriber(model_size="base", device="cuda") - self.processor = LLMProcessor() - self.rag_manager = RAGManager() + self.processor = LLMProcessor(**self.llm_config) + self.rag_manager = RAGManager(llm_config=self.llm_config) # Queues self.stt_to_clean_queue = asyncio.Queue() diff --git a/src/rag/manager.py b/src/rag/manager.py index 63eb641..9fde4a8 100644 --- a/src/rag/manager.py +++ b/src/rag/manager.py @@ -12,8 +12,9 @@ from src.llm.processor import LLMProcessor class RAGManager: - def __init__(self, persist_dir: str = "data/rag_index"): + def __init__(self, persist_dir: str = "data/rag_index", llm_config: Optional[dict] = None): self.persist_dir = persist_dir + self.llm_config = llm_config or {} self.db = chromadb.PersistentClient(path=self.persist_dir) self.collection_name = "phb_collection" @@ -110,7 +111,7 @@ class RAGManager: if not nodes: return [] - processor = LLMProcessor() + processor = LLMProcessor(**self.llm_config) # Construct the context from retrieved nodes context_text = "\n\n".join(