agno 2.3.22__py3-none-any.whl → 2.3.24__py3-none-any.whl

agno/agent/agent.py

@@ -61,6 +61,9 @@ from agno.run.agent import (
     RunOutput,
     RunOutputEvent,
 )
+from agno.run.cancel import (
+    acancel_run as acancel_run_global,
+)
 from agno.run.cancel import (
     acleanup_run,
     araise_if_cancelled,
@@ -2030,6 +2033,7 @@ class Agent:
                         run_response=run_response,
                         run_context=run_context,
                         session=agent_session,
+                        async_mode=True,
                     )
 
                     # 6. Prepare run messages
@@ -2377,6 +2381,7 @@ class Agent:
                         run_response=run_response,
                         run_context=run_context,
                         session=agent_session,
+                        async_mode=True,
                     )
 
                     # 6. Prepare run messages
@@ -3297,6 +3302,9 @@ class Agent:
                         tools=tools,
                         tool_choice=self.tool_choice,
                         tool_call_limit=self.tool_call_limit,
+                        run_response=run_response,
+                        send_media_to_model=self.send_media_to_model,
+                        compression_manager=self.compression_manager if self.compress_tool_results else None,
                     )
 
                     # Check for cancellation after model processing
@@ -3987,6 +3995,7 @@ class Agent:
                         run_response=run_response,
                         run_context=run_context,
                         session=agent_session,
+                        async_mode=True,
                     )
 
                     # 6. Prepare run messages
@@ -4009,6 +4018,9 @@ class Agent:
                         tools=_tools,
                         tool_choice=self.tool_choice,
                         tool_call_limit=self.tool_call_limit,
+                        run_response=run_response,
+                        send_media_to_model=self.send_media_to_model,
+                        compression_manager=self.compression_manager if self.compress_tool_results else None,
                     )
                     # Check for cancellation after model call
                     await araise_if_cancelled(run_response.run_id)  # type: ignore
@@ -4293,6 +4305,7 @@ class Agent:
                         run_response=run_response,
                         run_context=run_context,
                         session=agent_session,
+                        async_mode=True,
                     )
 
                     # 6. Prepare run messages
@@ -6514,6 +6527,7 @@ class Agent:
         run_response: RunOutput,
         run_context: RunContext,
         session: AgentSession,
+        async_mode: bool = False,
     ) -> List[Union[Function, dict]]:
         _function_names = []
         _functions: List[Union[Function, dict]] = []
@@ -6544,7 +6558,8 @@ class Agent:
 
                 elif isinstance(tool, Toolkit):
                     # For each function in the toolkit and process entrypoint
-                    for name, _func in tool.functions.items():
+                    toolkit_functions = tool.get_async_functions() if async_mode else tool.get_functions()
+                    for name, _func in toolkit_functions.items():
                         if name in _function_names:
                             continue
                         _function_names.append(name)
@@ -6774,6 +6789,18 @@ class Agent:
         """
         return cancel_run_global(run_id)
 
+    @staticmethod
+    async def acancel_run(run_id: str) -> bool:
+        """Cancel a running agent execution (async version).
+
+        Args:
+            run_id (str): The run_id to cancel.
+
+        Returns:
+            bool: True if the run was found and marked for cancellation, False otherwise.
+        """
+        return await acancel_run_global(run_id)
+
     # -*- Session Database Functions
     def _read_session(
         self, session_id: str, session_type: SessionType = SessionType.AGENT

agno/agent/remote.py

@@ -488,7 +488,7 @@ class RemoteAgent(BaseRemote):
         else:
             raise ValueError("No client available")
 
-    async def cancel_run(self, run_id: str, auth_token: Optional[str] = None) -> bool:
+    async def acancel_run(self, run_id: str, auth_token: Optional[str] = None) -> bool:
         """Cancel a running agent execution.
 
         Args:

agno/db/mongo/mongo.py

@@ -1,5 +1,6 @@
 import time
 from datetime import date, datetime, timedelta, timezone
+from importlib import metadata
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 from uuid import uuid4
 
@@ -31,10 +32,13 @@ try:
     from pymongo import MongoClient, ReturnDocument
     from pymongo.collection import Collection
     from pymongo.database import Database
