src/core/lfm_controller.py

import os
import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor
from llama_cpp import Llama

logging.basicConfig(level=logging.INFO, format='%(asctime)s - [LFMController] - %(levelname)s - %(message)s')
logger = logging.getLogger("LFMController")

class LFMController:
    def __init__(self, model_path: str = "", n_ctx: int = 4096, n_threads: int = 6, n_gpu_layers: int = -1):
        if not model_path:
            candidates = [
                os.path.expanduser("~/.vitalis/models/Llama-3.2-3B-Instruct-Q4_K_M.gguf"),
                os.path.expanduser("~/.vitalis/models/LFM2.5-1.2B-Instruct-Q4_K_M.gguf"),
                "Llama-3.2-3B-Instruct-Q4_K_M.gguf",
                "LFM2.5-1.2B-Instruct-Q4_K_M.gguf",
            ]
            for c in candidates:
                if os.path.exists(c):
                    model_path = c
                    break
            if not model_path:
                model_path = candidates[0]
        if not os.path.exists(model_path):
            logger.critical(f"Target model weights missing at path: {model_path}")
            raise FileNotFoundError(f"Model file target missing: {model_path}")

        logger.info(f"Initializing local model instance from {model_path}...")
        try:
            self.llm = Llama(
                model_path=model_path,
                n_ctx=n_ctx,
                n_threads=n_threads,
                n_gpu_layers=n_gpu_layers,
                verbose=False
            )
            logger.info("Model hardware acceleration context successfully initialized.")
        except Exception as e:
            logger.error(f"Failed to load hardware context for GGUF: {str(e)}")
            raise e

        self.executor = ThreadPoolExecutor(max_workers=1)

    def execute_raw(self, prompt: str, max_tokens: int = 1024, temperature: float = 0.2, top_p: float = 0.95) -> str:
        try:
            response = self.llm(
                prompt=prompt,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                stop=["<|endoftext|>", "###", "Instruction:", "Response:"]
            )
            output_text = response["choices"][0]["text"].strip()
            return output_text
        except Exception as e:
            logger.error(f"Error encountered during raw execution sequence: {str(e)}")
            return f"EXECUTION_ERROR: {str(e)}"

    async def generate_async(self, prompt: str, max_tokens: int = 1024, temperature: float = 0.2, top_p: float = 0.95) -> str:
        loop = asyncio.get_running_loop()
        try:
            return await loop.run_in_executor(
                self.executor,
                self.execute_raw,
                prompt, max_tokens, temperature, top_p
            )
        except Exception as e:
            logger.error(f"Async worker thread crashed: {str(e)}")
            return f"ASYNC_EXECUTION_ERROR: {str(e)}"

    def shutdown(self):
        self.executor.shutdown(wait=True)