aixtools 0.0.0__py3-none-any.whl

aixtools/db/vector_db.py

@@ -0,0 +1,115 @@
+"""
+Vector database implementation for embedding storage and similarity search.
+"""
+
+from langchain_chroma import Chroma
+from langchain_core.embeddings import Embeddings
+from langchain_ollama import OllamaEmbeddings
+from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
+
+from aixtools.logging.logging_config import get_logger
+from aixtools.utils.config import (
+    AZURE_OPENAI_API_KEY,
+    AZURE_VDB_EMBEDDINGS_MODEL_NAME,
+    OLLAMA_VDB_EMBEDDINGS_MODEL_NAME,
+    OPENAI_API_KEY,
+    OPENAI_VDB_EMBEDDINGS_MODEL_NAME,
+    VDB_CHROMA_PATH,
+    VDB_EMBEDDINGS_MODEL_FAMILY,
+)
+
+CREATE_DB = False
+
+_vector_dbs = {}
+
+logger = get_logger(__name__)
+
+
+def get_vdb_embedding(model_family=VDB_EMBEDDINGS_MODEL_FAMILY) -> Embeddings:
+    """Get the embedding model for vector storage"""
+    match model_family:
+        case "openai":
+            return OpenAIEmbeddings(model=OPENAI_VDB_EMBEDDINGS_MODEL_NAME, api_key=OPENAI_API_KEY)  # type: ignore
+        case "azure":
+            return AzureOpenAIEmbeddings(  # type: ignore
+                model=AZURE_VDB_EMBEDDINGS_MODEL_NAME, api_key=AZURE_OPENAI_API_KEY
+            )
+        case "ollama":
+            return OllamaEmbeddings(model=OLLAMA_VDB_EMBEDDINGS_MODEL_NAME)  # type: ignore
+        case _:
+            raise ValueError(f"Model family {model_family} not supported")
+
+
+def get_vector_db(collection_name: str) -> Chroma:
+    """Implement singleton pattern for database connections"""
+    # _vector_dbs will not be re-assigned, but it will be modified
+    global _vector_dbs  # noqa: PLW0602, pylint: disable=protected-access,global-variable-not-assigned
+    if collection_name not in _vector_dbs:
+        print(f"Creating new DB connection: {collection_name=}")
+        vdb = Chroma(
+            persist_directory=str(VDB_CHROMA_PATH),
+            collection_name=collection_name,
+            embedding_function=get_vdb_embedding(),
+        )
+        _vector_dbs[collection_name] = vdb
+    return _vector_dbs[collection_name]
+
+
+def vdb_add(vdb: Chroma, text: str, doc_id: str, meta=list[dict] | dict | None, force=False) -> str | None:
+    """
+    Add a document to the database if it's not already there.
+    """
+    if not force and vdb_has_id(vdb, doc_id):
+        return None  # Document already exists, return None
+    if isinstance(meta, list):
+        metadatas = meta
+    elif isinstance(meta, dict):
+        metadatas = [meta]
+    else:
+        metadatas = None
+    ids = vdb.add_texts(texts=[text], ids=[doc_id], metadatas=metadatas)  # type: ignore
+    if not ids:
+        return None
+    return ids[0]  # Return the id of the added document
+
+
+def vdb_get_by_id(vdb: Chroma, doc_id: str):
+    """Get document with by id"""
+    collection = vdb._collection  # pylint: disable=protected-access
+    return collection.get(ids=[doc_id])  # query by id
+
+
+def vdb_has_id(vdb: Chroma, doc_id: str):
+    """Check if a document with a given id exists in the database"""
+    result = vdb_get_by_id(vdb, doc_id)
+    return len(result["ids"]) > 0
+
+
+# Load database
+def vdb_query(  # noqa: PLR0913, pylint: disable=too-many-arguments,too-many-positional-arguments
+    vdb: Chroma,
+    query: str,
+    filter: dict[str, str] | None = None,  # pylint: disable=redefined-builtin
+    where_document: dict[str, str] | None = None,
+    max_items=10,
+    similarity_threshold=None,
+):
+    """
+    Query vector database with a given query, return top k results.
+    Args:
+        query: str, query string
+        max_items: int, maximum number of items to return
+        similarity_threshold: float, similarity threshold to filter the results
+    """
+    results = vdb.similarity_search_with_relevance_scores(
+        query, k=max_items, filter=filter, where_document=where_document
+    )
+    logger.debug(
+        "Got %s results before filter, first one's similarity score is: %s",
+        len(results),
+        results[0][1] if results else None,
+    )
+    if similarity_threshold is not None:
+        results = [(doc_id, score) for doc_id, score in results if score > similarity_threshold]
+        print(f"Got {len(results)} results after filter")
+    return results

