agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl

agno/db/postgres/schemas.py

@@ -4,7 +4,7 @@ from typing import Any
 
 try:
     from sqlalchemy.dialects.postgresql import JSONB
-    from sqlalchemy.types import JSON, BigInteger, Boolean, Date, String
+    from sqlalchemy.types import BigInteger, Boolean, Date, Integer, String, Text
 except ImportError:
     raise ImportError("`sqlalchemy` not installed. Please install it using `pip install sqlalchemy`")
 
@@ -15,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": [
@@ -34,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},
@@ -62,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},
@@ -84,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},
@@ -102,27 +104,151 @@ METRICS_TABLE_SCHEMA = {
 CULTURAL_KNOWLEDGE_TABLE_SCHEMA = {
     "id": {"type": String, "primary_key": True, "nullable": False},
     "name": {"type": String, "nullable": False, "index": True},
-    "summary": {"type": String, "nullable": True},
+    "summary": {"type": Text, "nullable": True},
     "content": {"type": JSONB, "nullable": True},
     "metadata": {"type": JSONB, "nullable": True},
-    "input": {"type": String, "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},
+}
+
+TRACE_TABLE_SCHEMA = {
+    "trace_id": {"type": String, "primary_key": True, "nullable": False},
+    "name": {"type": String, "nullable": False},
+    "status": {"type": String, "nullable": False, "index": True},
+    "start_time": {"type": String, "nullable": False, "index": True},  # ISO 8601 datetime string
+    "end_time": {"type": String, "nullable": False},  # ISO 8601 datetime string
+    "duration_ms": {"type": BigInteger, "nullable": False},
+    "run_id": {"type": String, "nullable": True, "index": True},
+    "session_id": {"type": String, "nullable": True, "index": True},
+    "user_id": {"type": String, "nullable": True, "index": True},
+    "agent_id": {"type": String, "nullable": True, "index": True},
+    "team_id": {"type": String, "nullable": True, "index": True},
+    "workflow_id": {"type": String, "nullable": True, "index": True},
+    "created_at": {"type": String, "nullable": False, "index": True},  # ISO 8601 datetime string
+}
+
+
+def _get_span_table_schema(traces_table_name: str = "agno_traces", db_schema: str = "agno") -> dict[str, Any]:
+    """Get the span table schema with the correct foreign key reference.
+
+    Args:
+        traces_table_name: The name of the traces table to reference in the foreign key.
+        db_schema: The database schema name.
+
+    Returns:
+        The span table schema dictionary.
+    """
+    return {
+        "span_id": {"type": String, "primary_key": True, "nullable": False},
+        "trace_id": {
+            "type": String,
+            "nullable": False,
+            "index": True,
+            "foreign_key": f"{db_schema}.{traces_table_name}.trace_id",
+        },
+        "parent_span_id": {"type": String, "nullable": True, "index": True},
+        "name": {"type": String, "nullable": False},
+        "span_kind": {"type": String, "nullable": False},
+        "status_code": {"type": String, "nullable": False},
+        "status_message": {"type": Text, "nullable": True},
+        "start_time": {"type": String, "nullable": False, "index": True},  # ISO 8601 datetime string
+        "end_time": {"type": String, "nullable": False},  # ISO 8601 datetime string
+        "duration_ms": {"type": BigInteger, "nullable": False},
+        "attributes": {"type": JSONB, "nullable": True},
+        "created_at": {"type": String, "nullable": False, "index": True},  # ISO 8601 datetime string
+    }
+
+
+COMPONENT_TABLE_SCHEMA = {
+    "component_id": {"type": String, "primary_key": True},
+    "component_type": {"type": String, "nullable": False, "index": True},  # agent|team|workflow
+    "name": {"type": String, "nullable": True, "index": True},
+    "description": {"type": Text, "nullable": True},
+    "current_version": {"type": Integer, "nullable": True, "index": True},
+    "metadata": {"type": JSONB, "nullable": True},
+    "created_at": {"type": BigInteger, "nullable": False, "index": True},
+    "updated_at": {"type": BigInteger, "nullable": True},
+    "deleted_at": {"type": BigInteger, "nullable": True},
+}
+
+COMPONENT_CONFIGS_TABLE_SCHEMA = {
+    "component_id": {"type": String, "primary_key": True, "foreign_key": "components.component_id"},
+    "version": {"type": Integer, "primary_key": True},
+    "label": {"type": String, "nullable": True},  # stable|v1.2.0|pre-refactor
+    "stage": {"type": String, "nullable": False, "default": "draft", "index": True},  # draft|published
+    "config": {"type": JSONB, "nullable": False},
+    "notes": {"type": Text, "nullable": True},
+    "created_at": {"type": BigInteger, "nullable": False, "index": True},
+    "updated_at": {"type": BigInteger, "nullable": True},
+    "deleted_at": {"type": BigInteger, "nullable": True},
+}
+
+COMPONENT_LINKS_TABLE_SCHEMA = {
+    "parent_component_id": {"type": String, "nullable": False},
+    "parent_version": {"type": Integer, "nullable": False},
+    "link_kind": {"type": String, "nullable": False, "index": True},
+    "link_key": {"type": String, "nullable": False},
+    "child_component_id": {"type": String, "nullable": False, "foreign_key": "components.component_id"},
+    "child_version": {"type": Integer, "nullable": True},
+    "position": {"type": Integer, "nullable": False},
+    "meta": {"type": JSONB, "nullable": True},
+    "created_at": {"type": BigInteger, "nullable": True, "index": True},
+    "updated_at": {"type": BigInteger, "nullable": True},
+    "__primary_key__": ["parent_component_id", "parent_version", "link_kind", "link_key"],
+    "__foreign_keys__": [
+        {
+            "columns": ["parent_component_id", "parent_version"],
+            "ref_table": "component_configs",
+            "ref_columns": ["component_id", "version"],
+        }
+    ],
+}
+LEARNINGS_TABLE_SCHEMA = {
+    "learning_id": {"type": String, "primary_key": True, "nullable": False},
+    "learning_type": {"type": String, "nullable": False, "index": True},
+    "namespace": {"type": String, "nullable": True, "index": True},
+    "user_id": {"type": String, "nullable": True, "index": True},
+    "agent_id": {"type": String, "nullable": True, "index": True},
+    "team_id": {"type": String, "nullable": True, "index": True},
+    "workflow_id": {"type": String, "nullable": True, "index": True},
+    "session_id": {"type": String, "nullable": True, "index": True},
+    "entity_id": {"type": String, "nullable": True, "index": True},
+    "entity_type": {"type": String, "nullable": True, "index": True},
+    "content": {"type": JSONB, "nullable": False},
+    "metadata": {"type": JSONB, "nullable": True},
+    "created_at": {"type": BigInteger, "nullable": False, "index": True},
+    "updated_at": {"type": BigInteger, "nullable": True},
+}
+
 
-def get_table_schema_definition(table_type: str) -> dict[str, Any]:
+def get_table_schema_definition(
+    table_type: str, traces_table_name: str = "agno_traces", db_schema: str = "agno"
+) -> dict[str, Any]:
     """
     Get the expected schema definition for the given table.
 
     Args:
         table_type (str): The type of table to get the schema for.
+        traces_table_name (str): The name of the traces table (used for spans foreign key).
+        db_schema (str): The database schema name (used for spans foreign key).
 
     Returns:
         Dict[str, Any]: Dictionary containing column definitions for the table
     """
+    # Handle spans table specially to resolve the foreign key reference
+    if table_type == "spans":
+        return _get_span_table_schema(traces_table_name, db_schema)
+
     schemas = {
         "sessions": SESSION_TABLE_SCHEMA,
         "evals": EVAL_TABLE_SCHEMA,
@@ -130,6 +256,12 @@ def get_table_schema_definition(table_type: str) -> dict[str, Any]:
         "memories": MEMORY_TABLE_SCHEMA,
         "knowledge": KNOWLEDGE_TABLE_SCHEMA,
         "culture": CULTURAL_KNOWLEDGE_TABLE_SCHEMA,
+        "versions": VERSIONS_TABLE_SCHEMA,
+        "traces": TRACE_TABLE_SCHEMA,
+        "components": COMPONENT_TABLE_SCHEMA,
+        "component_configs": COMPONENT_CONFIGS_TABLE_SCHEMA,
+        "component_links": COMPONENT_LINKS_TABLE_SCHEMA,
+        "learnings": LEARNINGS_TABLE_SCHEMA,
     }
 
     schema = schemas.get(table_type, {})

agno/db/postgres/utils.py

@@ -13,8 +13,9 @@ from agno.db.schemas.culture import CulturalKnowledge
 from agno.utils.log import log_debug, log_error, log_warning
 
 try:
-    from sqlalchemy import Table
+    from sqlalchemy import Table, func
     from sqlalchemy.dialects import postgresql
+    from sqlalchemy.exc import NoSuchTableError
     from sqlalchemy.inspection import inspect
     from sqlalchemy.orm import Session
     from sqlalchemy.sql.expression import text
@@ -34,6 +35,11 @@ def apply_sorting(stmt, table: Table, sort_by: Optional[str] = None, sort_order:
 
     Returns:
         The modified statement with sorting applied
+
+    Note:
+        For 'updated_at' sorting, uses COALESCE(updated_at, created_at) to fall back
+        to created_at when updated_at is NULL. This ensures pre-2.0 records (which may
+        have NULL updated_at) are sorted correctly by their creation time.
     """
     if sort_by is None:
         return stmt
@@ -42,8 +48,13 @@ def apply_sorting(stmt, table: Table, sort_by: Optional[str] = None, sort_order:
         log_debug(f"Invalid sort field: '{sort_by}'. Will not apply any sorting.")
         return stmt
 
-    # Apply the given sorting
-    sort_column = getattr(table.c, sort_by)
+    # For updated_at, use COALESCE to fall back to created_at if updated_at is NULL
+    # This handles pre-2.0 records that may have NULL updated_at values
+    if sort_by == "updated_at" and hasattr(table.c, "created_at"):
+        sort_column = func.coalesce(table.c.updated_at, table.c.created_at)
+    else:
+        sort_column = getattr(table.c, sort_by)
+
     if sort_order and sort_order == "asc":
         return stmt.order_by(sort_column.asc())
     else:
@@ -173,6 +184,9 @@ async def ais_valid_table(db_engine: AsyncEngine, table_name: str, table_type: s
             return False
 
         return True
+    except NoSuchTableError:
+        log_error(f"Table {db_schema}.{table_name} does not exist")
+        return False
     except Exception as e:
         log_error(f"Error validating table schema for {db_schema}.{table_name}: {e}")
         return False
@@ -298,16 +312,17 @@ def calculate_date_metrics(date_to_process: date, sessions_data: dict) -> dict:
         for session in sessions:
             if session.get("user_id"):
                 all_user_ids.add(session["user_id"])
-            metrics[runs_count_key] += len(session.get("runs", []))
-            if runs := session.get("runs", []):
+            runs = session.get("runs", []) or []
+            metrics[runs_count_key] += len(runs)
+            if runs:
                 for run in runs:
                     if model_id := run.get("model"):
                         model_provider = run.get("model_provider", "")
                         model_counts[f"{model_id}:{model_provider}"] = (
                             model_counts.get(f"{model_id}:{model_provider}", 0) + 1
                         )
-
-            session_metrics = session.get("session_data", {}).get("session_metrics", {})
+            session_data = session.get("session_data", {}) or {}
+            session_metrics = session_data.get("session_metrics", {}) or {}
             for field in token_metrics:
                 token_metrics[field] += session_metrics.get(field, 0)