aixtools 0.1.0__py3-none-any.whl

aixtools/logging/mcp_logger.py

@@ -0,0 +1,172 @@
+"""
+Logger implementations for MCP server.
+"""
+
+import json
+from abc import ABC, abstractmethod
+from datetime import datetime
+from pathlib import Path
+
+from aixtools.logging.logging_config import get_logger
+
+from .mcp_log_models import CodeLogEntry, CommandLogEntry, LogEntry, ServiceLogEntry, SystemLogEntry
+
+logger = get_logger(__name__)
+
+
+def log_with_default_logger(entry: LogEntry) -> None:
+    """
+    Formats a log entry into a human-readable string and logs it.
+    """
+    if isinstance(entry, SystemLogEntry):
+        logger.info("%s: %s", entry.event, entry.details or "")
+    elif isinstance(entry, ServiceLogEntry):
+        logger.info("%s: %s", entry.event, entry.details or "")
+    elif isinstance(entry, CodeLogEntry):
+        logger.info("%s code: %s", entry.language, entry.code)
+    elif isinstance(entry, CommandLogEntry):
+        logger.info("%s, CWD: %s", entry.command, entry.working_directory)
+    else:
+        logger.debug("Logging entry: %s", entry.model_dump_json(indent=2))
+
+
+class McpLogger(ABC):
+    """Abstract base class for loggers."""
+
+    @abstractmethod
+    def log(self, entry: LogEntry) -> None:
+        """Log an entry."""
+
+    @abstractmethod
+    def get_logs(  # noqa: PLR0913  # pylint: disable=too-many-arguments,too-many-positional-arguments
+        self,
+        user_id: str | None = None,
+        session_id: str | None = None,
+        container_id: str | None = None,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        limit: int = 100,
+    ) -> list[LogEntry]:
+        """Get logs with optional filters."""
+
+
+class JSONFileMcpLogger(McpLogger):
+    """Logger that stores logs in a single JSON file."""
+
+    def __init__(self, log_dir: str | Path):
+        """Initialize the logger."""
+        self.log_dir = Path(log_dir)
+        self.log_dir.mkdir(parents=True, exist_ok=True)
+        self.log_file_path = self.log_dir / "mcp_logs.jsonl"
+        self.log_file = open(self.log_file_path, "a", encoding="utf-8")  # pylint: disable=consider-using-with
+
+    def __del__(self):
+        """Ensure file is closed when the logger is destroyed."""
+        if hasattr(self, "log_file") and self.log_file and not self.log_file.closed:
+            self.log_file.close()
+
+    def _get_log_file_path(self) -> Path:
+        """Get the path to the log file."""
+        return self.log_file_path
+
+    def log(self, entry: LogEntry) -> None:
+        """Log an entry to the JSON file."""
+        log_with_default_logger(entry)
+        # Convert the entry to a JSON string
+        entry_json = entry.model_dump_json()
+        # Append the entry to the log file and flush to ensure it's written
+        self.log_file.write(entry_json + "\n")
+        self.log_file.flush()
+
+    def get_logs(  # noqa: PLR0913, PLR0912  # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-branches
+        self,
+        user_id: str | None = None,
+        session_id: str | None = None,
+        container_id: str | None = None,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        limit: int = 100,
+    ) -> list[LogEntry]:
+        """Get logs with optional filters.
+
+        Args:
+            user_id: Filter by user ID.
+            session_id: Filter by session ID.
+            container_id: Filter by container ID.
+            start_time: Filter by start time.
+            end_time: Filter by end time.
+            limit: Maximum number of logs to return.
+
+        Returns:
+            List of log entries.
+        """
+        logs: list[LogEntry] = []
+
+        # Ensure any pending writes are flushed to disk
+        self.log_file.flush()
+
+        # Read from the single log file
+        if self.log_file_path.exists():
+            with open(self.log_file_path, "r", encoding="utf-8") as f:
+                for line in f:
+                    try:
+                        # Parse the JSON entry
+                        entry_dict = json.loads(line)
+
+                        # Convert timestamp string to datetime
+                        entry_dict["timestamp"] = datetime.fromisoformat(entry_dict["timestamp"])
+
+                        # Apply filters
+                        if user_id and entry_dict.get("user_id") != user_id:
+                            continue
+                        if session_id and entry_dict.get("session_id") != session_id:
+                            continue
+                        if container_id and entry_dict.get("container_id") != container_id:
+                            continue
+                        if start_time and entry_dict["timestamp"] < start_time:
+                            continue
+                        if end_time and entry_dict["timestamp"] > end_time:
+                            continue
+
+                        # Create the appropriate log entry object based on log_type
+                        log_type = entry_dict["log_type"]
+                        if log_type == "command":
+                            entry = CommandLogEntry(**entry_dict)
+                        elif log_type == "code":
+                            entry = CodeLogEntry(**entry_dict)
+                        elif log_type == "system":
+                            entry = SystemLogEntry(**entry_dict)
+                        else:
+                            continue
+
+                        logs.append(entry)
+                    except (json.JSONDecodeError, KeyError) as e:
+                        logger.error("Error parsing log entry: %s", e)
+                        continue
+
+                    # Check if we've reached the limit
+                    if len(logs) >= limit:
+                        break
+
+        # Sort logs by timestamp (newest first)
+        logs.sort(key=lambda x: x.timestamp, reverse=True)
+
+        return logs[:limit]
+
+
+# Global logger instance
+_mcp_logger: McpLogger | None = None
+
+
+def initialize_mcp_logger(mcp_logger: McpLogger) -> None:
+    """Initialize the MCP logger"""
+    global _mcp_logger  # noqa: PLW0603, pylint: disable=global-statement
+    _mcp_logger = mcp_logger
+
+
+def get_mcp_logger() -> McpLogger:
+    """Get the global logger for MCP server."""
+    if _mcp_logger is None:
+        raise RuntimeError("Logger not initialized")
+
+    return _mcp_logger