aixtools/google/client.py

@@ -0,0 +1,25 @@
+import os
+from pathlib import Path
+
+from google import genai
+
+from aixtools.logging.logging_config import get_logger
+from aixtools.utils.config import GOOGLE_CLOUD_LOCATION, GOOGLE_CLOUD_PROJECT, GOOGLE_GENAI_USE_VERTEXAI
+
+logger = get_logger(__name__)
+
+
+def get_genai_client(service_account_key_path: Path | None = None) -> genai.Client:
+    """Initialize and return a Google GenAI client using Vertex AI / Gemini Developer API."""
+    assert GOOGLE_CLOUD_PROJECT, "GOOGLE_CLOUD_PROJECT is not set"
+    assert GOOGLE_CLOUD_LOCATION, "GOOGLE_CLOUD_LOCATION is not set"
+    if service_account_key_path:
+        if not service_account_key_path.exists():
+            raise FileNotFoundError(f"Service account key file not found: {service_account_key_path}")
+        logger.info(f"✅ GCP Service Account Key File: {service_account_key_path}")
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = str(service_account_key_path)
+    return genai.Client(
+        vertexai=GOOGLE_GENAI_USE_VERTEXAI,
+        project=GOOGLE_CLOUD_PROJECT,
+        location=GOOGLE_CLOUD_LOCATION,
+    )

aixtools/log_view/__init__.py

@@ -0,0 +1,17 @@
+"""
+Streamlit application to visualize agent nodes from log files.
+
+This package provides tools to:
+- View the most recent log file by default
+- Open and analyze other log files
+- Visualize nodes from agent runs with expandable/collapsible sections
+- Filter nodes by various criteria
+- Export visualizations
+"""
+
+from aixtools.log_view.app import main, main_cli
+
+__all__ = [
+    "main",
+    "main_cli",
+]

aixtools/log_view/app.py