+    from pymongo.driver_info import DriverInfo
     from pymongo.errors import OperationFailure
 except ImportError:
     raise ImportError("`pymongo` not installed. Please install it using `pip install pymongo`")
 
+DRIVER_METADATA = DriverInfo(name="Agno", version=metadata.version("agno"))
+
 
 class MongoDb(BaseDb):
     def __init__(
@@ -92,10 +96,14 @@ class MongoDb(BaseDb):
 
         _client: Optional[MongoClient] = db_client
         if _client is None and db_url is not None:
-            _client = MongoClient(db_url)
+            _client = MongoClient(db_url, driver=DRIVER_METADATA)
         if _client is None:
             raise ValueError("One of db_url or db_client must be provided")
 
+        # append_metadata was added in PyMongo 4.14.0, but is a valid database name on earlier versions
+        if callable(_client.append_metadata):
+            _client.append_metadata(DRIVER_METADATA)
+
         self.db_url: Optional[str] = db_url
         self.db_client: MongoClient = _client
 

agno/db/mysql/async_mysql.py

@@ -158,7 +158,10 @@ class AsyncMySQLDb(AsyncBaseDb):
             Table: SQLAlchemy Table object
         """
         try:
-            table_schema = get_table_schema_definition(table_type).copy()
+            # Pass traces_table_name and db_schema for spans table foreign key resolution
+            table_schema = get_table_schema_definition(
+                table_type, traces_table_name=self.trace_table_name, db_schema=self.db_schema
+            ).copy()
 
             log_debug(f"Creating table {self.db_schema}.{table_name} with schema: {table_schema}")
 
@@ -183,12 +186,7 @@ class AsyncMySQLDb(AsyncBaseDb):
 
                 # Handle foreign key constraint
                 if "foreign_key" in col_config:
-                    fk_ref = col_config["foreign_key"]
-                    # For spans table, dynamically replace the traces table reference
-                    # with the actual trace table name configured for this db instance
-                    if table_type == "spans" and "trace_id" in fk_ref:
-                        fk_ref = f"{self.db_schema}.{self.trace_table_name}.trace_id"
-                    column_args.append(ForeignKey(fk_ref))
+                    column_args.append(ForeignKey(col_config["foreign_key"]))
 
                 columns.append(Column(*column_args, **column_kwargs))  # type: ignore
 

agno/db/mysql/mysql.py

@@ -155,7 +155,10 @@ class MySQLDb(BaseDb):
             Table: SQLAlchemy Table object
         """
         try:
-            table_schema = get_table_schema_definition(table_type).copy()
+            # Pass traces_table_name and db_schema for spans table foreign key resolution
+            table_schema = get_table_schema_definition(
+                table_type, traces_table_name=self.trace_table_name, db_schema=self.db_schema
+            ).copy()
 
             columns: List[Column] = []
             indexes: List[str] = []
@@ -178,12 +181,7 @@ class MySQLDb(BaseDb):
 
                 # Handle foreign key constraint
                 if "foreign_key" in col_config:
-                    fk_ref = col_config["foreign_key"]
-                    # For spans table, dynamically replace the traces table reference
-                    # with the actual trace table name configured for this db instance
-                    if table_type == "spans" and "trace_id" in fk_ref:
-                        fk_ref = f"{self.db_schema}.{self.trace_table_name}.trace_id"
-                    column_args.append(ForeignKey(fk_ref))
+                    column_args.append(ForeignKey(col_config["foreign_key"]))
 
                 columns.append(Column(*column_args, **column_kwargs))  # type: ignore
 

agno/db/mysql/schemas.py

@@ -136,37 +136,56 @@ TRACE_TABLE_SCHEMA = {
     "created_at": {"type": lambda: String(128), "nullable": False, "index": True},  # ISO 8601 datetime string
 }
 
-SPAN_TABLE_SCHEMA = {
-    "span_id": {"type": lambda: String(128), "primary_key": True, "nullable": False},
-    "trace_id": {
-        "type": lambda: String(128),
-        "nullable": False,
-        "index": True,
-        "foreign_key": "agno_traces.trace_id",  # Foreign key to traces table
-    },
-    "parent_span_id": {"type": lambda: String(128), "nullable": True, "index": True},
-    "name": {"type": lambda: String(255), "nullable": False},
-    "span_kind": {"type": lambda: String(50), "nullable": False},
-    "status_code": {"type": lambda: String(50), "nullable": False},
-    "status_message": {"type": Text, "nullable": True},
-    "start_time": {"type": lambda: String(128), "nullable": False, "index": True},  # ISO 8601 datetime string
-    "end_time": {"type": lambda: String(128), "nullable": False},  # ISO 8601 datetime string
-    "duration_ms": {"type": BigInteger, "nullable": False},
-    "attributes": {"type": JSON, "nullable": True},
-    "created_at": {"type": lambda: String(128), "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": lambda: String(128), "primary_key": True, "nullable": False},
+        "trace_id": {
+            "type": lambda: String(128),
+            "nullable": False,
+            "index": True,
+            "foreign_key": f"{db_schema}.{traces_table_name}.trace_id",
+        },
+        "parent_span_id": {"type": lambda: String(128), "nullable": True, "index": True},
+        "name": {"type": lambda: String(255), "nullable": False},
+        "span_kind": {"type": lambda: String(50), "nullable": False},
+        "status_code": {"type": lambda: String(50), "nullable": False},
+        "status_message": {"type": Text, "nullable": True},
+        "start_time": {"type": lambda: String(128), "nullable": False, "index": True},  # ISO 8601 datetime string
+        "end_time": {"type": lambda: String(128), "nullable": False},  # ISO 8601 datetime string
+        "duration_ms": {"type": BigInteger, "nullable": False},
+        "attributes": {"type": JSON, "nullable": True},
+        "created_at": {"type": lambda: String(128), "nullable": False, "index": True},  # ISO 8601 datetime string
+    }
 
-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,
@@ -176,7 +195,6 @@ def get_table_schema_definition(table_type: str) -> dict[str, Any]:
         "culture": CULTURAL_KNOWLEDGE_TABLE_SCHEMA,
         "versions": VERSIONS_TABLE_SCHEMA,
         "traces": TRACE_TABLE_SCHEMA,
-        "spans": SPAN_TABLE_SCHEMA,
     }
 
     schema = schemas.get(table_type, {})