aixtools/logging/model_patch_logging.py

@@ -0,0 +1,87 @@
+"""
+Logging utilities for model patching and request/response tracking.
+"""
+
+import functools
+from contextlib import asynccontextmanager
+from uuid import uuid4
+
+from aixtools.logging.logging_config import get_logger
+from aixtools.model_patch.model_patch import (
+    ModelRawRequest,
+    ModelRawRequestResult,
+    ModelRawRequestYieldItem,
+    get_request_fn,
+    get_request_stream_fn,
+    model_patch,
+)
+
+logger = get_logger(__name__)
+
+
+def log_async_method(fn, agent_logger):
+    """Log async method calls"""
+
+    @functools.wraps(fn)
+    async def model_request_logger_wrapper(*args, **kwargs):
+        # Log request
+        uuid = str(uuid4())  # Create a unique ID for this request
+        log_object = ModelRawRequest(method_name=fn.__name__, request_id=uuid, args=args, kwargs=kwargs)
+        agent_logger.log(log_object)
+        # Invoke the original method
+        try:
+            result = await fn(*args, **kwargs)
+            # Log results
+            log_object = ModelRawRequestResult(method_name=fn.__name__, request_id=uuid, result=result)
+            agent_logger.log(log_object)
+        except Exception as e:
+            # Log exception
+            agent_logger.log(e)
+            raise e
+        return result
+
+    return model_request_logger_wrapper
+
+
+def log_async_stream(fn, agent_logger):
+    """Log async streaming method calls with individual item tracking."""
+
+    @functools.wraps(fn)
+    @asynccontextmanager
+    async def model_request_stream_logger_wrapper(*args, **kwargs):
+        # Log request
+        uuid = str(uuid4())  # Create a unique ID for this request
+        log_object = ModelRawRequest(method_name=fn.__name__, request_id=uuid, args=args, kwargs=kwargs)
+        agent_logger.log(log_object)
+        # Invoke the original method
+        async with fn(*args, **kwargs) as stream:
+
+            async def gen():
+                item_num = 0
+                try:
+                    async for item in stream:
+                        # Log yielded items
+                        log_object = ModelRawRequestYieldItem(
+                            method_name=fn.__name__, request_id=uuid, item=item, item_num=item_num
+                        )
+                        agent_logger.log(log_object)
+                        item_num += 1
+                        yield item
+                except Exception as e:
+                    # Log exception
+                    agent_logger.log(e)
+                    raise e
+
+            yield gen()
+
+    return model_request_stream_logger_wrapper
+
+
+def model_patch_logging(model, agent_logger):
+    """Patch model with logging methods"""
+    logger.debug("Patching model with logging")
+    return model_patch(
+        model,
+        request_method=log_async_method(get_request_fn(model), agent_logger),
+        request_stream_method=log_async_stream(get_request_stream_fn(model), agent_logger),
+    )

aixtools/logging/open_telemetry.py

