aixtools 0.1.0__py3-none-any.whl

aixtools/db/database.py

@@ -0,0 +1,110 @@
+"""
+Database Interface for Clinical Trials Information.
+
+This module provides a database interface for querying clinical trials data
+from the SQLite database.
+"""
+
+import sqlite3
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+from aixtools.logging.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class DatabaseError(Exception):
+    """Exception raised for database-related errors."""
+
+
+class SqliteDb:
+    """
+    Database interface.
+    """
+
+    def __init__(self, db_path: str | Path):
+        """Initialize the database interface"""
+        self.db_path = Path(db_path)
+        if not self.db_path.exists():
+            raise FileNotFoundError(f"Database file not found: {self.db_path}")
+        # Test connection
+        with self.connection() as conn:
+            logger.info("Connected to database: %s, connection: %s", self.db_path, conn)
+
+    @contextmanager
+    def connection(self):
+        """
+        Context manager for database connections.
+
+        Yields:
+            sqlite3.Connection: An active database connection
+        """
+        conn = None
+        try:
+            conn = sqlite3.connect(self.db_path)
+            # Enable dictionary row factory
+            conn.row_factory = sqlite3.Row
+            yield conn
+        except sqlite3.Error as e:
+            raise DatabaseError(f"Database error: {e}") from e
+        finally:
+            if conn:
+                conn.close()
+
+    def query(self, query: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+        """
+        Execute a SQL query and return the results as a list of dictionaries.
+
+        Args:
+            query: SQL query to execute
+            params: Parameters for the query
+
+        Returns:
+            List of dictionaries representing the query results
+        """
+        with self.connection() as conn:
+            cursor = conn.cursor()
+            if params:
+                cursor.execute(query, params)
+            else:
+                cursor.execute(query)
+
+            results = cursor.fetchall()
+            # Convert sqlite3.Row objects to dictionaries
+            return [dict(row) for row in results]
+
+    def query_df(self, query: str, params: dict[str, Any] | None = None) -> pd.DataFrame:
+        """
+        Execute a SQL query and return the results as a pandas DataFrame.
+
+        Args:
+            query: SQL query to execute.
+            params: Parameters to substitute in the query.
+
+        Returns:
+            A pandas DataFrame containing the query results.
+        """
+        with self.connection() as conn:
+            if params:
+                df = pd.read_sql_query(query, conn, params=params)
+            else:
+                df = pd.read_sql_query(query, conn)
+            return df
+
+    def validate(self, query) -> str | None:
+        """
+        Validate the SQL query by executing an EXPLAIN QUERY PLAN statement.
+        Returns the error string if there is an issue, otherwise returns None
+        """
+        with self.connection() as conn:
+            try:
+                cursor = conn.cursor()
+                cursor.execute(f"EXPLAIN QUERY PLAN\n{query}")
+                cursor.fetchall()
+                return None
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                return str(e)

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/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)