agno/db/postgres/async_postgres.py

@@ -31,7 +31,7 @@ from agno.utils.log import log_debug, log_error, log_info, log_warning
 from agno.utils.string import sanitize_postgres_string, sanitize_postgres_strings
 
 try:
-    from sqlalchemy import Index, String, Table, UniqueConstraint, and_, case, func, or_, update
+    from sqlalchemy import ForeignKey, Index, String, Table, UniqueConstraint, and_, case, func, or_, update
     from sqlalchemy.dialects import postgresql
     from sqlalchemy.dialects.postgresql import TIMESTAMP
     from sqlalchemy.exc import ProgrammingError
@@ -192,7 +192,10 @@ class AsyncPostgresDb(AsyncBaseDb):
             Table: SQLAlchemy Table object
         """
         try:
-            table_schema = get_table_schema_definition(table_type).copy()
+            # Pass traces_table_name and db_schema for spans table foreign key resolution
+            table_schema = get_table_schema_definition(
+                table_type, traces_table_name=self.trace_table_name, db_schema=self.db_schema
+            ).copy()
 
             columns: List[Column] = []
             indexes: List[str] = []
@@ -212,6 +215,11 @@ class AsyncPostgresDb(AsyncBaseDb):
                 if col_config.get("unique", False):
                     column_kwargs["unique"] = True
                     unique_constraints.append(col_name)
+
+                # Handle foreign key constraint
+                if "foreign_key" in col_config:
+                    column_args.append(ForeignKey(col_config["foreign_key"]))
+
                 columns.append(Column(*column_args, **column_kwargs))  # type: ignore
 
             # Create the table object

agno/db/postgres/postgres.py

@@ -176,7 +176,10 @@ class PostgresDb(BaseDb):
             Table: SQLAlchemy Table object
         """
         try:
-            table_schema = get_table_schema_definition(table_type).copy()
+            # Pass traces_table_name and db_schema for spans table foreign key resolution
+            table_schema = get_table_schema_definition(
+                table_type, traces_table_name=self.trace_table_name, db_schema=self.db_schema
+            ).copy()
 
             columns: List[Column] = []
             indexes: List[str] = []
@@ -199,12 +202,7 @@ class PostgresDb(BaseDb):
 
                 # Handle foreign key constraint
                 if "foreign_key" in col_config:
-                    fk_ref = col_config["foreign_key"]
-                    # For spans table, dynamically replace the traces table reference
-                    # with the actual trace table name configured for this db instance
-                    if table_type == "spans" and "trace_id" in fk_ref:
-                        fk_ref = f"{self.db_schema}.{self.trace_table_name}.trace_id"
-                    column_args.append(ForeignKey(fk_ref))
+                    column_args.append(ForeignKey(col_config["foreign_key"]))
 
                 columns.append(Column(*column_args, **column_kwargs))  # type: ignore
 

agno/db/postgres/schemas.py

@@ -137,37 +137,56 @@ TRACE_TABLE_SCHEMA = {
     "created_at": {"type": String, "nullable": False, "index": True},  # ISO 8601 datetime string
 }
 
-SPAN_TABLE_SCHEMA = {
-    "span_id": {"type": String, "primary_key": True, "nullable": False},
-    "trace_id": {
-        "type": String,
-        "nullable": False,
-        "index": True,
-        "foreign_key": "agno_traces.trace_id",  # Foreign key to traces table
-    },
-    "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
-}
 
+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
+    }
 
-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,
@@ -177,7 +196,6 @@ def get_table_schema_definition(table_type: str) -> dict[str, Any]:
         "culture": CULTURAL_KNOWLEDGE_TABLE_SCHEMA,
         "versions": VERSIONS_TABLE_SCHEMA,
         "traces": TRACE_TABLE_SCHEMA,
-        "spans": SPAN_TABLE_SCHEMA,
     }
 
     schema = schemas.get(table_type, {})

agno/db/singlestore/schemas.py

@@ -131,35 +131,56 @@ TRACE_TABLE_SCHEMA = {
     "created_at": {"type": lambda: String(64), "nullable": False, "index": True},  # ISO 8601 datetime string
 }
 
-SPAN_TABLE_SCHEMA = {
-    "span_id": {"type": lambda: String(128), "primary_key": True, "nullable": False},
-    "trace_id": {
-        "type": lambda: String(128),
-        "nullable": False,
-        "index": True,
-        "foreign_key": "agno_traces.trace_id",  # Foreign key to traces table
-    },
-    "parent_span_id": {"type": lambda: String(128), "nullable": True, "index": True},
-    "name": {"type": lambda: String(512), "nullable": False},
-    "span_kind": {"type": lambda: String(50), "nullable": False},
-    "status_code": {"type": lambda: String(20), "nullable": False},
-    "status_message": {"type": Text, "nullable": True},
-    "start_time": {"type": lambda: String(64), "nullable": False, "index": True},  # ISO 8601 datetime string
-    "end_time": {"type": lambda: String(64), "nullable": False},  # ISO 8601 datetime string
-    "duration_ms": {"type": BigInteger, "nullable": False},
-    "attributes": {"type": JSON, "nullable": True},
-    "created_at": {"type": lambda: String(64), "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": lambda: String(128), "primary_key": True, "nullable": False},
+        "trace_id": {
+            "type": lambda: String(128),
+            "nullable": False,
+            "index": True,
+            "foreign_key": f"{db_schema}.{traces_table_name}.trace_id",
+        },
+        "parent_span_id": {"type": lambda: String(128), "nullable": True, "index": True},
+        "name": {"type": lambda: String(512), "nullable": False},
+        "span_kind": {"type": lambda: String(50), "nullable": False},
+        "status_code": {"type": lambda: String(20), "nullable": False},
+        "status_message": {"type": Text, "nullable": True},
+        "start_time": {"type": lambda: String(64), "nullable": False, "index": True},  # ISO 8601 datetime string
+        "end_time": {"type": lambda: String(64), "nullable": False},  # ISO 8601 datetime string
+        "duration_ms": {"type": BigInteger, "nullable": False},
+        "attributes": {"type": JSON, "nullable": True},
+        "created_at": {"type": lambda: String(64), "nullable": False, "index": True},  # ISO 8601 datetime string
+    }
 