@@ -0,0 +1,195 @@
+"""
+Main application module for the Agent Log Viewer.
+"""
+
+import argparse
+import os
+import subprocess
+from pathlib import Path
+
+import streamlit as st
+
+from aixtools.log_view.display import display_node
+from aixtools.log_view.export import export_nodes_to_json
+from aixtools.log_view.filters import filter_nodes
+from aixtools.log_view.log_utils import format_timestamp_from_filename, get_log_files
+from aixtools.log_view.node_summary import NodeTitle, extract_node_types
+
+# Now we can import our modules
+from aixtools.logging.log_objects import load_from_log
+from aixtools.utils.config import LOGS_DIR
+
+
+def main(log_dir: Path | None = None):  # noqa: PLR0915, pylint: disable=too-many-locals,too-many-statements
+    """Main function to run the Streamlit app."""
+    st.set_page_config(
+        page_title="Agent Log Viewer",
+        layout="wide",
+    )
+
+    st.title("Agent Log Viewer")
+
+    # Use provided log directory or default
+    if log_dir is None:
+        log_dir = LOGS_DIR
+
+    # Create the logs directory if it doesn't exist
+    os.makedirs(log_dir, exist_ok=True)
+
+    st.sidebar.header("Settings")
+
+    # Allow user to select a different log directory
+    custom_log_dir = st.sidebar.text_input("Log Directory", value=str(log_dir))
+    if custom_log_dir and custom_log_dir != str(log_dir):
+        log_dir = Path(custom_log_dir)
+        # Create the custom directory if it doesn't exist
+        os.makedirs(log_dir, exist_ok=True)
+
+    # Get log files
+    log_files = get_log_files(log_dir)
+
+    if not log_files:
+        st.warning(f"No log files found in {log_dir}")
+        st.info("Run an agent with logging enabled to create log files.")
+        return
+
+    # Create a dictionary of log files with formatted timestamps as display names
+    log_file_options = {f"{format_timestamp_from_filename(f.name)} - {f.name}": f for f in log_files}
+
+    # Select log file (default to most recent)
+    selected_log_file_name = st.sidebar.selectbox(
+        "Select Log File",
+        options=list(log_file_options.keys()),
+        index=0,
+    )
+
+    selected_log_file = log_file_options[selected_log_file_name]
+
+    st.sidebar.info(f"Selected: {selected_log_file.name}")
+
+    # Load nodes
+    try:
+        with st.spinner("Loading log file..."):
+            nodes = load_from_log(selected_log_file)
+
+        st.success(f"Loaded {len(nodes)} nodes from {selected_log_file.name}")
+
+        # Create filter section in sidebar
+        st.sidebar.header("Filters")
+
+        # Text filter
+        filter_text = st.sidebar.text_input("Text Search", help="Filter nodes containing this text")
+
+        # Extract node types for filtering
+        node_types = extract_node_types(nodes)
+
+        # Type filter
+        selected_types = st.sidebar.multiselect(
+            "Node Types", options=sorted(node_types), default=[], help="Select node types to display"
+        )
+
+        # Attribute filter
+        filter_attribute = st.sidebar.text_input("Has Attribute", help="Filter nodes that have this attribute")
+
+        # Regex filter
+        filter_regex = st.sidebar.text_input("Regex Pattern", help="Filter nodes matching this regex pattern")
+
+        # Combine all filters
+        filters = {"text": filter_text, "types": selected_types, "attribute": filter_attribute, "regex": filter_regex}
+
+        # Apply filters
+        filtered_nodes = filter_nodes(nodes, filters)
+
+        # Show filter results
+        if len(filtered_nodes) != len(nodes):
+            st.info(f"Filtered to {len(filtered_nodes)} of {len(nodes)} nodes")
+
+        # Display options
+        st.sidebar.header("Display Options")
+
+        # Option to expand all nodes by default
+        expand_all = st.sidebar.checkbox("Expand All Nodes", value=False)
+
+        # Option to select output format
+        display_format = st.sidebar.radio(
+            "Display Format",
+            options=["Markdown", "Rich", "JSON"],
+            index=0,
+            help="Select the format for displaying node content",
+        )
+
+        # Export options
+        st.sidebar.header("Export")
+
+        # Export to JSON
+        if st.sidebar.button("Export to JSON"):
+            json_str = export_nodes_to_json(filtered_nodes)
+            st.sidebar.download_button(
+                label="Download JSON",
+                data=json_str,
+                file_name=f"agent_nodes_{selected_log_file.stem}.json",
+                mime="application/json",
+            )
+
+        # Main content area - display nodes
+        if filtered_nodes:
+            node_title = NodeTitle()
+            # Display nodes with proper formatting
+            for i, node in enumerate(filtered_nodes):
+                # Create a header for each node
+                node_header = f"{i}: {node_title.summary(node)}"
+
+                # Display the node with proper formatting
+                with st.expander(node_header, expanded=expand_all):
+                    try:
+                        display_node(node, display_format=display_format)
+                    except Exception as e:  # pylint: disable=broad-exception-caught
+                        st.error(f"Error displaying node: {e}")
+                        st.exception(e)
+        else:
+            st.warning("No nodes match the current filters")
+
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        st.error(f"Error loading or processing log file: {e}")
+        st.exception(e)
+
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(description="Agent Log Viewer")
+    parser.add_argument("log_dir", nargs="?", type=Path, help="Directory containing log files (default: DATA_DIR/logs)")
+    return parser.parse_args()
+
+
+def main_cli():
+    """Entry point for the command-line tool."""
+    cmd_args = parse_args()
+
+    # Print a message to indicate the app is starting
+    print("Starting Agent Log Viewer...")
+    print(f"Log directory: {cmd_args.log_dir or LOGS_DIR}")
+
+    # Launch the Streamlit app
+
+    # Get the path to this script
+    script_path = Path(__file__).resolve()
+
+    # Use streamlit run to start the app
+    cmd = ["streamlit", "run", str(script_path)]
+
+    # Add log_dir argument if provided
+    if cmd_args.log_dir:
+        cmd.extend(["--", str(cmd_args.log_dir)])
+
+    # Run the command
+    try:
+        subprocess.run(cmd, check=False)
+    except KeyboardInterrupt:
+        print("\nShutting down Agent Log Viewer...")
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        print(f"Error running Streamlit app: {e}")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args.log_dir)

