Add LLM configuration and pipeline execution
This commit is contained in:
@@ -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")
|
||||
|
||||
@@ -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()
|
||||
|
||||
+3
-2
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user