basion-agent 0.4.0__py3-none-any.whl

basion_agent/gateway_pb2.py

@@ -0,0 +1,73 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: gateway.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'gateway.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\rgateway.proto\x12\x07gateway\"\xf6\x01\n\rClientMessage\x12$\n\x04\x61uth\x18\x01 \x01(\x0b\x32\x14.gateway.AuthRequestH\x00\x12.\n\tsubscribe\x18\x02 \x01(\x0b\x32\x19.gateway.SubscribeRequestH\x00\x12\x32\n\x0bunsubscribe\x18\x03 \x01(\x0b\x32\x1b.gateway.UnsubscribeRequestH\x00\x12*\n\x07produce\x18\x04 \x01(\x0b\x32\x17.gateway.ProduceRequestH\x00\x12$\n\x04ping\x18\x05 \x01(\x0b\x32\x14.gateway.PingRequestH\x00\x42\t\n\x07payload\"\xd4\x02\n\rServerMessage\x12.\n\rauth_response\x18\x01 \x01(\x0b\x32\x15.gateway.AuthResponseH\x00\x12(\n\x07message\x18\x02 \x01(\x0b\x32\x15.gateway.KafkaMessageH\x00\x12*\n\x0bproduce_ack\x18\x03 \x01(\x0b\x32\x13.gateway.ProduceAckH\x00\x12.\n\rsubscribe_ack\x18\x04 \x01(\x0b\x32\x15.gateway.SubscribeAckH\x00\x12\x32\n\x0funsubscribe_ack\x18\x05 \x01(\x0b\x32\x17.gateway.UnsubscribeAckH\x00\x12%\n\x04pong\x18\x06 \x01(\x0b\x32\x15.gateway.PongResponseH\x00\x12\'\n\x05\x65rror\x18\x07 \x01(\x0b\x32\x16.gateway.ErrorResponseH\x00\x42\t\n\x07payload\"3\n\x0b\x41uthRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12\x13\n\x0b\x61gent_names\x18\x02 \x03(\t\"@\n\x0c\x41uthResponse\x12\x15\n\rconnection_id\x18\x01 \x01(\t\x12\x19\n\x11subscribed_topics\x18\x02 \x03(\t\"&\n\x10SubscribeRequest\x12\x12\n\nagent_name\x18\x01 \x01(\t\"\x1d\n\x0cSubscribeAck\x12\r\n\x05topic\x18\x01 \x01(\t\"(\n\x12UnsubscribeRequest\x12\x12\n\nagent_name\x18\x01 \x01(\t\"$\n\x0eUnsubscribeAck\x12\x12\n\nagent_name\x18\x01 \x01(\t\"\xb9\x01\n\x0eProduceRequest\x12\r\n\x05topic\x18\x01 \x01(\t\x12\x0b\n\x03key\x18\x02 \x01(\t\x12\x35\n\x07headers\x18\x03 \x03(\x0b\x32$.gateway.ProduceRequest.HeadersEntry\x12\x0c\n\x04\x62ody\x18\x04 \x01(\x0c\x12\x16\n\x0e\x63orrelation_id\x18\x05 \x01(\t\x1a.\n\x0cHeadersEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"V\n\nProduceAck\x12\r\n\x05topic\x18\x01 \x01(\t\x12\x11\n\tpartition\x18\x02 \x01(\x05\x12\x0e\n\x06offset\x18\x03 \x01(\x03\x12\x16\n\x0e\x63orrelation_id\x18\x04 \x01(\t\"\xd3\x01\n\x0cKafkaMessage\x12\r\n\x05topic\x18\x01 \x01(\t\x12\x11\n\tpartition\x18\x02 \x01(\x05\x12\x0e\n\x06offset\x18\x03 \x01(\x03\x12\x0b\n\x03key\x18\x04 \x01(\t\x12\x33\n\x07headers\x18\x05 \x03(\x0b\x32\".gateway.KafkaMessage.HeadersEntry\x12\x0c\n\x04\x62ody\x18\x06 \x01(\x0c\x12\x11\n\ttimestamp\x18\x07 \x01(\x03\x1a.\n\x0cHeadersEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\" \n\x0bPingRequest\x12\x11\n\ttimestamp\x18\x01 \x01(\x03\"B\n\x0cPongResponse\x12\x18\n\x10\x63lient_timestamp\x18\x01 \x01(\x03\x12\x18\n\x10server_timestamp\x18\x02 \x01(\x03\"F\n\rErrorResponse\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x16\n\x0e\x63orrelation_id\x18\x03 \x01(\t2Q\n\x0c\x41gentGateway\x12\x41\n\x0b\x41gentStream\x12\x16.gateway.ClientMessage\x1a\x16.gateway.ServerMessage(\x01\x30\x01\x42 Z\x1e\x61gent-gateway/internal/grpc/pbb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'gateway_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'Z\036agent-gateway/internal/grpc/pb'
+  _globals['_PRODUCEREQUEST_HEADERSENTRY']._loaded_options = None
+  _globals['_PRODUCEREQUEST_HEADERSENTRY']._serialized_options = b'8\001'
+  _globals['_KAFKAMESSAGE_HEADERSENTRY']._loaded_options = None
+  _globals['_KAFKAMESSAGE_HEADERSENTRY']._serialized_options = b'8\001'
+  _globals['_CLIENTMESSAGE']._serialized_start=27
+  _globals['_CLIENTMESSAGE']._serialized_end=273
+  _globals['_SERVERMESSAGE']._serialized_start=276
+  _globals['_SERVERMESSAGE']._serialized_end=616
+  _globals['_AUTHREQUEST']._serialized_start=618
+  _globals['_AUTHREQUEST']._serialized_end=669
+  _globals['_AUTHRESPONSE']._serialized_start=671
+  _globals['_AUTHRESPONSE']._serialized_end=735
+  _globals['_SUBSCRIBEREQUEST']._serialized_start=737
+  _globals['_SUBSCRIBEREQUEST']._serialized_end=775
+  _globals['_SUBSCRIBEACK']._serialized_start=777
+  _globals['_SUBSCRIBEACK']._serialized_end=806
+  _globals['_UNSUBSCRIBEREQUEST']._serialized_start=808
+  _globals['_UNSUBSCRIBEREQUEST']._serialized_end=848
+  _globals['_UNSUBSCRIBEACK']._serialized_start=850
+  _globals['_UNSUBSCRIBEACK']._serialized_end=886
+  _globals['_PRODUCEREQUEST']._serialized_start=889
+  _globals['_PRODUCEREQUEST']._serialized_end=1074
+  _globals['_PRODUCEREQUEST_HEADERSENTRY']._serialized_start=1028
+  _globals['_PRODUCEREQUEST_HEADERSENTRY']._serialized_end=1074
+  _globals['_PRODUCEACK']._serialized_start=1076
+  _globals['_PRODUCEACK']._serialized_end=1162
+  _globals['_KAFKAMESSAGE']._serialized_start=1165
+  _globals['_KAFKAMESSAGE']._serialized_end=1376
+  _globals['_KAFKAMESSAGE_HEADERSENTRY']._serialized_start=1028
+  _globals['_KAFKAMESSAGE_HEADERSENTRY']._serialized_end=1074
+  _globals['_PINGREQUEST']._serialized_start=1378
+  _globals['_PINGREQUEST']._serialized_end=1410
+  _globals['_PONGRESPONSE']._serialized_start=1412
+  _globals['_PONGRESPONSE']._serialized_end=1478
+  _globals['_ERRORRESPONSE']._serialized_start=1480
+  _globals['_ERRORRESPONSE']._serialized_end=1550
+  _globals['_AGENTGATEWAY']._serialized_start=1552
+  _globals['_AGENTGATEWAY']._serialized_end=1633
+# @@protoc_insertion_point(module_scope)