aixtools/log_view/display.py

@@ -0,0 +1,285 @@
+"""
+Functions for displaying nodes in the Streamlit interface.
+Provides enhanced display capabilities for various object types,
+including dataclasses, with proper handling of nested structures.
+"""
+
+import inspect
+import json
+from dataclasses import fields as dataclass_fields
+from dataclasses import is_dataclass
+
+import pandas as pd
+import streamlit as st
+from rich.console import Console
+
+from aixtools.utils.utils import prepend_all_lines
+
+# Toggle for using markdown display instead of JSON
+USE_MARKDOWN = True
+
+
+def filter_private_fields(data_dict: dict) -> dict:
+    """Filter out private fields from the data dictionary."""
+    return {k: v for k, v in data_dict.items() if not k.startswith("_")}
+
+
+def filter_private_attributes(obj) -> dict:
+    """
+    Filter out private attributes and methods from an object.
+    Returns a dictionary of public attributes and their values.
+    """
+    if not hasattr(obj, "__dict__"):
+        return {}
+
+    result = {}
+    for attr, value in vars(obj).items():
+        if not attr.startswith("_"):
+            result[attr] = value
+
+    return result
+
+
+def is_method(obj, attr_name: str) -> bool:
+    """Check if an attribute is a method."""
+    try:
+        attr = getattr(obj, attr_name)
+        return inspect.ismethod(attr) or inspect.isfunction(attr)
+    except (AttributeError, TypeError):
+        return False
+
+
+def get_object_type_str(obj) -> str:  # noqa: PLR0911, pylint: disable=too-many-return-statements
+    """Get a string representation of the object's type."""
+    if obj is None:
+        return "null"
+    if isinstance(obj, bool):
+        return "bool"
+    if isinstance(obj, int):
+        return "int"
+    if isinstance(obj, float):
+        return "float"
+    if isinstance(obj, str):
+        return "str"
+    if isinstance(obj, list):
+        return f"list[{len(obj)}]"
+    if isinstance(obj, tuple):
+        return f"tuple[{len(obj)}]"
+    if isinstance(obj, dict):
+        return f"dict[{len(obj)}]"
+    if isinstance(obj, set):
+        return f"set[{len(obj)}]"
+    if is_dataclass(obj):
+        return f"dataclass:{type(obj).__name__}"
+    if hasattr(obj, "__dict__"):
+        return type(obj).__name__
+
+    return type(obj).__name__
+
+
+def object_to_json_with_types(obj, max_depth: int = 5, current_depth: int = 0):  # noqa: PLR0911, PLR0912, pylint: disable=too-many-return-statements,too-many-branches
+    """
+    Convert an object to a JSON-serializable dictionary with type information.
+    Handles nested objects up to max_depth.
+    """
+    # Prevent infinite recursion
+    if current_depth > max_depth:
+        return {"__type": get_object_type_str(obj), "__value": str(obj)}
+
+    # Handle None
+    if obj is None:
+        return None
+
+    # Handle basic types
+    if isinstance(obj, (bool, int, float, str)):
+        return obj
+
+    # Handle lists and tuples
+    if isinstance(obj, (list, tuple)):
+        items = []
+        for item in obj:
+            items.append(object_to_json_with_types(item, max_depth, current_depth + 1))
+        return items
+
+    # Handle dictionaries
+    if isinstance(obj, dict):
+        result = {}
+        for key, value in filter_private_fields(obj).items():
+            result[key] = object_to_json_with_types(value, max_depth, current_depth + 1)
+        return result
+
+    # Handle sets
+    if isinstance(obj, set):
+        items = []
+        for item in obj:
+            items.append(object_to_json_with_types(item, max_depth, current_depth + 1))
+        return {"__type": "set", "__items": items}
+
+    # Handle dataclasses
+    if is_dataclass(obj):
+        result = {"__type": f"dataclass:{type(obj).__name__}"}
+        for field in dataclass_fields(obj):
+            if field.name.startswith("_"):  # Skip private fields
+                continue
+            if not hasattr(obj, field.name):  # Skip not found
+                continue
+            value = getattr(obj, field.name)
+            result[field.name] = object_to_json_with_types(value, max_depth, current_depth + 1)
+        return result
+
+    # Handle objects with __dict__
+    if hasattr(obj, "__dict__"):
+        result = {"__type": type(obj).__name__}
+        for attr, value in filter_private_attributes(obj).items():
+            if not is_method(obj, attr):  # Skip methods
+                result[attr] = object_to_json_with_types(value, max_depth, current_depth + 1)
+        return result
+
+    # Handle other types
+    return {"__type": get_object_type_str(obj), "__value": str(obj)}
+
+
+def object_to_markdown(  # noqa: PLR0911, PLR0912, PLR0915, pylint: disable=too-many-locals,too-many-return-statements,too-many-branches,too-many-statements
+    obj, max_depth: int = 5, current_depth: int = 0, indent: str = ""
+) -> str:
+    """
+    Convert an object to a compact markdown representation.
+    Handles nested objects up to max_depth.
+    """
+    max_display_items = 10  # Show only first MAX_DISPLAY_ITEMS items for large collections, dicts, and sets
+
+    # Prevent infinite recursion
+    if current_depth > max_depth:
+        return f"`{get_object_type_str(obj)}`: {str(obj)}"
+
+    # Handle None
+    if obj is None:
+        return "`None`"
+
+    # Handle basic types
+    if isinstance(obj, bool):
+        return f"`{str(obj).lower()}`"
+
+    if isinstance(obj, (int, float)):
+        return f"`{obj}`"
+
+    if isinstance(obj, str):
+        lines = str(obj).splitlines()
+        if len(lines) > 1:
+            return f"\n{indent}```\n{prepend_all_lines(obj, prepend=indent)}\n{indent}```\n"
+        return obj
+
+    # Handle lists and tuples
+    if isinstance(obj, (list, tuple)):
+        if not obj:  # Empty collection
+            return f"`{get_object_type_str(obj)}`: empty"
+
+        max_inline_length = 3  # For small collections, show inline
+        if (
+            len(obj) <= max_inline_length
+            and current_depth > 0
+            and all(isinstance(x, (bool, int, float, str, type(None))) for x in obj)
+        ):
+            items = [object_to_markdown(item, max_depth, current_depth + 1) for item in obj]
+            return f"`{get_object_type_str(obj)}`: [{', '.join(items)}]"
+
+        # For larger collections, use bullet points
+        result = [f"`{get_object_type_str(obj)}`:"]
+        for i, item in enumerate(obj):
+            if i >= max_display_items and len(obj) > max_display_items + 2:
+                result.append(f"{indent}* ... ({len(obj) - 10} more items)")
+                break
+            item_md = object_to_markdown(item, max_depth, current_depth + 1, indent + "  ")
+            result.append(f"{indent}* {item_md}")
+        return "\n".join(result)
+
+    # Handle dictionaries
+    if isinstance(obj, dict):
+        if not obj:  # Empty dict
+            return "`dict`: empty"
+
+        result = [f"`dict[{len(obj)}]`:"]
+        for i, (key, value) in enumerate(filter_private_fields(obj).items()):
+            if i >= max_display_items and len(obj) > max_display_items + 2:
+                result.append(f"{indent}* ... ({len(obj) - 10} more items)")
+                break
+            value_md = object_to_markdown(value, max_depth, current_depth + 1, indent + "  ")
+            result.append(f"{indent}* **{key}**: {value_md}")
+        return "\n".join(result)
+
+    # Handle sets
+    if isinstance(obj, set):
+        if not obj:  # Empty set
+            return "`set`: empty"
+
+        result = [f"`set[{len(obj)}]`:"]
+        for i, item in enumerate(obj):
+            if i >= max_display_items and len(obj) > max_display_items + 2:
+                result.append(f"{indent}* ... ({len(obj) - 10} more items)")
+                break
+            item_md = object_to_markdown(item, max_depth, current_depth + 1, indent + "  ")
+            result.append(f"{indent}* {item_md}")
+        return "\n".join(result)
+
+    # Handle dataclasses
+    if is_dataclass(obj):
+        result = [f"`{type(obj).__name__}:`"]
+        for field in dataclass_fields(obj):
+            if field.name.startswith("_"):  # Skip private fields
+                continue
+            if not hasattr(obj, field.name):  # Skip not found
+                continue
+            value = getattr(obj, field.name)
+            value_md = object_to_markdown(value, max_depth, current_depth + 1, indent + "  ")
+            result.append(f"{indent}* **{field.name}**: {value_md}")
+        return "\n".join(result)
+
+    # Handle objects with __dict__
+    if hasattr(obj, "__dict__"):
+        attrs = filter_private_attributes(obj)
+        if not attrs:  # No public attributes
+            return f"`{type(obj).__name__}`: (no public attributes)"
+
+        result = [f"`{type(obj).__name__}`:"]
+        for attr, value in attrs.items():
+            if not is_method(obj, attr):  # Skip methods
+                value_md = object_to_markdown(value, max_depth, current_depth + 1, indent + "  ")
+                result.append(f"{indent}* **{attr}**: {value_md}")
+        return "\n".join(result)
+
+    # Handle other types
+    return f"`{get_object_type_str(obj)}`: {str(obj)}"
+
+
+def format_json_for_display(json_obj) -> str:
+    """Format a JSON object for display with proper indentation."""
+    return json.dumps(json_obj, indent=2, default=str)
+
+
+def display_node(node, display_format: str) -> None:
+    """
+    Display node content based on its type, with enhanced formatting.
+    """
+    # Special handling for specific types
+    if isinstance(node, pd.DataFrame):
+        st.dataframe(node)
+        return
+
+    # Toggle between markdown and JSON display
+    match display_format:
+        case "Markdown":
+            st.markdown(object_to_markdown(node))
+        case "JSON":
+            st.json(object_to_json_with_types(node))
+        case "Rich":
+            st.write(rich_print(node))
+        case _:
+            raise ValueError(f"Unsupported display format: {display_format}")
+
+
+def rich_print(node):
+    """Display a node using rich print."""
+    console = Console(color_system=None)
+    with console.capture() as capture:
+        console.print(node)
+    return f"```\n{capture.get()}\n```"

