memgraph-context-graph 0.5.0__py3-none-any.whl

memgraph_context_graph/events.py

@@ -0,0 +1,210 @@
+"""Common event protocol for agent runtime actions.
+
+This module defines runtime-agnostic event types that any adapter can emit
+and any graph connector can consume. Events capture the essential
+information about agent lifecycle moments without coupling to a specific runtime.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+
+
+class EventType(str, Enum):
+    """Types of events that can occur during an agent run."""
+
+    SESSION_START = "session_start"
+    SESSION_END = "session_end"
+    TURN_END = "turn_end"
+    TOOL_START = "tool_start"
+    TOOL_END = "tool_end"
+    AGENT_START = "agent_start"
+    AGENT_END = "agent_end"
+    LLM_START = "llm_start"
+    LLM_END = "llm_end"
+    HANDOFF = "handoff"
+    MESSAGE = "message"
+    ERROR = "error"
+
+
+@dataclass
+class Event:
+    """Base event emitted by runtime adapters.
+
+    Every concrete event subclass carries a fixed ``event_type`` and
+    runtime-agnostic payload fields. Adapters populate these from their
+    respective runtime callbacks or hook payloads, and connectors decide which fields are
+    relevant to them.
+    """
+
+    event_type: EventType
+    session_id: str
+    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    #: The SDK that produced this event (e.g. "claude", "openai").
+    source_sdk: str = ""
+
+
+# ------------------------------------------------------------------
+# Session events
+# ------------------------------------------------------------------
+
+
+@dataclass
+class SessionStartEvent(Event):
+    """Fired when an agent session begins."""
+
+    event_type: EventType = field(default=EventType.SESSION_START, init=False)
+    model: str | None = None
+    working_directory: str | None = None
+    user_id: str | None = None
+
+
+@dataclass
+class SessionEndEvent(Event):
+    """Fired when an agent session ends."""
+
+    event_type: EventType = field(default=EventType.SESSION_END, init=False)
+    status: str = "completed"
+    total_cost_usd: float | None = None
+    total_input_tokens: int | None = None
+    total_output_tokens: int | None = None
+
+
+@dataclass
+class TurnEndEvent(Event):
+    """Fired when the agent finishes one turn and the session stays open.
+
+    Runtimes without a session-end hook (Codex, Antigravity, ``opencode run``)
+    only ever signal this, so connectors must not treat it as the session
+    ending. ``reason`` is the runtime's own stop reason, when it gives one.
+    """
+
+    event_type: EventType = field(default=EventType.TURN_END, init=False)
+    reason: str | None = None
+
+
+# ------------------------------------------------------------------
+# Tool events
+# ------------------------------------------------------------------
+
+
+@dataclass
+class ToolStartEvent(Event):
+    """Fired before a tool/function is invoked."""
+
+    event_type: EventType = field(default=EventType.TOOL_START, init=False)
+    tool_name: str = ""
+    tool_input: Any = None
+    tool_use_id: str | None = None
+    agent_name: str | None = None
+
+
+@dataclass
+class ToolEndEvent(Event):
+    """Fired after a tool/function returns."""
+
+    event_type: EventType = field(default=EventType.TOOL_END, init=False)
+    tool_name: str = ""
+    tool_use_id: str | None = None
+    result: Any = None
+    is_error: bool = False
+    error_message: str | None = None
+    agent_name: str | None = None
+
+
+# ------------------------------------------------------------------
+# Agent events
+# ------------------------------------------------------------------
+
+
+@dataclass
+class AgentStartEvent(Event):
+    """Fired when an agent (or subagent) begins execution."""
+
+    event_type: EventType = field(default=EventType.AGENT_START, init=False)
+    agent_name: str = ""
+    agent_type: str = ""
+    parent_agent_name: str | None = None
+
+
+@dataclass
+class AgentEndEvent(Event):
+    """Fired when an agent (or subagent) finishes execution."""
+
+    event_type: EventType = field(default=EventType.AGENT_END, init=False)
+    agent_name: str = ""
+    agent_type: str = ""
+    output: Any = None
+
+
+# ------------------------------------------------------------------
+# LLM events
+# ------------------------------------------------------------------
+
+
+@dataclass
+class LLMStartEvent(Event):
+    """Fired before an LLM call is made."""
+
+    event_type: EventType = field(default=EventType.LLM_START, init=False)
+    agent_name: str | None = None
+    system_prompt: str | None = None
+    input_items: list[Any] = field(default_factory=list)
+
+
+@dataclass
+class LLMEndEvent(Event):
+    """Fired after an LLM call returns."""
+
+    event_type: EventType = field(default=EventType.LLM_END, init=False)
+    agent_name: str | None = None
+    model: str | None = None
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+    response: Any = None
+
+
+# ------------------------------------------------------------------
+# Communication events
+# ------------------------------------------------------------------
+
+
+@dataclass
+class HandoffEvent(Event):
+    """Fired when one agent hands off to another."""
+
+    event_type: EventType = field(default=EventType.HANDOFF, init=False)
+    from_agent: str = ""
+    to_agent: str = ""
+
+
+@dataclass
+class MessageEvent(Event):
+    """Fired when a message is sent (user, assistant, system)."""
+
+    event_type: EventType = field(default=EventType.MESSAGE, init=False)
+    role: str = ""
+    content: Any = None
+    model: str | None = None
+    agent_name: str | None = None
+
+
+# ------------------------------------------------------------------
+# Error events
+# ------------------------------------------------------------------
+
+
+@dataclass
+class ErrorOccurredEvent(Event):
+    """Fired when an error occurs during agent execution."""
+
+    event_type: EventType = field(default=EventType.ERROR, init=False)
+    error_type: str = ""
+    error_message: str = ""
+    error_details: dict[str, Any] = field(default_factory=dict)
+    recoverable: bool = True

memgraph_context_graph/hooks/__init__.py

@@ -0,0 +1 @@
+"""Command hook entry points for memgraph-context-graph runtimes."""

memgraph_context_graph/hooks/cli.py

@@ -0,0 +1,184 @@
+"""Generic command hook CLI for memgraph-context-graph runtimes."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+_HELP = """usage: memgraph-context-graph hook <command> [options]
+
+Bridge agent command hooks to memgraph-context-graph.
+
+Commands:
+  init <runtime>   Generate a runtime's private hook config, if it supports one.
+  run <runtime>    Run a runtime's command-hook adapter.
+
+Runtimes are discovered via the `memgraph_context_graph.runtimes` entry point --
+run `memgraph-context-graph hook run --help` (or see `runtime_plugin.py`) to see
+what's registered in this environment.
+"""
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Dispatch to a runtime-specific command hook adapter."""
+    args = list(argv) if argv is not None else sys.argv[1:]
+    if not args:
+        print(_HELP)
+        return 2
+    if args[0] in {"-h", "--help"}:
+        print(_HELP)
+        return 0
+
+    command = args[0]
+    if command == "init":
+        return _init(args[1:])
+
+    if command == "run":
+        if len(args) == 1:
+            print("usage: memgraph-context-graph hook run <runtime> [options]", file=sys.stderr)
+            return 2
+        runtime = args[1]
+        runtime_args = args[2:]
+        return _run_runtime(runtime, runtime_args)
+
+    # Backward-compatible module form:
+    # python -m memgraph_context_graph.hooks.cli codex --connector skills-graph
+    runtime = command
+    runtime_args = args[1:]
+    return _run_runtime(runtime, runtime_args)
+
+
+def _run_runtime(runtime: str, runtime_args: list[str]) -> int:
+    from memgraph_context_graph.hooks.runner import run_hook
+    from memgraph_context_graph.hooks.runtime_plugin import UnknownRuntimeError, get_runtime_plugin
+
+    try:
+        plugin = get_runtime_plugin(runtime)
+    except UnknownRuntimeError as exc:
+        print(str(exc), file=sys.stderr)
+        return 2
+
+    return run_hook(plugin, runtime_args)
+
+
+def _init(argv: list[str]) -> int:
+    if not argv:
+        print("usage: memgraph-context-graph hook init <runtime> [options]", file=sys.stderr)
+        return 2
+
+    from memgraph_context_graph.hooks.runtime_plugin import UnknownRuntimeError, get_runtime_plugin
+
+    runtime, init_argv = argv[0], argv[1:]
+    try:
+        plugin = get_runtime_plugin(runtime)
+    except UnknownRuntimeError as exc:
+        print(str(exc), file=sys.stderr)
+        return 2
+
+    if any(arg.split("=", 1)[0] == "--memory-backend" for arg in init_argv):
+        return _switch_memory_backend(plugin, init_argv)
+
+    init = getattr(plugin, "init", None)
+    if init is None:
+        print(f"{plugin.name} hook setup is not implemented yet.", file=sys.stderr)
+        return 2
+
+    return _run_generic_init(plugin.name, init, init_argv)
+
+
+def _switch_memory_backend(plugin: Any, argv: list[str]) -> int:
+    """``setup <runtime> --memory-backend context-graph|native``: whose memory the harness's model uses."""
+    from memgraph_context_graph.adapters._identity import MEMORY_BACKEND_GRAPH, MEMORY_BACKEND_NATIVE
+    from memgraph_context_graph.memory_backend import MemoryBackendError, use_graph_memory, use_native_memory
+
+    parser = argparse.ArgumentParser(
+        prog=f"memgraph-context-graph setup {plugin.name}",
+        description=(
+            f"Make Context Graph {plugin.name}'s memory (turning its built-in memory off and importing it), "
+            "or hand memory back to the built-in one."
+        ),
+    )
+    parser.add_argument("--memory-backend", required=True, choices=[MEMORY_BACKEND_GRAPH, MEMORY_BACKEND_NATIVE])
+    args = parser.parse_args(argv)
+    switch = use_graph_memory if args.memory_backend == MEMORY_BACKEND_GRAPH else use_native_memory
+    try:
+        lines = switch(plugin)
+    except MemoryBackendError as exc:
+        print(str(exc), file=sys.stderr)
+        return 1
+    print("\n".join(lines))
+    return 0
+
+
+def _run_generic_init(runtime_name: str, init: Any, argv: list[str]) -> int:
+    """Parse the shared ``hook init`` flags and delegate to *init*.
+
+    The flag set is generic across runtimes. Each plugin's ``init`` writes its
+    own config from the project dir, connectors, hook command, timeout, and
+    force flag; ``--setup-schema`` (with the ``--memgraph-*`` overrides) is
+    handled here, identically for every runtime.
+    """
+    parser = argparse.ArgumentParser(description=f"Generate a private {runtime_name} hook config.")
+    parser.add_argument(
+        "--connector",
+        action="append",
+        default=None,
+        help="Graph connector to enable. Defaults to skills-graph, actions-graph, sessions-graph.",
+    )
+    parser.add_argument(
+        "--project-dir",
+        default=".",
+        help="Project directory where the runtime's config should be generated.",
+    )
+    parser.add_argument(
+        "--hook-command",
+        default=None,
+        help="Full command to place in the runtime's hook config. Defaults to this installed CLI.",
+    )
+    parser.add_argument("--memgraph-url", default=None, help="Memgraph Bolt URL for --setup-schema.")
+    parser.add_argument("--memgraph-user", default=None, help="Memgraph username for --setup-schema.")
+    parser.add_argument(
+        "--memgraph-password", default=None, help="Memgraph password for --setup-schema. Never written to disk."
+    )
+    parser.add_argument("--memgraph-database", default=None, help="Memgraph database for --setup-schema.")
+    parser.add_argument(
+        "--setup-schema",
+        action="store_true",
+        help="Connect to Memgraph now and initialize enabled graph connector schemas.",
+    )
+    parser.add_argument("--timeout", type=int, default=30, help="Hook timeout in seconds.")
+    parser.add_argument("--force", action="store_true", help="Overwrite existing generated config.")
+    args = parser.parse_args(argv)
+
+    connectors = args.connector or ["skills-graph", "actions-graph", "sessions-graph"]
+    project_dir = Path(args.project_dir).expanduser().resolve()
+
+    try:
+        init(project_dir, connectors, hook_command=args.hook_command, timeout=args.timeout, force=args.force)
+    except FileExistsError as exc:
+        print(str(exc), file=sys.stderr)
+        print("Re-run with --force to replace generated files.", file=sys.stderr)
+        return 1
+
+    if args.setup_schema:
+        from memgraph_context_graph.adapters._identity import resolve_memgraph_env
+        from memgraph_context_graph.hooks.runner import setup_connector_schemas
+
+        memgraph_env = resolve_memgraph_env(
+            url=args.memgraph_url,
+            user=args.memgraph_user,
+            password=args.memgraph_password,
+            database=args.memgraph_database,
+        )
+        setup_connector_schemas(connectors, memgraph_env=memgraph_env)
+        print(f"Initialized schema on {memgraph_env['MEMGRAPH_URL']} ({memgraph_env['MEMGRAPH_DATABASE']})")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())