@@ -0,0 +1,36 @@
+"""
+OpenTelemetry integration for logging and tracing agent operations.
+"""
+
+import os
+
+import logfire  # pylint: disable=import-error
+from pydantic_ai import Agent
+
+from aixtools.utils.config import LOGFIRE_TOKEN, LOGFIRE_TRACES_ENDPOINT
+
+
+def open_telemetry_on():
+    """Configure and enable OpenTelemetry tracing with LogFire integration."""
+    service_name = "agent_poc"
+
+    if LOGFIRE_TRACES_ENDPOINT:
+        os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = LOGFIRE_TRACES_ENDPOINT
+        logfire.configure(
+            service_name=service_name,
+            # Sending to Logfire is on by default regardless of the OTEL env vars.
+            # Keep this line here if you don't want to send to both Jaeger and Logfire.
+            send_to_logfire=False,
+        )
+        Agent.instrument_all(True)
+        return
+
+    if LOGFIRE_TOKEN:
+        logfire.configure(
+            token=LOGFIRE_TOKEN,
+            service_name=service_name,
+        )
+        Agent.instrument_all(True)
+        return
+
+    print("OpenTelemetry is not enabled. Set the LOGFIRE_TOKEN or LOGFIRE_TRACES_ENDPOINT environment variable.")

aixtools/mcp/__init__.py

@@ -0,0 +1,9 @@
+"""
+Model Context Protocol (MCP) implementation for AI agent communication.
+"""
+
+from aixtools.mcp.fast_mcp_log import FastMcpLog
+
+__all__ = [
+    "FastMcpLog",
+]

aixtools/mcp/example_client.py

@@ -0,0 +1,30 @@
+"""
+Example client implementation for Model Context Protocol (MCP) servers.
+"""
+
+import asyncio
+
+from pydantic_ai.mcp import MCPServerSSE, MCPServerStdio
+
+from aixtools.agents import get_agent, run_agent
+
+USE_SEE = False
+
+
+if USE_SEE:
+    server = MCPServerSSE(url="http://127.0.0.1:8000/sse")
+else:
+    server = MCPServerStdio(command="fastmcp", args=["run", "aixtools/mcp/example_server.py"])
+
+
+async def main(agent, prompt):  # pylint: disable=redefined-outer-name
+    """Run an agent with MCP servers and display the result."""
+    async with agent:
+        ret = await run_agent(agent, prompt)
+        print(f"Agent returned: {ret}")
+
+
+if __name__ == "__main__":
+    agent = get_agent(mcp_servers=[server])  # pylint: disable=unexpected-keyword-arg
+    print(f"Agent created: {agent}")
+    asyncio.run(main(agent, "What is the add of 923048502345 and 795467090123481926349123941 ?"))

aixtools/mcp/example_server.py

@@ -0,0 +1,22 @@
+"""
+Example server implementation for Model Context Protocol (MCP).
+"""
+
+from aixtools.mcp.fast_mcp_log import FastMcpLog
+
+mcp = FastMcpLog("Demo")
+# mcp = FastMCP("Demo")
+
+
+# Add an addition tool
+@mcp.tool()
+def add(a: int, b: int) -> int:
+    """Add two numbers"""
+    return a + b
+
+
+# Add a dynamic greeting resource
+@mcp.resource("greeting://{name}")
+def get_greeting(name: str) -> str:
+    """Get a personalized greeting"""
+    return f"Hello, {name}!"

aixtools/mcp/fast_mcp_log.py

@@ -0,0 +1,31 @@
+"""
+FastMCP logging implementation for Model Context Protocol.
+"""
+
+import sys
+from typing import Any
+
+from fastmcp import FastMCP
+from pydantic import AnyUrl
+
+
+class FastMcpLog(FastMCP):
+    """A FastMCP with hooks for logging."""
+
+    async def _call_tool(self, key: str, arguments: dict[str, Any]):
+        print(f"Calling tool: {key} with arguments: {arguments}", file=sys.stderr)
+        ret = await super()._call_tool(key, arguments)
+        print(f"Tool returned: {ret}", file=sys.stderr)
+        return ret
+
+    async def _read_resource(self, uri: AnyUrl | str):
+        print(f"Reading resource: {uri}", file=sys.stderr)
+        ret = await super()._read_resource(uri)
+        print(f"Resource contents: {ret}", file=sys.stderr)
+        return ret
+
+    async def get_prompt(self, key: str, arguments: dict[str, Any] | None = None):
+        print(f"Getting prompt: {key} with arguments: {arguments}", file=sys.stderr)
+        ret = await super()._get_prompt(key, arguments)
+        print(f"Prompt result: {ret}", file=sys.stderr)
+        return ret