kairo-code 0.1.0__py3-none-any.whl

kairo_code/heartbeat.py

@@ -0,0 +1,63 @@
+"""Agent heartbeat for Kairo Code CLI.
+
+Sends periodic heartbeat to the Kairo backend to register this CLI
+instance as an active agent. Requires a valid API key and agent ID.
+"""
+
+import threading
+import time
+import logging
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+HEARTBEAT_INTERVAL = 30  # seconds
+BACKEND_URL = "https://app.kaironlabs.io"
+
+
+class Heartbeat:
+    """Background heartbeat sender."""
+
+    def __init__(self, api_key: str, agent_id: str):
+        self.api_key = api_key
+        self.agent_id = agent_id
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        """Start the heartbeat background thread."""
+        if not self.api_key or not self.agent_id:
+            return
+
+        self._thread = threading.Thread(target=self._loop, daemon=True)
+        self._thread.start()
+        logger.info("Heartbeat started for agent %s", self.agent_id)
+
+    def stop(self) -> None:
+        """Stop the heartbeat."""
+        self._stop.set()
+        if self._thread:
+            self._thread.join(timeout=5)
+        logger.info("Heartbeat stopped")
+
+    def _loop(self) -> None:
+        """Heartbeat loop running in background thread."""
+        while not self._stop.is_set():
+            self._stop.wait(HEARTBEAT_INTERVAL)
+            if self._stop.is_set():
+                break
+
+            try:
+                resp = httpx.post(
+                    f"{BACKEND_URL}/agents/heartbeat",
+                    json={"agent_id": self.agent_id, "status": "online"},
+                    headers={"Authorization": f"Bearer {self.api_key}"},
+                    timeout=10,
+                )
+                if resp.status_code == 200:
+                    logger.debug("Heartbeat OK for agent %s", self.agent_id)
+                else:
+                    logger.warning("Heartbeat returned %s", resp.status_code)
+            except Exception as e:
+                logger.warning("Heartbeat failed: %s", e)

kairo_code/llm.py

