Add LLM configuration and pipeline execution
This commit is contained in:
@@ -1,10 +1,15 @@
|
|||||||
import argparse
|
import argparse
|
||||||
|
import asyncio
|
||||||
|
import os
|
||||||
|
|
||||||
|
from src.pipeline.orchestrator import PipelineOrchestrator
|
||||||
from src.rag.manager import RAGManager
|
from src.rag.manager import RAGManager
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
parser = argparse.ArgumentParser(description="D&D Helpers CLI")
|
parser = argparse.ArgumentParser(description="D&D Helpers CLI")
|
||||||
|
|
||||||
|
# RAG Ingestion Arguments
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--ingest-pdf",
|
"--ingest-pdf",
|
||||||
type=str,
|
type=str,
|
||||||
@@ -21,9 +26,65 @@ def main():
|
|||||||
help="Path to a directory of markdown files to ingest into the RAG system",
|
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()
|
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:
|
if args.ingest_pdf:
|
||||||
print(f"Ingesting PDF: {args.ingest_pdf}...")
|
print(f"Ingesting PDF: {args.ingest_pdf}...")
|
||||||
@@ -40,8 +101,8 @@ def main():
|
|||||||
rag_manager.ingest_directory(args.ingest_dir)
|
rag_manager.ingest_directory(args.ingest_dir)
|
||||||
print("Directory ingestion complete.")
|
print("Directory ingestion complete.")
|
||||||
|
|
||||||
if not any([args.ingest_pdf, args.ingest_file, args.ingest_dir]):
|
if not any([args.ingest_pdf, args.ingest_file, args.ingest_dir, args.run_pipeline]):
|
||||||
print("Hello from dnd-helpers!")
|
print("Hello from dnd-helpers! Use --help to see available commands.")
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ class LLMProcessor:
|
|||||||
api_key: Optional[str] = None,
|
api_key: Optional[str] = None,
|
||||||
base_url: Optional[str] = None,
|
base_url: Optional[str] = None,
|
||||||
model: Optional[str] = None,
|
model: Optional[str] = None,
|
||||||
|
backend: Optional[str] = None,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Initializes the LLMProcessor.
|
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 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 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 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
|
# Ollama's OpenAI-compatible API
|
||||||
final_base_url = base_url or "http://localhost:11434/v1"
|
final_base_url = base_url or "http://localhost:11434/v1"
|
||||||
final_api_key = api_key or "ollama"
|
final_api_key = api_key or "ollama"
|
||||||
elif backend == "vllm":
|
elif backend_env == "vllm":
|
||||||
# Remote vLLM server
|
# Remote vLLM server
|
||||||
final_base_url = base_url or os.environ.get("OPENAI_BASE_URL")
|
final_base_url = base_url or os.environ.get("OPENAI_BASE_URL")
|
||||||
final_api_key = api_key or os.environ.get("OPENAI_API_KEY")
|
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_base_url = base_url or os.environ.get("OPENAI_BASE_URL")
|
||||||
final_api_key = api_key or os.environ.get("OPENAI_API_KEY")
|
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:
|
try:
|
||||||
self.client = OpenAI(
|
self.client = OpenAI(
|
||||||
api_key=final_api_key,
|
api_key=final_api_key,
|
||||||
base_url=final_base_url,
|
base_url=final_base_url,
|
||||||
)
|
)
|
||||||
# Simple connectivity check for local backends
|
# Simple connectivity check for local backends
|
||||||
if backend == "ollama":
|
if backend_env == "ollama":
|
||||||
# We can't easily check connectivity without making a call,
|
# We can't easily check connectivity without making a call,
|
||||||
# but we can ensure the client is initialized.
|
# but we can ensure the client is initialized.
|
||||||
pass
|
pass
|
||||||
except Exception as e:
|
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
|
raise
|
||||||
|
|
||||||
self.model = model or os.environ.get("LLM_MODEL", "gpt-4o")
|
self.model = model or os.environ.get("LLM_MODEL", "gpt-4o")
|
||||||
|
|||||||
@@ -39,14 +39,15 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
|
|
||||||
class PipelineOrchestrator:
|
class PipelineOrchestrator:
|
||||||
def __init__(self, loop: asyncio.AbstractEventLoop):
|
def __init__(self, loop: asyncio.AbstractEventLoop, llm_config: Optional[dict] = None):
|
||||||
self.loop = loop
|
self.loop = loop
|
||||||
|
self.llm_config = llm_config or {}
|
||||||
|
|
||||||
# Modules
|
# Modules
|
||||||
self.listener = AudioListener(loop=self.loop)
|
self.listener = AudioListener(loop=self.loop)
|
||||||
self.transcriber = Transcriber(model_size="base", device="cuda")
|
self.transcriber = Transcriber(model_size="base", device="cuda")
|
||||||
self.processor = LLMProcessor()
|
self.processor = LLMProcessor(**self.llm_config)
|
||||||
self.rag_manager = RAGManager()
|
self.rag_manager = RAGManager(llm_config=self.llm_config)
|
||||||
|
|
||||||
# Queues
|
# Queues
|
||||||
self.stt_to_clean_queue = asyncio.Queue()
|
self.stt_to_clean_queue = asyncio.Queue()
|
||||||
|
|||||||
+3
-2
@@ -12,8 +12,9 @@ from src.llm.processor import LLMProcessor
|
|||||||
|
|
||||||
|
|
||||||
class RAGManager:
|
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.persist_dir = persist_dir
|
||||||
|
self.llm_config = llm_config or {}
|
||||||
self.db = chromadb.PersistentClient(path=self.persist_dir)
|
self.db = chromadb.PersistentClient(path=self.persist_dir)
|
||||||
self.collection_name = "phb_collection"
|
self.collection_name = "phb_collection"
|
||||||
|
|
||||||
@@ -110,7 +111,7 @@ class RAGManager:
|
|||||||
if not nodes:
|
if not nodes:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
processor = LLMProcessor()
|
processor = LLMProcessor(**self.llm_config)
|
||||||
|
|
||||||
# Construct the context from retrieved nodes
|
# Construct the context from retrieved nodes
|
||||||
context_text = "\n\n".join(
|
context_text = "\n\n".join(
|
||||||
|
|||||||
Reference in New Issue
Block a user