aixtools/log_view/export.py

@@ -0,0 +1,51 @@
+"""
+Functions for exporting nodes to various formats.
+"""
+
+import json
+
+
+def export_nodes_to_json(nodes: list) -> str:
+    """Export nodes to a JSON string for download."""
+    # Convert nodes to a serializable format
+    serializable_nodes = []
+
+    for node in nodes:
+        if hasattr(node, "__dict__"):
+            # For objects with attributes
+            node_dict = {
+                "type": type(node).__name__,
+                "attributes": {
+                    attr: str(value) if not isinstance(value, (dict, list, int, float, bool, type(None))) else value
+                    for attr, value in vars(node).items()
+                    if not attr.startswith("_")
+                },
+            }
+            serializable_nodes.append(node_dict)
+        elif isinstance(node, dict):
+            # For dictionaries
+            serializable_nodes.append(
+                {
+                    "type": "dict",
+                    "content": {
+                        str(k): str(v) if not isinstance(v, (dict, list, int, float, bool, type(None))) else v
+                        for k, v in node.items()
+                    },
+                }
+            )
+        elif isinstance(node, (list, tuple)):
+            # For lists and tuples
+            serializable_nodes.append(
+                {
+                    "type": "list" if isinstance(node, list) else "tuple",
+                    "content": [
+                        str(item) if not isinstance(item, (dict, list, int, float, bool, type(None))) else item
+                        for item in node
+                    ],
+                }
+            )
+        else:
+            # For primitive types
+            serializable_nodes.append({"type": type(node).__name__, "value": str(node)})
+
+    return json.dumps(serializable_nodes, indent=2)