@@ -0,0 +1,291 @@
+"""Cloud LLM wrapper with streaming support via OpenAI-compatible API (vLLM)"""
+
+import time
+from typing import Generator, Any
+
+from openai import OpenAI
+
+from .config import Config
+from .logging_config import get_llm_logger, log_model_selection
+
+logger = get_llm_logger()
+
+
+class LLM:
+    """Wrapper around OpenAI-compatible API (vLLM) with streaming and model switching."""
+
+    def __init__(self):
+        self.config = Config()
+        self.client = OpenAI(
+            base_url=self.config.cloud_endpoint,
+            api_key=self.config.cloud_api_key,
+        )
+        self._available_models: list[str] | None = None
+        # Don't cache resolved models - always re-resolve from config
+        self._resolved_models: dict[str, str] = {}
+
+    def check_health(self) -> bool:
+        """Check if the vLLM backend is reachable."""
+        import httpx
+        try:
+            base = self.config.cloud_endpoint.rstrip("/v1").rstrip("/")
+            resp = httpx.get(f"{base}/health", timeout=5.0)
+            return resp.status_code == 200
+        except Exception:
+            return False
+
+    def clear_cache(self):
+        """Clear model cache to force re-resolution."""
+        self._available_models = None
+        self._resolved_models = {}
+
+    def _get_available_models(self) -> list[str]:
+        """Get list of available models (cached)."""
+        if self._available_models is None:
+            self._available_models = self.list_models()
+        return self._available_models
+
+    def _model_matches(self, available: str, wanted: str) -> bool:
+        """Check if an available model matches a wanted model name."""
+        # Exact match
+        if available == wanted:
+            return True
+        # Partial match (e.g., "deepseek-coder-33b" matches "deepseek")
+        if available.startswith(wanted) or wanted in available:
+            return True
+        return False
+
+    def select_model(self, category: str) -> str:
+        """
+        Select the best available model for a category.
+
+        Args:
+            category: One of 'coder', 'router', 'chat'
+
+        Returns:
+            The best available model name
+        """
+        logger.debug(f"Selecting model for category: {category}")
+
+        # Check cache first
+        if category in self._resolved_models:
+            logger.debug(f"Using cached model: {self._resolved_models[category]}")
+            return self._resolved_models[category]
+
+        # Get configured model for this category
+        configured = self.config.get(f"models.{category}", "auto")
+        logger.debug(f"Configured model: {configured}")
+
+        # If not "auto", use the configured model directly
+        if configured != "auto":
+            self._resolved_models[category] = configured
+            logger.info(f"Model selection: {category} -> {configured} (explicit)")
+            return configured
+
+        # Auto-select from preferences
+        preferences = self.config.get(f"models.preferences.{category}", [])
+        available = self._get_available_models()
+        logger.debug(f"Available models: {available}")
+        logger.debug(f"Preferences for {category}: {preferences}")
+
+        for preferred in preferences:
+            for avail in available:
+                if self._model_matches(avail, preferred):
+                    self._resolved_models[category] = avail
+                    logger.info(f"Model selection: {category} -> {avail} (matched {preferred})")
+                    log_model_selection(category, avail, available)
+                    return avail
+
+        # Fallback: use any available model
+        if available:
+            fallback = available[0]
+            self._resolved_models[category] = fallback
+            logger.warning(f"Model selection: {category} -> {fallback} (fallback)")
+            return fallback
+
+        # Ultimate fallback: use default model from config
+        default_model = self.config.get("cloud.default_model", "deepseek-ai/deepseek-coder-33b-instruct")
+        self._resolved_models[category] = default_model
+        logger.warning(f"Model selection: {category} -> {default_model} (default)")
+        return default_model
+
+    def generate(
+        self,
+        prompt: str,
+        model: str | None = None,
+        system: str | None = None,
+        stream: bool = True,
+    ) -> Generator[str, None, None] | str:
+        """
+        Generate a response from the LLM.
+
+        Args:
+            prompt: The user prompt
+            model: Model to use (defaults to coder model)
+            system: System prompt (defaults to config system prompt)
+            stream: Whether to stream the response
+
+        Returns:
+            Generator of tokens if streaming, otherwise full response string
+        """
+        model = model or self.select_model("coder")
+        system = system or self.config.system_prompt
+
+        # Convert to chat format for OpenAI API
+        messages = []
+        if system:
+            messages.append({"role": "system", "content": system})
+        messages.append({"role": "user", "content": prompt})
+
+        if stream:
+            return self._stream_chat(messages, model)
+        else:
+            response = self.client.chat.completions.create(
+                model=model,
+                messages=messages,
+                max_tokens=self.config.get("cloud.max_tokens", 4096),
+                temperature=self.config.get("cloud.temperature", 0.1),
+            )
+            return response.choices[0].message.content
+
+    def _stream_generate(
+        self,
+        prompt: str,
+        model: str,
+        system: str,
+    ) -> Generator[str, None, None]:
+        """Stream tokens from the LLM."""
+        messages = []
+        if system:
+            messages.append({"role": "system", "content": system})
+        messages.append({"role": "user", "content": prompt})
+
+        return self._stream_chat(messages, model)
+
+    def chat(
+        self,
+        messages: list[dict[str, str]],
+        model: str | None = None,
+        stream: bool = True,
+    ) -> Generator[str, None, None] | str:
+        """
+        Chat with conversation history.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            model: Model to use (defaults to coder model)
+            stream: Whether to stream the response
+
+        Returns:
+            Generator of tokens if streaming, otherwise full response string
+        """
+        model = model or self.select_model("coder")
+
+        if stream:
+            return self._stream_chat(messages, model)
+        else:
+            response = self.client.chat.completions.create(
+                model=model,
+                messages=messages,
+                max_tokens=self.config.get("cloud.max_tokens", 4096),
+                temperature=self.config.get("cloud.temperature", 0.1),
+            )
+            return response.choices[0].message.content
+
+    def _stream_chat(
+        self,
+        messages: list[dict[str, str]],
+        model: str,
+    ) -> Generator[str, None, None]:
+        """Stream chat response."""
+        logger.info(f"Starting streaming chat with model: {model}")
+        logger.debug(f"Message count: {len(messages)}")
+        if messages:
+            last_msg = messages[-1]
+            logger.debug(f"Last message ({last_msg.get('role')}): {last_msg.get('content', '')[:200]}...")
+
+        start_time = time.time()
+        token_count = 0
+
+        try:
+            logger.debug("Calling cloud API...")
+            stream = self.client.chat.completions.create(
+                model=model,
+                messages=messages,
+                max_tokens=self.config.get("cloud.max_tokens", 4096),
+                temperature=self.config.get("cloud.temperature", 0.1),
+                stream=True,
+            )
+
+            logger.debug("Stream started, receiving tokens...")
+            for chunk in stream:
+                if chunk.choices and chunk.choices[0].delta.content:
+                    token_count += 1
+                    yield chunk.choices[0].delta.content
+
+            elapsed = time.time() - start_time
+            logger.info(f"Chat completed: {token_count} chunks in {elapsed:.2f}s")
+
+        except Exception as e:
+            elapsed = time.time() - start_time
+            logger.error(f"Chat failed after {elapsed:.2f}s: {e}")
+            raise
+
+    def classify(self, prompt: str, categories: list[str] | None = None) -> str:
+        """
+        Use the router model to classify intent.
+
+        Args:
+            prompt: The user's input to classify
+            categories: Optional list of valid categories
+
+        Returns:
+            The classified intent category
+        """
+        logger.info(f"Classifying intent for: {prompt[:100]}...")
+        system = self.config.router_prompt
+        router_model = self.select_model("router")
+
+        start_time = time.time()
+        logger.debug(f"Calling router model: {router_model}")
+
+        messages = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.client.chat.completions.create(
+            model=router_model,
+            messages=messages,
+            max_tokens=50,
+            temperature=0.0,
+        )
+
+        elapsed = time.time() - start_time
+        result = response.choices[0].message.content.strip().lower()
+        logger.info(f"Classification result: '{result}' (took {elapsed:.2f}s)")
+
+        # Validate against categories if provided
+        if categories:
+            for cat in categories:
+                if cat.lower() in result:
+                    return cat.lower()
+            # Default to chat if no match
+            return "chat"
+
+        return result
+
+    def list_models(self) -> list[str]:
+        """List available models from the cloud endpoint."""
+        try:
+            result = self.client.models.list()
+            return [m.id for m in result.data]
+        except Exception as e:
+            logger.warning(f"Failed to list models: {e}")
+            # Return default model if listing fails
+            return [self.config.get("cloud.default_model", "deepseek-ai/deepseek-coder-33b-instruct")]
+
+    def check_model(self, model: str) -> bool:
+        """Check if a model is available."""
+        available = self.list_models()
+        return any(model in m for m in available)