basion_agent/gateway_pb2_grpc.py

@@ -0,0 +1,101 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+from . import gateway_pb2 as gateway__pb2
+
+GRPC_GENERATED_VERSION = '1.76.0'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + ' but the generated code in gateway_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )
+
+
+class AgentGatewayStub(object):
+    """Bidirectional streaming service for agent communication
+    """
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.AgentStream = channel.stream_stream(
+                '/gateway.AgentGateway/AgentStream',
+                request_serializer=gateway__pb2.ClientMessage.SerializeToString,
+                response_deserializer=gateway__pb2.ServerMessage.FromString,
+                _registered_method=True)
+
+
+class AgentGatewayServicer(object):
+    """Bidirectional streaming service for agent communication
+    """
+
+    def AgentStream(self, request_iterator, context):
+        """Main bidirectional stream for all agent communication
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_AgentGatewayServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'AgentStream': grpc.stream_stream_rpc_method_handler(
+                    servicer.AgentStream,
+                    request_deserializer=gateway__pb2.ClientMessage.FromString,
+                    response_serializer=gateway__pb2.ServerMessage.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'gateway.AgentGateway', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('gateway.AgentGateway', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class AgentGateway(object):
+    """Bidirectional streaming service for agent communication
+    """
+
+    @staticmethod
+    def AgentStream(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(
+            request_iterator,
+            target,
+            '/gateway.AgentGateway/AgentStream',
+            gateway__pb2.ClientMessage.SerializeToString,
+            gateway__pb2.ServerMessage.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

basion_agent/heartbeat.py

@@ -0,0 +1,84 @@
+"""Heartbeat manager for keeping agent status active."""
+
+import logging
+import time
+import threading
+from typing import Optional, TYPE_CHECKING
+
+from .exceptions import HeartbeatError
+
+if TYPE_CHECKING:
+    from .gateway_client import GatewayClient
+
+logger = logging.getLogger(__name__)
+
+
+class HeartbeatManager:
+    """Manages periodic heartbeat signals to the API via the gateway proxy."""
+
+    def __init__(
+        self,
+        gateway_client: "GatewayClient",
+        agent_name: str,
+        interval: int = 60
+    ):
+        self.gateway_client = gateway_client
+        self.agent_name = agent_name
+        self.interval = interval
+        self._running = False
+        self._thread: Optional[threading.Thread] = None
+
+    def start(self):
+        """Start the heartbeat thread."""
+        if self._running:
+            raise HeartbeatError("Heartbeat already running")
+
+        self._running = True
+        self._thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
+        self._thread.start()
+        logger.info(f"Heartbeat started for agent '{self.agent_name}'")
+
+    def stop(self):
+        """Stop the heartbeat thread."""
+        if not self._running:
+            return
+
+        self._running = False
+        if self._thread:
+            self._thread.join(timeout=5.0)
+        logger.info("Heartbeat stopped")
+
+    def _heartbeat_loop(self):
+        """Background loop that sends heartbeat periodically."""
+        while self._running:
+            try:
+                self._send_heartbeat()
+            except Exception as e:
+                logger.error(f"Heartbeat send failed: {e}", exc_info=True)
+
+            # Sleep in small intervals to allow quick shutdown
+            for _ in range(self.interval):
+                if not self._running:
+                    break
+                time.sleep(1)
+
+    def _send_heartbeat(self):
+        """Send a single heartbeat to the API via gateway proxy."""
+        payload = {"name": self.agent_name}
+
+        try:
+            response = self.gateway_client.http_post(
+                "ai-inventory",
+                "/api/v1/health/heartbeat",
+                json=payload,
+                timeout=5.0
+            )
+            if response.status_code == 404:
+                logger.warning(f"Agent '{self.agent_name}' not found. Re-registration may be needed.")
+            elif response.status_code not in (200, 201):
+                logger.warning(f"Heartbeat returned status {response.status_code}: {response.text}")
+            else:
+                logger.debug(f"Heartbeat sent successfully for agent '{self.agent_name}'")
+        except Exception as e:
+            logger.error(f"Heartbeat request failed: {e}")
+            raise HeartbeatError(f"Failed to send heartbeat: {e}")

basion_agent/loki_handler.py

@@ -0,0 +1,355 @@
+"""
+Loki log handler for remote log collection.
+
+This module provides a Python logging.Handler that batches logs
+and sends them to a centralized Loki instance via the agent-gateway.
+"""
+
+import atexit
+import gzip
+import json
+import logging
+import os
+import queue
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Any, TYPE_CHECKING
+
+import requests
+
+if TYPE_CHECKING:
+    from .gateway_client import GatewayClient
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LogEntry:
+    """A single log entry to be sent to Loki."""
+    timestamp: int  # Unix timestamp in nanoseconds
+    level: str
+    message: str
+    extra: Dict[str, str] = field(default_factory=dict)
+
+
+class LokiLogHandler(logging.Handler):
+    """
+    A logging handler that batches and sends logs to Loki via the agent-gateway.
+
+    Features:
+    - Batches logs to reduce HTTP overhead
+    - Background thread for non-blocking log sending
+    - Graceful degradation when Loki is unavailable
+    - Bounded buffer to prevent memory exhaustion
+
+    Example:
+        from basion_agent.loki_handler import LokiLogHandler
+
+        handler = LokiLogHandler(
+            gateway_client=gateway_client,
+            agent_name="my-agent",
+            batch_size=100,
+            flush_interval=5.0,
+        )
+
+        # Add to root logger or specific logger
+        logging.getLogger().addHandler(handler)
+    """
+
+    def __init__(
+        self,
+        gateway_client: "GatewayClient",
+        agent_name: str,
+        batch_size: int = 100,
+        flush_interval: float = 5.0,
+        max_buffer_size: int = 10000,
+        retry_max_delay: float = 30.0,
+        level: int = logging.DEBUG,
+        enable_compression: bool = True,
+    ):
+        """
+        Initialize the Loki log handler.
+
+        Args:
+            gateway_client: The GatewayClient for HTTP communication
+            agent_name: Name of the agent (used as a label in Loki)
+            batch_size: Number of logs to batch before sending (default: 100)
+            flush_interval: Maximum seconds between flushes (default: 5.0)
+            max_buffer_size: Maximum logs to buffer (default: 10000)
+            retry_max_delay: Maximum retry delay in seconds (default: 30.0)
+            level: Minimum log level to handle (default: DEBUG)
+            enable_compression: Use gzip compression for log batches (default: True)
+        """
+        super().__init__(level=level)
+
+        self.gateway_client = gateway_client
+        self.agent_name = agent_name
+        self.batch_size = batch_size
+        self.flush_interval = flush_interval
+        self.max_buffer_size = max_buffer_size
+        self.retry_max_delay = retry_max_delay
+        self.enable_compression = enable_compression
+
+        # Environment-based log level override for production
+        if os.getenv("ENVIRONMENT") == "production":
+            production_level = max(level, logging.INFO)
+            self.setLevel(production_level)
+            if production_level != level:
+                logger.info(f"Production mode: Elevated log level from {level} to {production_level}")
+
+        # Thread-safe queue for log entries
+        self._queue: queue.Queue[LogEntry] = queue.Queue(maxsize=max_buffer_size)
+
+        # Background flush thread
+        self._running = False
+        self._thread: Optional[threading.Thread] = None
+        self._lock = threading.Lock()
+
+        # Retry state
+        self._retry_delay = 1.0
+        self._last_error_time: Optional[float] = None
+
+        # Start background thread
+        self._start_flush_thread()
+
+        # Register cleanup on exit
+        atexit.register(self.close)
+
+    def _start_flush_thread(self):
+        """Start the background flush thread."""
+        with self._lock:
+            if self._running:
+                return
+
+            self._running = True
+            self._thread = threading.Thread(
+                target=self._flush_loop,
+                daemon=True,
+                name=f"LokiLogHandler-{self.agent_name}"
+            )
+            self._thread.start()
+
+    def _flush_loop(self):
+        """Background loop that periodically flushes logs."""
+        last_flush = time.time()
+
+        while self._running:
+            try:
+                # Calculate time until next flush
+                elapsed = time.time() - last_flush
+                wait_time = max(0, self.flush_interval - elapsed)
+
+                # Try to collect logs up to batch_size or until timeout
+                entries: List[LogEntry] = []
+                deadline = time.time() + wait_time
+
+                while len(entries) < self.batch_size:
+                    remaining = max(0, deadline - time.time())
+                    if remaining <= 0 and entries:
+                        break
+
+                    try:
+                        entry = self._queue.get(timeout=min(remaining, 0.5) if remaining > 0 else 0.1)
+                        entries.append(entry)
+                        self._queue.task_done()
+                    except queue.Empty:
+                        if time.time() >= deadline:
+                            break
+
+                # Flush if we have entries
+                if entries:
+                    self._flush_entries(entries)
+                    last_flush = time.time()
+
+            except Exception as e:
+                # Log locally but don't crash
+                logger.warning(f"Error in Loki flush loop: {e}")
+                time.sleep(1.0)
+
+    def _flush_entries(self, entries: List[LogEntry]):
+        """Send a batch of log entries to Loki via the gateway."""
+        if not entries:
+            return
+
+        payload = {
+            "agent_name": self.agent_name,
+            "entries": [
+                {
+                    "timestamp": e.timestamp,
+                    "level": e.level,
+                    "message": e.message,
+                    "extra": e.extra,
+                }
+                for e in entries
+            ]
+        }
+
+        try:
+            # Prepare headers
+            headers = {"X-API-Key": self.gateway_client.api_key}
+
+            # Prepare body (with optional compression)
+            json_payload = json.dumps(payload).encode('utf-8')
+            body = json_payload
+
+            if self.enable_compression:
+                body = gzip.compress(json_payload)
+                headers["Content-Encoding"] = "gzip"
+                headers["Content-Type"] = "application/json"
+            else:
+                headers["Content-Type"] = "application/json"
+
+            # Use gateway client's HTTP URL
+            response = requests.post(
+                f"{self._get_http_url()}/api/v1/logs",
+                data=body,
+                headers=headers,
+                timeout=10.0,
+            )
+
+            if response.status_code in (200, 204):
+                # Success - reset retry delay
+                self._retry_delay = 1.0
+                self._last_error_time = None
+            else:
+                raise Exception(f"Gateway returned status {response.status_code}")
+
+        except Exception as e:
+            # Log warning locally (but avoid recursion by using direct stderr)
+            current_time = time.time()
+            if self._last_error_time is None or current_time - self._last_error_time > 60:
+                import sys
+                print(
+                    f"[LokiLogHandler] Failed to send {len(entries)} logs: {e}. "
+                    f"Retry delay: {self._retry_delay}s",
+                    file=sys.stderr
+                )
+                self._last_error_time = current_time
+
+            # Exponential backoff
+            time.sleep(self._retry_delay)
+            self._retry_delay = min(self._retry_delay * 2, self.retry_max_delay)
+
+            # Re-queue entries that failed (if buffer has space)
+            for entry in entries:
+                try:
+                    self._queue.put_nowait(entry)
+                except queue.Full:
+                    break  # Buffer full, drop oldest logs
+
+    def _get_http_url(self) -> str:
+        """Get HTTP URL from gateway client (handles protocol conversion)."""
+        if hasattr(self.gateway_client, 'http_url'):
+            return self.gateway_client.http_url
+
+        # Fallback: convert gRPC URL to HTTP URL
+        gateway_url = self.gateway_client.gateway_url
+        return gateway_url.replace("grpc://", "http://").replace("grpcs://", "https://")
+
+    def is_healthy(self) -> bool:
+        """
+        Check if handler is functioning (not backlogged).
+
+        Returns:
+            True if handler is running and buffer is not >80% full
+        """
+        return (
+            self._running and
+            self._queue.qsize() < self.max_buffer_size * 0.8
+        )
+
+    def emit(self, record: logging.LogRecord):
+        """
+        Emit a log record.
+
+        This method is called by the logging framework for each log message.
+        It converts the record to a LogEntry and queues it for sending.
+        """
+        try:
+            # Convert LogRecord to LogEntry
+            entry = self._record_to_entry(record)
+
+            # Try to add to queue (non-blocking)
+            try:
+                self._queue.put_nowait(entry)
+            except queue.Full:
+                # Buffer full - drop this log (graceful degradation)
+                pass
+
+        except Exception:
+            # Never crash due to logging
+            self.handleError(record)
+
+    def _record_to_entry(self, record: logging.LogRecord) -> LogEntry:
+        """Convert a logging.LogRecord to a LogEntry."""
+        # Get timestamp in nanoseconds
+        timestamp = int(record.created * 1_000_000_000)
+
+        # Normalize level name
+        level = record.levelname.lower()
+        if level == "warning":
+            level = "warn"
+
+        # Format message
+        message = self.format(record) if self.formatter else record.getMessage()
+
+        # Extract extra fields
+        extra: Dict[str, str] = {}
+
+        # Check for conversation_id in record
+        if hasattr(record, "conversation_id"):
+            extra["conversation_id"] = str(record.conversation_id)
+
+        # Check for user_id in record
+        if hasattr(record, "user_id"):
+            extra["user_id"] = str(record.user_id)
+
+        # Include logger name
+        extra["logger"] = record.name
+
+        # Include module/function info for debugging
+        if record.funcName and record.funcName != "<module>":
+            extra["function"] = record.funcName
+
+        # Include exception info if present
+        if record.exc_info and record.exc_info[0] is not None:
+            extra["exception_type"] = record.exc_info[0].__name__
+
+        return LogEntry(
+            timestamp=timestamp,
+            level=level,
+            message=message,
+            extra=extra,
+        )
+
+    def flush(self):
+        """Flush all buffered logs immediately."""
+        entries: List[LogEntry] = []
+        while True:
+            try:
+                entry = self._queue.get_nowait()
+                entries.append(entry)
+                self._queue.task_done()
+            except queue.Empty:
+                break
+
+        if entries:
+            self._flush_entries(entries)
+
+    def close(self):
+        """Close the handler and flush remaining logs."""
+        with self._lock:
+            if not self._running:
+                return
+
+            self._running = False
+
+        # Wait for flush thread to finish
+        if self._thread and self._thread.is_alive():
+            self._thread.join(timeout=5.0)
+
+        # Final flush
+        self.flush()
+
+        super().close()

basion_agent/memory.py

@@ -0,0 +1,73 @@
+"""Memory context for accessing AI memory from message handler."""
+
+from typing import List, Optional
+
+from .memory_client import MemoryClient, MemorySearchResult, UserSummary
+
+
+class Memory:
+    """User-facing memory context, attached to Message.
+
+    Provides easy access to memory functions scoped to the current user/conversation.
+    """
+
+    def __init__(self, user_id: str, conversation_id: str, client: MemoryClient):
+        self.user_id = user_id
+        self.conversation_id = conversation_id
+        self._client = client
+
+    async def query_about_user(
+        self,
+        query: str,
+        limit: int = 10,
+        threshold: int = 70,
+        context_messages: int = 0,
+    ) -> List[MemorySearchResult]:
+        """Query user's long-term memory with semantic search.
+
+        Args:
+            query: Semantic search query
+            limit: Maximum results (1-100)
+            threshold: Similarity threshold 0-100
+            context_messages: Surrounding messages to include (0-20)
+
+        Returns:
+            List of MemorySearchResult
+        """
+        return await self._client.search(
+            query=query,
+            user_id=self.user_id,
+            limit=limit,
+            min_similarity=threshold,
+            context_messages=context_messages,
+        )
+
+    async def query_about_conversation(
+        self,
+        query: str,
+        limit: int = 10,
+        threshold: int = 70,
+        context_messages: int = 0,
+    ) -> List[MemorySearchResult]:
+        """Query conversation's long-term memory with semantic search.
+
+        Args:
+            query: Semantic search query
+            limit: Maximum results (1-100)
+            threshold: Similarity threshold 0-100
+            context_messages: Surrounding messages to include (0-20)
+
+        Returns:
+            List of MemorySearchResult
+        """
+        return await self._client.search(
+            query=query,
+            conversation_id=self.conversation_id,
+            limit=limit,
+            min_similarity=threshold,
+            context_messages=context_messages,
+        )
+
+    async def get_user_summary(self) -> Optional[UserSummary]:
+        """Get user's summary (aggregated across all conversations)."""
+        return await self._client.get_user_summary(self.user_id)