agno 2.0.1__py3-none-any.whl → 2.3.0__py3-none-any.whl

agno/db/postgres/schemas.py

@@ -3,7 +3,8 @@
 from typing import Any
 
 try:
-    from sqlalchemy.types import JSON, BigInteger, Boolean, Date, String
+    from sqlalchemy.dialects.postgresql import JSONB
+    from sqlalchemy.types import BigInteger, Boolean, Date, String, Text
 except ImportError:
     raise ImportError("`sqlalchemy` not installed. Please install it using `pip install sqlalchemy`")
 
@@ -14,13 +15,13 @@ SESSION_TABLE_SCHEMA = {
     "team_id": {"type": String, "nullable": True},
     "workflow_id": {"type": String, "nullable": True},
     "user_id": {"type": String, "nullable": True},
-    "session_data": {"type": JSON, "nullable": True},
-    "agent_data": {"type": JSON, "nullable": True},
-    "team_data": {"type": JSON, "nullable": True},
-    "workflow_data": {"type": JSON, "nullable": True},
-    "metadata": {"type": JSON, "nullable": True},
-    "runs": {"type": JSON, "nullable": True},
-    "summary": {"type": JSON, "nullable": True},
+    "session_data": {"type": JSONB, "nullable": True},
+    "agent_data": {"type": JSONB, "nullable": True},
+    "team_data": {"type": JSONB, "nullable": True},
+    "workflow_data": {"type": JSONB, "nullable": True},
+    "metadata": {"type": JSONB, "nullable": True},
+    "runs": {"type": JSONB, "nullable": True},
+    "summary": {"type": JSONB, "nullable": True},
     "created_at": {"type": BigInteger, "nullable": False, "index": True},
     "updated_at": {"type": BigInteger, "nullable": True},
     "_unique_constraints": [
@@ -33,20 +34,22 @@ SESSION_TABLE_SCHEMA = {
 
 MEMORY_TABLE_SCHEMA = {
     "memory_id": {"type": String, "primary_key": True, "nullable": False},
-    "memory": {"type": JSON, "nullable": False},
-    "input": {"type": String, "nullable": True},
+    "memory": {"type": JSONB, "nullable": False},
+    "feedback": {"type": Text, "nullable": True},
+    "input": {"type": Text, "nullable": True},
     "agent_id": {"type": String, "nullable": True},
     "team_id": {"type": String, "nullable": True},
     "user_id": {"type": String, "nullable": True, "index": True},
-    "topics": {"type": JSON, "nullable": True},
+    "topics": {"type": JSONB, "nullable": True},
+    "created_at": {"type": BigInteger, "nullable": False, "index": True},
     "updated_at": {"type": BigInteger, "nullable": True, "index": True},
 }
 
 EVAL_TABLE_SCHEMA = {
     "run_id": {"type": String, "primary_key": True, "nullable": False},
     "eval_type": {"type": String, "nullable": False},
-    "eval_data": {"type": JSON, "nullable": False},
-    "eval_input": {"type": JSON, "nullable": False},
+    "eval_data": {"type": JSONB, "nullable": False},
+    "eval_input": {"type": JSONB, "nullable": False},
     "name": {"type": String, "nullable": True},
     "agent_id": {"type": String, "nullable": True},
     "team_id": {"type": String, "nullable": True},
@@ -61,14 +64,14 @@ EVAL_TABLE_SCHEMA = {
 KNOWLEDGE_TABLE_SCHEMA = {
     "id": {"type": String, "primary_key": True, "nullable": False},
     "name": {"type": String, "nullable": False},
-    "description": {"type": String, "nullable": False},
-    "metadata": {"type": JSON, "nullable": True},
+    "description": {"type": Text, "nullable": False},
+    "metadata": {"type": JSONB, "nullable": True},
     "type": {"type": String, "nullable": True},
     "size": {"type": BigInteger, "nullable": True},
     "linked_to": {"type": String, "nullable": True},
     "access_count": {"type": BigInteger, "nullable": True},
     "status": {"type": String, "nullable": True},
-    "status_message": {"type": String, "nullable": True},
+    "status_message": {"type": Text, "nullable": True},
     "created_at": {"type": BigInteger, "nullable": True},
     "updated_at": {"type": BigInteger, "nullable": True},
     "external_id": {"type": String, "nullable": True},
@@ -83,8 +86,8 @@ METRICS_TABLE_SCHEMA = {
     "team_sessions_count": {"type": BigInteger, "nullable": False, "default": 0},
     "workflow_sessions_count": {"type": BigInteger, "nullable": False, "default": 0},
     "users_count": {"type": BigInteger, "nullable": False, "default": 0},
-    "token_metrics": {"type": JSON, "nullable": False, "default": {}},
-    "model_metrics": {"type": JSON, "nullable": False, "default": {}},
+    "token_metrics": {"type": JSONB, "nullable": False, "default": {}},
+    "model_metrics": {"type": JSONB, "nullable": False, "default": {}},
     "date": {"type": Date, "nullable": False, "index": True},
     "aggregation_period": {"type": String, "nullable": False},
     "created_at": {"type": BigInteger, "nullable": False},
@@ -98,6 +101,26 @@ METRICS_TABLE_SCHEMA = {
     ],
 }
 
+CULTURAL_KNOWLEDGE_TABLE_SCHEMA = {
+    "id": {"type": String, "primary_key": True, "nullable": False},
+    "name": {"type": String, "nullable": False, "index": True},
+    "summary": {"type": Text, "nullable": True},
+    "content": {"type": JSONB, "nullable": True},
+    "metadata": {"type": JSONB, "nullable": True},
+    "input": {"type": Text, "nullable": True},
+    "created_at": {"type": BigInteger, "nullable": True},
+    "updated_at": {"type": BigInteger, "nullable": True},
+    "agent_id": {"type": String, "nullable": True},
+    "team_id": {"type": String, "nullable": True},
+}
+
+VERSIONS_TABLE_SCHEMA = {
+    "table_name": {"type": String, "nullable": False, "primary_key": True},
+    "version": {"type": String, "nullable": False},
+    "created_at": {"type": String, "nullable": False, "index": True},
+    "updated_at": {"type": String, "nullable": True},
+}
+
 
 def get_table_schema_definition(table_type: str) -> dict[str, Any]:
     """
@@ -115,6 +138,8 @@ def get_table_schema_definition(table_type: str) -> dict[str, Any]:
         "metrics": METRICS_TABLE_SCHEMA,
         "memories": MEMORY_TABLE_SCHEMA,
         "knowledge": KNOWLEDGE_TABLE_SCHEMA,
+        "culture": CULTURAL_KNOWLEDGE_TABLE_SCHEMA,
+        "versions": VERSIONS_TABLE_SCHEMA,
     }
 
     schema = schemas.get(table_type, {})

agno/db/postgres/utils.py

@@ -6,8 +6,10 @@ from typing import Any, Dict, List, Optional
 from uuid import uuid4
 
 from sqlalchemy import Engine
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
 
 from agno.db.postgres.schemas import get_table_schema_definition
+from agno.db.schemas.culture import CulturalKnowledge
 from agno.utils.log import log_debug, log_error, log_warning
 
 try:
@@ -62,6 +64,20 @@ def create_schema(session: Session, db_schema: str) -> None:
         log_warning(f"Could not create schema {db_schema}: {e}")
 
 
+async def acreate_schema(session: AsyncSession, db_schema: str) -> None:
+    """Create the database schema if it doesn't exist.
+
+    Args:
+        session: The SQLAlchemy session to use
+        db_schema (str): The definition of the database schema to create
+    """
+    try:
+        log_debug(f"Creating schema if not exists: {db_schema}")
+        await session.execute(text(f"CREATE SCHEMA IF NOT EXISTS {db_schema};"))
+    except Exception as e:
+        log_warning(f"Could not create schema {db_schema}: {e}")
+
+
 def is_table_available(session: Session, table_name: str, db_schema: str) -> bool:
     """
     Check if a table with the given name exists in the given schema.
@@ -81,6 +97,24 @@ def is_table_available(session: Session, table_name: str, db_schema: str) -> boo
         return False
 
 
+async def ais_table_available(session: AsyncSession, table_name: str, db_schema: str) -> bool:
+    """
+    Check if a table with the given name exists in the given schema.
+
+    Returns:
+        bool: True if the table exists, False otherwise.
+    """
+    try:
+        exists_query = text(
+            "SELECT 1 FROM information_schema.tables WHERE table_schema = :schema AND table_name = :table"
+        )
+        exists = (await session.execute(exists_query, {"schema": db_schema, "table": table_name})).scalar() is not None
+        return exists
+    except Exception as e:
+        log_error(f"Error checking if table exists: {e}")
+        return False
+
+
 def is_valid_table(db_engine: Engine, table_name: str, table_type: str, db_schema: str) -> bool:
     """
     Check if the existing table has the expected column names.
@@ -113,6 +147,44 @@ def is_valid_table(db_engine: Engine, table_name: str, table_type: str, db_schem
         return False
 
 
+async def ais_valid_table(db_engine: AsyncEngine, table_name: str, table_type: str, db_schema: str) -> bool:
+    """
+    Check if the existing table has the expected column names.
+
+    Args:
+        table_name (str): Name of the table to validate
+        schema (str): Database schema name
+
+    Returns:
+        bool: True if table has all expected columns, False otherwise
+    """
+    try:
+        expected_table_schema = get_table_schema_definition(table_type)
+        expected_columns = {col_name for col_name in expected_table_schema.keys() if not col_name.startswith("_")}
+
+        # Get existing columns from the async engine
+        async with db_engine.connect() as conn:
+            existing_columns = await conn.run_sync(_get_table_columns, table_name, db_schema)
+
+        # Check if all expected columns exist
+        missing_columns = expected_columns - existing_columns
+        if missing_columns:
+            log_warning(f"Missing columns {missing_columns} in table {db_schema}.{table_name}")
+            return False
+
+        return True
+    except Exception as e:
+        log_error(f"Error validating table schema for {db_schema}.{table_name}: {e}")
+        return False
+
+
+def _get_table_columns(conn, table_name: str, db_schema: str) -> set[str]:
+    """Helper function to get table columns using sync inspector."""
+    inspector = inspect(conn)
+    columns_info = inspector.get_columns(table_name, schema=db_schema)
+    return {col["name"] for col in columns_info}
+
+
 # -- Metrics util methods --
 def bulk_upsert_metrics(session: Session, table: Table, metrics_records: list[dict]) -> list[dict]:
     """Bulk upsert metrics into the database.
@@ -147,6 +219,39 @@ def bulk_upsert_metrics(session: Session, table: Table, metrics_records: list[di
     return results  # type: ignore
 
 
+async def abulk_upsert_metrics(session: AsyncSession, table: Table, metrics_records: list[dict]) -> list[dict]:
+    """Bulk upsert metrics into the database.
+
+    Args:
+        table (Table): The table to upsert into.
+        metrics_records (list[dict]): The metrics records to upsert.
+
+    Returns:
+        list[dict]: The upserted metrics records.
+    """
+    if not metrics_records:
+        return []
+
+    results = []
+    stmt = postgresql.insert(table)
+
+    # Columns to update in case of conflict
+    update_columns = {
+        col.name: stmt.excluded[col.name]
+        for col in table.columns
+        if col.name not in ["id", "date", "created_at", "aggregation_period"]
+    }
+
+    stmt = stmt.on_conflict_do_update(index_elements=["date", "aggregation_period"], set_=update_columns).returning(  # type: ignore
+        table
+    )
+    result = await session.execute(stmt, metrics_records)
+    results = [row._mapping for row in result.fetchall()]
+    await session.commit()
+
+    return results  # type: ignore
+
+
 def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
     """Calculate metrics for the given single date.
 
@@ -187,7 +292,7 @@ def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
     all_user_ids = set()
 
     for session_type, sessions_count_key, runs_count_key in session_types:
-        sessions = sessions_data.get(session_type, [])
+        sessions = sessions_data.get(session_type, []) or []
         metrics[sessions_count_key] = len(sessions)
 
         for session in sessions:
@@ -208,7 +313,7 @@ def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
 
     model_metrics = []
     for model, count in model_counts.items():
-        model_id, model_provider = model.split(":")
+        model_id, model_provider = model.rsplit(":", 1)
         model_metrics.append({"model_id": model_id, "model_provider": model_provider, "count": count})
 
     metrics["users_count"] = len(all_user_ids)
@@ -278,3 +383,60 @@ def get_dates_to_calculate_metrics_for(starting_date: date) -> list[date]:
     if days_diff <= 0:
         return []
     return [starting_date + timedelta(days=x) for x in range(days_diff)]
+
+
+# -- Cultural Knowledge util methods --
+def serialize_cultural_knowledge(cultural_knowledge: CulturalKnowledge) -> Dict[str, Any]:
+    """Serialize a CulturalKnowledge object for database storage.
+
+    Converts the model's separate content, categories, and notes fields
+    into a single JSON dict for the database content column.
+
+    Args:
+        cultural_knowledge (CulturalKnowledge): The cultural knowledge object to serialize.
+
+    Returns:
+        Dict[str, Any]: A dictionary with the content field as JSON containing content, categories, and notes.
+    """
+    content_dict: Dict[str, Any] = {}
+    if cultural_knowledge.content is not None:
+        content_dict["content"] = cultural_knowledge.content
+    if cultural_knowledge.categories is not None:
+        content_dict["categories"] = cultural_knowledge.categories
+    if cultural_knowledge.notes is not None:
+        content_dict["notes"] = cultural_knowledge.notes
+
+    return content_dict if content_dict else {}
+
+
+def deserialize_cultural_knowledge(db_row: Dict[str, Any]) -> CulturalKnowledge:
+    """Deserialize a database row to a CulturalKnowledge object.
+
+    The database stores content as a JSON dict containing content, categories, and notes.
+    This method extracts those fields and converts them back to the model format.
+
+    Args:
+        db_row (Dict[str, Any]): The database row as a dictionary.
+
+    Returns:
+        CulturalKnowledge: The cultural knowledge object.
+    """
+    # Extract content, categories, and notes from the JSON content field
+    content_json = db_row.get("content", {}) or {}
+
+    return CulturalKnowledge.from_dict(
+        {
+            "id": db_row.get("id"),
+            "name": db_row.get("name"),
+            "summary": db_row.get("summary"),
+            "content": content_json.get("content"),
+            "categories": content_json.get("categories"),
+            "notes": content_json.get("notes"),
+            "metadata": db_row.get("metadata"),
+            "input": db_row.get("input"),
+            "created_at": db_row.get("created_at"),
+            "updated_at": db_row.get("updated_at"),
+            "agent_id": db_row.get("agent_id"),
+            "team_id": db_row.get("team_id"),
+        }
+    )