kairo_code/logging_config.py

@@ -0,0 +1,156 @@
+"""Logging configuration for Kairo Code"""
+
+import logging
+import os
+import stat
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Optional
+
+
+# Log directory
+LOG_DIR = Path.home() / ".kairo_code" / "logs"
+
+
+def setup_logging(
+    level: int = logging.DEBUG,
+    console_level: int = logging.WARNING,
+    log_file: Optional[Path] = None,
+) -> logging.Logger:
+    """
+    Setup logging for Kairo Code.
+
+    Args:
+        level: File logging level (default DEBUG for full detail)
+        console_level: Console logging level (default WARNING to not clutter output)
+        log_file: Optional custom log file path
+
+    Returns:
+        Configured logger instance
+    """
+    # Create log directory with secure permissions
+    LOG_DIR.mkdir(parents=True, exist_ok=True)
+    try:
+        os.chmod(LOG_DIR, stat.S_IRWXU)  # 0700 — owner only
+    except OSError:
+        pass
+
+    # Create session log file with timestamp
+    if log_file is None:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        log_file = LOG_DIR / f"kairo_{timestamp}.log"
+
+    # Get or create logger
+    logger = logging.getLogger("kairo_code")
+    logger.setLevel(logging.DEBUG)  # Capture everything, handlers filter
+
+    # Clear existing handlers
+    logger.handlers.clear()
+
+    # File handler - detailed logging
+    file_handler = logging.FileHandler(log_file, encoding="utf-8")
+    file_handler.setLevel(level)
+    file_formatter = logging.Formatter(
+        "%(asctime)s | %(levelname)-8s | %(name)s.%(funcName)s:%(lineno)d | %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S"
+    )
+    file_handler.setFormatter(file_formatter)
+    logger.addHandler(file_handler)
+
+    # Console handler - only warnings/errors (to not clutter UI)
+    console_handler = logging.StreamHandler(sys.stderr)
+    console_handler.setLevel(console_level)
+    console_formatter = logging.Formatter("[%(levelname)s] %(message)s")
+    console_handler.setFormatter(console_formatter)
+    logger.addHandler(console_handler)
+
+    # Store log file path for reference
+    logger.log_file = log_file
+
+    logger.info(f"=== Kairo Code Session Started ===")
+    logger.info(f"Log file: {log_file}")
+
+    return logger
+
+
+def get_logger(name: str = "kairo_code") -> logging.Logger:
+    """Get a logger instance (child of main kairo_code logger)."""
+    return logging.getLogger(name)
+
+
+# Convenience loggers for different components
+def get_llm_logger() -> logging.Logger:
+    return logging.getLogger("kairo_code.llm")
+
+def get_tool_logger() -> logging.Logger:
+    return logging.getLogger("kairo_code.tools")
+
+def get_agent_logger() -> logging.Logger:
+    return logging.getLogger("kairo_code.agents")
+
+def get_router_logger() -> logging.Logger:
+    return logging.getLogger("kairo_code.router")
+
+
+class LogContext:
+    """Context manager for logging sections with clear boundaries."""
+
+    def __init__(self, logger: logging.Logger, section: str, level: int = logging.DEBUG):
+        self.logger = logger
+        self.section = section
+        self.level = level
+
+    def __enter__(self):
+        self.logger.log(self.level, f">>> START: {self.section}")
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type:
+            self.logger.error(f"<<< END: {self.section} (ERROR: {exc_val})")
+        else:
+            self.logger.log(self.level, f"<<< END: {self.section}")
+        return False
+
+
+def log_tool_call(tool_name: str, params: dict, result: "Any") -> None:
+    """Log a tool call with its parameters and result."""
+    logger = get_tool_logger()
+    logger.info(f"TOOL CALL: {tool_name}")
+    logger.debug(f"  Params: {params}")
+    if result.success:
+        logger.debug(f"  Result: SUCCESS")
+        # Truncate long output
+        output = result.output[:500] + "..." if len(result.output) > 500 else result.output
+        logger.debug(f"  Output: {output}")
+    else:
+        logger.warning(f"  Result: FAILED - {result.error}")
+
+
+def log_llm_request(model: str, messages: list, system: str = None) -> None:
+    """Log an LLM request."""
+    logger = get_llm_logger()
+    logger.info(f"LLM REQUEST: model={model}")
+    if system:
+        logger.debug(f"  System prompt: {system[:200]}...")
+    logger.debug(f"  Messages: {len(messages)} total")
+    for i, msg in enumerate(messages[-3:]):  # Log last 3 messages
+        content = msg.get('content', '')[:200]
+        logger.debug(f"  [{i}] {msg.get('role')}: {content}...")
+
+
+def log_llm_response(model: str, response: str, tokens: int = None) -> None:
+    """Log an LLM response."""
+    logger = get_llm_logger()
+    truncated = response[:500] + "..." if len(response) > 500 else response
+    logger.info(f"LLM RESPONSE: model={model}, length={len(response)}")
+    logger.debug(f"  Content: {truncated}")
+    if tokens:
+        logger.debug(f"  Tokens: {tokens}")
+
+
+def log_model_selection(category: str, selected: str, available: list) -> None:
+    """Log model auto-selection."""
+    logger = get_llm_logger()
+    logger.info(f"MODEL SELECTION: {category} -> {selected}")
+    logger.debug(f"  Available models: {available}")