-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,
@@ -169,7 +190,6 @@ def get_table_schema_definition(table_type: str) -> dict[str, Any]:
         "culture": CULTURAL_KNOWLEDGE_TABLE_SCHEMA,
         "versions": VERSIONS_TABLE_SCHEMA,
         "traces": TRACE_TABLE_SCHEMA,
-        "spans": SPAN_TABLE_SCHEMA,
     }
     schema = schemas.get(table_type, {})
 

agno/db/singlestore/singlestore.py

@@ -31,7 +31,7 @@ from agno.utils.log import log_debug, log_error, log_info, log_warning
 from agno.utils.string import generate_id
 
 try:
-    from sqlalchemy import Index, UniqueConstraint, and_, func, select, update
+    from sqlalchemy import ForeignKey, Index, UniqueConstraint, and_, func, select, update
     from sqlalchemy.dialects import mysql
     from sqlalchemy.engine import Engine, create_engine
     from sqlalchemy.orm import scoped_session, sessionmaker
@@ -151,7 +151,10 @@ class SingleStoreDb(BaseDb):
             Table: SQLAlchemy Table object with column definitions
         """
         try:
-            table_schema = get_table_schema_definition(table_type)
+            # Pass traces_table_name and db_schema for spans table foreign key resolution
+            table_schema = get_table_schema_definition(
+                table_type, traces_table_name=self.trace_table_name, db_schema=self.db_schema or "agno"
+            )
 
             columns: List[Column] = []
             # Get the columns from the table schema
@@ -207,7 +210,10 @@ class SingleStoreDb(BaseDb):
         """
         table_ref = f"{self.db_schema}.{table_name}" if self.db_schema else table_name
         try:
-            table_schema = get_table_schema_definition(table_type)
+            # Pass traces_table_name and db_schema for spans table foreign key resolution
+            table_schema = get_table_schema_definition(
+                table_type, traces_table_name=self.trace_table_name, db_schema=self.db_schema or "agno"
+            ).copy()
 
             columns: List[Column] = []
             indexes: List[str] = []
@@ -227,6 +233,11 @@ class SingleStoreDb(BaseDb):
                 if col_config.get("unique", False):
                     column_kwargs["unique"] = True
                     unique_constraints.append(col_name)
+
+                # Handle foreign key constraint
+                if "foreign_key" in col_config:
+                    column_args.append(ForeignKey(col_config["foreign_key"]))
+
                 columns.append(Column(*column_args, **column_kwargs))
 
             # Create the table object

agno/db/sqlite/async_sqlite.py

@@ -31,7 +31,7 @@ from agno.utils.log import log_debug, log_error, log_info, log_warning
 from agno.utils.string import generate_id
 
 try:
-    from sqlalchemy import Column, MetaData, String, Table, func, select, text
+    from sqlalchemy import Column, ForeignKey, MetaData, String, Table, func, select, text
     from sqlalchemy.dialects import sqlite
     from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
     from sqlalchemy.schema import Index, UniqueConstraint
@@ -174,7 +174,8 @@ class AsyncSqliteDb(AsyncBaseDb):
             Table: SQLAlchemy Table object
         """
         try:
-            table_schema = get_table_schema_definition(table_type)
+            # Pass traces_table_name for spans table foreign key resolution
+            table_schema = get_table_schema_definition(table_type, traces_table_name=self.trace_table_name).copy()
 
             columns: List[Column] = []
             indexes: List[str] = []
@@ -196,6 +197,10 @@ class AsyncSqliteDb(AsyncBaseDb):
                     column_kwargs["unique"] = True
                     unique_constraints.append(col_name)
 
+                # Handle foreign key constraint
+                if "foreign_key" in col_config:
+                    column_args.append(ForeignKey(col_config["foreign_key"]))
+
                 columns.append(Column(*column_args, **column_kwargs))  # type: ignore
 
             # Create the table object