aixtools 0.0.0__py3-none-any.whl

aixtools/agents/agent_batch.py

@@ -0,0 +1,71 @@
+"""
+Batch processing functionality for running multiple agent queries in parallel.
+"""
+
+import asyncio
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+from aixtools.agents.agent import get_agent, run_agent
+
+
+class AgentQueryParams(BaseModel):
+    """Parameters for configuring agent queries in batch processing."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    id: str = ""  # Unique identifier for the query
+    prompt: str | list[str]
+    agent: Any = None
+    model: Any = None
+    debug: bool = False
+    output_type: Any = str
+    tools: list | None = []
+
+    async def run(self):
+        """Query the LLM"""
+        agent = self.agent
+        if agent is None:
+            agent = get_agent(
+                system_prompt=self.prompt, model=self.model, tools=self.tools, output_type=self.output_type
+            )
+        return await run_agent(agent=agent, prompt=self.prompt, debug=self.debug)
+
+
+async def run_agent_batch(query_parameters: list[AgentQueryParams], batch_size=10):
+    """
+    Run multiple queries simultanously in batches of at most batch_size
+    and yield the results as they come in.
+
+        Usage example:
+        query_parameters = [
+            AgentQueryParams(prompt="What is the meaning of life")
+            AgentQueryParams(prompt="Who is the prime minister of Canada")
+        ]
+
+        async for result in agent_batch(query_parameters):
+            print(result)
+    """
+    tasks = []
+    batch_num, total = 1, len(query_parameters)
+    for i, qp in enumerate(query_parameters):
+        tasks.append(qp.run())
+        if len(tasks) >= batch_size:
+            # Run a batch of tasks
+            print(f"Running batch {batch_num}, {i + 1} / {total}")
+            tasks_results = await asyncio.gather(
+                *tasks
+            )  # Returns a list of results, each one is a tuple (result, nodes)
+            # Yield the results
+            for r, _ in tasks_results:
+                yield r
+            tasks = []
+            batch_num += 1
+    # Run the last batch of tasks
+    if tasks:
+        print(f"Running final batch {batch_num}")
+        tasks_results = await asyncio.gather(*tasks)
+        for r, _ in tasks_results:
+            yield r
+    print("Done")

aixtools/agents/prompt.py

@@ -0,0 +1,97 @@
+"""Prompt building utilities for Pydantic AI agent, including file handling and context management."""
+
+import mimetypes
+from pathlib import Path, PurePosixPath
+
+from pydantic_ai import BinaryContent
+
+from aixtools.context import SessionIdTuple
+from aixtools.server import container_to_host_path
+from aixtools.utils.files import is_text_content
+
+CLAUDE_MAX_FILE_SIZE_IN_CONTEXT = 4 * 1024 * 1024  # Claude limit 4.5 MB for PDF files
+CLAUDE_IMAGE_MAX_FILE_SIZE_IN_CONTEXT = (
+    5 * 1024 * 1024
+)  # Claude limit 5 MB for images, to avoid large image files in context
+
+
+def should_be_included_into_context(
+    file_content: BinaryContent | str | None,
+    file_size: int,
+    *,
+    max_img_size_bytes: int = CLAUDE_IMAGE_MAX_FILE_SIZE_IN_CONTEXT,
+    max_file_size_bytes: int = CLAUDE_MAX_FILE_SIZE_IN_CONTEXT,
+) -> bool:
+    """Decide whether a file content should be included into the model context based on its type and size."""
+    if not isinstance(file_content, BinaryContent):
+        return False
+
+    if file_content.media_type.startswith("text/"):
+        return False
+
+    # Exclude archive files as they're not supported by OpenAI models
+    archive_types = {
+        "application/zip",
+        "application/x-tar",
+        "application/gzip",
+        "application/x-gzip",
+        "application/x-rar-compressed",
+        "application/x-7z-compressed",
+    }
+    if file_content.media_type in archive_types:
+        return False
+
+    if file_content.is_image and file_size < max_img_size_bytes:
+        return True
+
+    return file_size < max_file_size_bytes
+
+
+def file_to_binary_content(file_path: str | Path, mime_type: str = "") -> str | BinaryContent:
+    """
+    Read a file and return its content as either a UTF-8 string (for text files)
+    or BinaryContent (for binary files).
+    """
+    with open(file_path, "rb") as f:
+        data = f.read()
+
+    if not mime_type:
+        mime_type, _ = mimetypes.guess_type(file_path)
+        mime_type = mime_type or "application/octet-stream"
+
+    if is_text_content(data, mime_type):
+        return data.decode("utf-8")
+
+    return BinaryContent(data=data, media_type=mime_type)
+
+
+def build_user_input(
+    session_tuple: SessionIdTuple,
+    user_text: str,
+    file_paths: list[Path],
+) -> str | list[str | BinaryContent]:
+    """Build user input for the Pydantic AI agent, including file attachments if provided."""
+    if not file_paths:
+        return user_text
+
+    attachment_info_lines = []
+    binary_attachments = []
+
+    for workspace_path in file_paths:
+        host_path = container_to_host_path(PurePosixPath(workspace_path), ctx=session_tuple)
+        file_size = host_path.stat().st_size
+        mime_type, _ = mimetypes.guess_type(host_path)
+        mime_type = mime_type or "application/octet-stream"
+
+        attachment_info = f"* {workspace_path.name} (file_size={file_size} bytes) (path in workspace: {workspace_path})"
+        binary_content = file_to_binary_content(host_path, mime_type)
+
+        if should_be_included_into_context(binary_content, file_size):
+            binary_attachments.append(binary_content)
+            attachment_info += f" -- provided to model context at index {len(binary_attachments) - 1}"
+
+        attachment_info_lines.append(attachment_info)
+
+    full_prompt = user_text + "\nAttachments:\n" + "\n".join(attachment_info_lines)
+
+    return [full_prompt] + binary_attachments

aixtools/app.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+
+"""
+Simple Chainlit app example
+"""
+
+import traceback
+
+import chainlit as cl
+from pydantic_graph import End
+
+from aixtools.agents.agent import get_agent
+from aixtools.logging.logging_config import get_logger
+from aixtools.utils.chainlit import cl_agent_show
+
+logger = get_logger(__name__)
+
+HISTORY = "history"
+
+SYSTEM_PROMPT = """
+You are a helpful assistant.
+"""
+
+
+@cl.step
+async def greet_tool(msg: str) -> str:
+    """A simple greeting tool"""
+    return f"Hello! You said: {msg}"
+
+
+async def parse_user_message(message):
+    """Parse user message and check if it is a command"""
+    # When we type something that starts with ':', we are using a "command" (i.e. it does not go to the agent)
+    command = str(message.content).strip().lower()
+    if command.startswith(":"):
+        logger.debug("Received command: %s", command)
+        match command:
+            case ":clear":
+                # Clear the history
+                cl.user_session.set(HISTORY, [])
+                return None
+            case ":help":
+                # Show help
+                help_message = """
+                Available commands:
+                - :clear: Clear the chat history
+                - :help: Show this help message
+                """
+                await cl.Message(content=help_message).send()
+                return None
+            case ":history":
+                # Show history
+                history = cl.user_session.get(HISTORY)
+                if history:
+                    history_message = "\n".join(history)
+                    await cl.Message(content=f"Chat history:\n{history_message}").send()
+                else:
+                    await cl.Message(content="No history available.").send()
+                return None
+            case _:
+                # Unknown command
+                await cl.Message(content=f"Unknown command: {command}").send()
+                return None
+    else:
+        user_message = message.content
+        logger.debug("User message: %s", user_message)
+    return user_message
+
+
+async def run_agent(messages):
+    """Run the agent with the given messages"""
+    agent = get_agent(system_prompt=SYSTEM_PROMPT, tools=[greet_tool])
+    ret = ""
+    msg = cl.Message(content="")
+    await msg.send()
+    try:
+        ret = await cl_agent_show.show_run(agent=agent, prompt=messages, msg=msg, debug=False)
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        msg.elements.append(cl.Text(name="Error", content=f"Error: {e}", type="error"))  # pylint: disable=unexpected-keyword-arg
+        logger.error("Error: %s", e)
+        # Log the full stack trace for debugging
+        stack_trace = traceback.format_exc()
+        logger.error("Stack tarace:\n%s", stack_trace)
+        logger.error("Stack trace:\n%s", stack_trace)
+        msg.elements.append(cl.Text(name="Stack Trace", content=stack_trace, language="python"))
+        ret = f"Internal server error: {e}"
+    await msg.send()
+    return ret
+
+
+def update_history(history, user_message=None, run_return=None):
+    """Update history with user message and model run output"""
+    assert user_message is not None or run_return is not None, "Either user message or run return must be provided"
+    if user_message is not None:
+        logger.debug("Updating history: Got user message type %s: %s", type(user_message), user_message)
+        assert isinstance(user_message, str)
+        history.append(user_message)
+    if run_return is not None:
+        logger.debug("Updating history: Got agent output type %s: %s", type(run_return), run_return)
+        latest_item = ""
+        if isinstance(run_return, list):
+            # If it is a list of 'node' items, the last element is the 'end_message' with the final result
+            end_message: End = run_return[-1]
+            final_result = end_message.data
+            latest_item = str(final_result.data)
+        else:
+            latest_item = str(run_return)
+        # Update history and store it
+        logger.debug("Updating history: Adding to history type %s: %s", type(latest_item), latest_item)
+        history.append(latest_item)
+    return history
+
+
+@cl.set_starters
+async def set_starters():
+    """Set the starters"""
+    return [
+        cl.Starter(label="Message", message="Hello world!"),
+    ]
+
+
+@cl.on_chat_start
+async def on_chat_start():
+    """Initialize chat session by resetting history when a new chat starts."""
+    # Reset history
+    logger.debug("On chat start")
+    cl.user_session.set(HISTORY, [])
+
+
+@cl.on_message
+async def on_message(message: cl.Message):
+    """Process incoming chat messages and generate responses using the agent."""
+    history = cl.user_session.get(HISTORY)  # Get user message and history
+    user_message = await parse_user_message(message)  # Parse user message
+    # Check if user message is None (e.g. if it is a command)
+    if user_message is None:
+        return
+    messages = update_history(history, user_message=user_message)  # Update history with user message
+    # Run the agent
+    run_return = await run_agent(messages)
+    # Update history and store it
+    history = update_history(history, run_return=run_return)
+    cl.user_session.set(HISTORY, messages)

aixtools/chainlit.md

@@ -0,0 +1,14 @@
+# Welcome to Chainlit! 🚀🤖
+
+Hi there, Developer! 👋 We're excited to have you on board. Chainlit is a powerful tool designed to help you prototype, debug and share applications built on top of LLMs.
+
+## Useful Links 🔗
+
+- **Documentation:** Get started with our comprehensive [Chainlit Documentation](https://docs.chainlit.io) 📚
+- **Discord Community:** Join our friendly [Chainlit Discord](https://discord.gg/k73SQ3FyUh) to ask questions, share your projects, and connect with other developers! 💬
+
+We can't wait to see what you create with Chainlit! Happy coding! 💻😊
+
+## Welcome screen
+
+To modify the welcome screen, edit the `chainlit.md` file at the root of your project. If you do not want a welcome screen, just leave this file empty.

aixtools/compliance/__init__.py

@@ -0,0 +1,9 @@
+"""
+Compliance module for aixtools.
+
+This module provides utilities for managing compliance-related data and operations.
+"""
+
+from .private_data import PrivateData
+
+__all__ = ["PrivateData"]

aixtools/compliance/private_data.py

@@ -0,0 +1,138 @@
+"""Private data management module for aixtools compliance."""
+
+import json
+from pathlib import Path
+
+from fastmcp import Context
+
+from aixtools.server.path import get_workspace_path
+
+PRIVATE_DATA_FILE = ".private_data"
+
+
+class PrivateData:
+    """
+    Class to manage private data file in the workspace.
+
+    The information is stored in a JSON file named `.private_data` within the workspace directory.
+    If the file does not exist, it indicates that there is no private data.
+
+    IMPORTANT: All modifications save the data to the file immediately.
+
+    FIXME: We should add some level of mutex/locking to prevent concurrent writes.
+    """
+
+    def __init__(self, ctx: Context | None = None):
+        self.ctx: Context | None = ctx
+        self._has_private_data: bool = False  # Flag indicating if private data exists
+        self._private_datasets: list[str] = []  # List of private datasets
+        self._idap_datasets: list[str] = []  # List of dataset with IDAP
+        self.load()
+
+    def add_private_dataset(self, dataset_name: str) -> None:
+        """
+        Add a private dataset to the list.
+        Save the state after modification.
+        """
+        if dataset_name not in self._private_datasets:
+            self._private_datasets.append(dataset_name)
+            self._has_private_data = True
+            self.save()
+
+    def add_idap_dataset(self, dataset_name: str) -> None:
+        """
+        Add a dataset with IDAP to the list.
+        This also adds it to the private datasets if not already present.
+        Save the state after modification.
+        """
+        if not self.has_idap_dataset(dataset_name):
+            self._idap_datasets.append(dataset_name)
+            self._has_private_data = True
+            # An IDAP dataset is also a private dataset
+            if not self.has_private_dataset(dataset_name):
+                self._private_datasets.append(dataset_name)
+            self.save()
+
+    def get_private_datasets(self) -> list[str]:
+        """Get the list of private datasets as a copy (to avoid modification)."""
+        return list(self._private_datasets)
+
+    def get_idap_datasets(self) -> list[str]:
+        """Get the list of datasets with IDAP as a copy (to avoid modification)."""
+        return list(self._idap_datasets)
+
+    def has_private_dataset(self, dataset_name: str) -> bool:
+        """Check if a specific private dataset exists."""
+        return dataset_name in self._private_datasets
+
+    def has_idap_dataset(self, dataset_name: str) -> bool:
+        """Check if a specific dataset with IDAP exists."""
+        return dataset_name in self._idap_datasets
+
+    @property
+    def has_private_data(self) -> bool:
+        """Check if private data exists."""
+        return self._has_private_data
+
+    @has_private_data.setter
+    def has_private_data(self, value: bool) -> None:
+        """
+        Set the flag indicating if private data exists.
+        Save the state after modification.
+        """
+        self._has_private_data = value
+        if not value:
+            self._private_datasets = []
+            self._idap_datasets = []
+        self.save()
+
+    def _get_private_data_path(self) -> Path:
+        """Get the path to the private data file in the workspace."""
+        return get_workspace_path(service_name=None, ctx=self.ctx) / PRIVATE_DATA_FILE
+
+    def _has_private_data_file(self) -> bool:
+        """Check if the private data file exists in the workspace."""
+        private_data_path = self._get_private_data_path()
+        return private_data_path.exists()
+
+    def save(self) -> None:
+        """Save content to the private data file in the workspace."""
+        private_data_path = self._get_private_data_path()
+        # No private data? Delete the file if it exists
+        if not self.has_private_data:
+            private_data_path.unlink(missing_ok=True)
+            return
+        # If there is private data, serialize this object as JSON
+        private_data_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(private_data_path, "w", encoding="utf-8") as f:
+            # Dump class as JSON, excluding the context
+            data_dict = self.__dict__.copy()
+            data_dict["ctx"] = None
+            json_data = json.dumps(data_dict, indent=4)
+            f.write(json_data)
+
+    def load(self) -> None:
+        """Load content from the private data file in the workspace."""
+        private_data_path = self._get_private_data_path()
+        if not private_data_path.exists():
+            # No private data file
+            self.has_private_data = False
+            self._private_datasets = []
+            self._idap_datasets = []
+            return
+        with open(private_data_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+            self.has_private_data = data.get("_has_private_data", False)
+            self._private_datasets = data.get("_private_datasets", [])
+            self._idap_datasets = data.get("_idap_datasets", [])
+
+    def __repr__(self) -> str:
+        return (
+            f"PrivateData(has_private_data={self.has_private_data}, "
+            f"private_datasets={self._private_datasets}, "
+            f"idap_datasets={self._idap_datasets}), "
+            f"file_path={self._get_private_data_path()})"
+        )
+
+    def __str__(self) -> str:
+        return self.__repr__()

aixtools/context.py

@@ -0,0 +1,17 @@
+"""
+This module defines global context variables for request-specific information
+that can be used for logging, tracing, and other purposes across applications
+that use aixtools.
+"""
+
+from contextvars import ContextVar
+
+# Define context variables with default values.
+# These can be populated by middleware or where they are initialized
+session_id_var: ContextVar[str | None] = ContextVar("session_id", default=None)
+user_id_var: ContextVar[str | None] = ContextVar("user_id", default=None)
+
+DEFAULT_USER_ID = "default_user"
+DEFAULT_SESSION_ID = "default_session"
+
+SessionIdTuple = tuple[str, str]

aixtools/db/__init__.py

@@ -0,0 +1,17 @@
+"""
+Database module for vector storage and retrieval.
+"""
+
+from aixtools.db.database import DatabaseError, SqliteDb
+from aixtools.db.vector_db import get_vdb_embedding, get_vector_db, vdb_add, vdb_get_by_id, vdb_has_id, vdb_query
+
+__all__ = [
+    "DatabaseError",
+    "SqliteDb",
+    "get_vdb_embedding",
+    "get_vector_db",
+    "vdb_add",
+    "vdb_get_by_id",
+    "vdb_has_id",
+    "vdb_query",
+]

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)