MindsDB 25.1.2.0__py3-none-any.whl → 25.1.5.0__py3-none-any.whl

mindsdb/integrations/utilities/rag/retrievers/sql_retriever.py

@@ -1,17 +1,37 @@
 import json
-from typing import List, Optional
+import re
+from pydantic import BaseModel, Field
+from typing import Any, List, Optional
 
 from langchain.chains.llm import LLMChain
 from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
 from langchain_core.documents.base import Document
 from langchain_core.embeddings import Embeddings
+from langchain_core.exceptions import OutputParserException
 from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.output_parsers import PydanticOutputParser
 from langchain_core.prompts import PromptTemplate
 from langchain_core.retrievers import BaseRetriever
 
 from mindsdb.api.executor.data_types.response_type import RESPONSE_TYPE
+from mindsdb.integrations.libs.response import HandlerResponse
 from mindsdb.integrations.libs.vectordatabase_handler import DistanceFunction, VectorStoreHandler
 from mindsdb.integrations.utilities.rag.settings import LLMExample, MetadataSchema, SearchKwargs
+from mindsdb.utilities import log
+
+logger = log.getLogger(__name__)
+
+
+class MetadataFilter(BaseModel):
+    '''Represents an LLM generated metadata filter to apply to a PostgreSQL query.'''
+    attribute: str = Field(description="Database column to apply filter to")
+    comparator: str = Field(description="PostgreSQL comparator to use to filter database column")
+    value: Any = Field(description="Value to use to filter database column")
+
+
+class MetadataFilters(BaseModel):
+    '''List of LLM generated metadata filters to apply to a PostgreSQL query.'''
+    filters: List[MetadataFilter] = Field(description="List of PostgreSQL metadata filters to apply for user query")
 
 
 class SQLRetriever(BaseRetriever):
@@ -22,32 +42,34 @@ class SQLRetriever(BaseRetriever):
     1. Use a LLM to rewrite the user input to something more suitable for retrieval. For example:
     "Show me documents containing how to finetune a LLM please" --> "how to finetune a LLM"
 
-    2. Use a LLM to generate a pgvector query with metadata filters based on the user input. Provided
-    metadata schemas & examples are used as additional context to generate the query.
+    2. Use a LLM to generate structured metadata filters based on the user input. Provided
+    metadata schemas & examples are used as additional context.
 
-    3. Use a LLM to double check the generated pgvector query is correct.
+    3. Generate a prepared PostgreSQL query from the structured metadata filters.
 
     4. Actually execute the query against our vector database to retrieve documents & return them.
     '''
+    fallback_retriever: BaseRetriever
     vector_store_handler: VectorStoreHandler
     metadata_schemas: Optional[List[MetadataSchema]] = None
     examples: Optional[List[LLMExample]] = None
 
-    embeddings_model: Embeddings
     rewrite_prompt_template: str
-    sql_prompt_template: str
-    query_checker_template: str
+    metadata_filters_prompt_template: str
+    embeddings_model: Embeddings
+    num_retries: int
     embeddings_table: str
     source_table: str
+    source_id_column: str = 'Id'
     distance_function: DistanceFunction
     search_kwargs: SearchKwargs
 
     llm: BaseChatModel
 
-    def _prepare_sql_prompt(self) -> PromptTemplate:
+    def _prepare_metadata_prompt(self) -> PromptTemplate:
         base_prompt_template = PromptTemplate(
-            input_variables=['dialect', 'input', 'embeddings_table', 'source_table', 'embeddings', 'distance_function', 'schema', 'examples'],
-            template=self.sql_prompt_template
+            input_variables=['format_instructions', 'schema', 'examples', 'input', 'embeddings'],
+            template=self.metadata_filters_prompt_template
         )
         schema_prompt_str = ''
         if self.metadata_schemas is not None:
@@ -61,7 +83,7 @@ class SQLRetriever(BaseRetriever):
                     if column.values is not None:
                         column_mapping[column.name]['values'] = column.values
                 column_mapping_json_str = json.dumps(column_mapping, indent=4)
-                schema_str = f'''{i+2}. {schema.table} - {schema.description}
+                schema_str = f'''{i+1}. {schema.table} - {schema.description}
 
 Columns:
 ```json
@@ -80,7 +102,7 @@ Output:
 {example.output}
 
 '''
-            examples_prompt_str += example_str
+                examples_prompt_str += example_str
         return base_prompt_template.partial(
             schema=schema_prompt_str,
             examples=examples_prompt_str
@@ -94,51 +116,94 @@ Output:
         rewrite_chain = LLMChain(llm=self.llm, prompt=rewrite_prompt)
         return rewrite_chain.predict(input=query)
 
-    def _prepare_pgvector_query(self, query: str, run_manager: CallbackManagerForRetrieverRun) -> str:
-        # Incorporate metadata schemas & examples into prompt.
-        sql_prompt = self._prepare_sql_prompt()
-        sql_chain = LLMChain(llm=self.llm, prompt=sql_prompt)
-        # Generate the initial pgvector query.
-        sql_query = sql_chain.predict(
-            # Only pgvector & similarity search is supported.
-            dialect='postgres',
-            input=query,
-            embeddings_table=self.embeddings_table,
-            source_table=self.source_table,
-            distance_function=self.distance_function.value[0],
-            k=self.search_kwargs.k,
-            callbacks=run_manager.get_child() if run_manager else None
-        )
-        query_checker_prompt = PromptTemplate(
-            input_variables=['dialect', 'query'],
-            template=self.query_checker_template
-        )
-        query_checker_chain = LLMChain(llm=self.llm, prompt=query_checker_prompt)
-        # Check the query & return the final result to be executed.
-        return query_checker_chain.predict(
-            dialect='postgres',
-            query=sql_query
+    def _prepare_pgvector_query(self, metadata_filters: List[MetadataFilter]) -> str:
+        # Base select JOINed with document source table.
+        base_query = f'''SELECT * FROM {self.embeddings_table} AS e INNER JOIN {self.source_table} AS s ON (e.metadata->>'original_row_id')::int = s."{self.source_id_column}" '''
+        col_to_schema = {}
+        if not self.metadata_schemas:
+            return ''
+        for schema in self.metadata_schemas:
+            for col in schema.columns:
+                col_to_schema[col.name] = schema
+        joined_schemas = set()
+        for filter in metadata_filters:
+            # Join schemas before filtering.
+            schema = col_to_schema.get(filter.attribute)
+            if schema is None or schema.table in joined_schemas or schema.table == self.source_table:
+                continue
+            joined_schemas.add(schema.table)
+            base_query += schema.join + ' '
+        # Actually construct WHERE conditions from metadata filters.
+        if metadata_filters:
+            base_query += 'WHERE '
+        for i, filter in enumerate(metadata_filters):
+            value = filter.value
+            if isinstance(value, str):
+                value = f"'{value}'"
+            base_query += f'"{filter.attribute}" {filter.comparator} {value}'
+            if i < len(metadata_filters) - 1:
+                base_query += ' AND '
+        base_query += f" ORDER BY e.embeddings {self.distance_function.value[0]} '{{embeddings}}' LIMIT {self.search_kwargs.k};"
+        return base_query
+
+    def _generate_metadata_filters(self, query: str) -> List[MetadataFilter]:
+        parser = PydanticOutputParser(pydantic_object=MetadataFilters)
+        metadata_prompt = self._prepare_metadata_prompt()
+        metadata_filters_chain = LLMChain(llm=self.llm, prompt=metadata_prompt)
+        metadata_filters_output = metadata_filters_chain.predict(
+            format_instructions=parser.get_format_instructions(),
+            input=query
         )
+        # If the LLM outputs raw JSON, use it as-is.
+        # If the LLM outputs anything including a json markdown section, use the last one.
+        json_markdown_output = re.findall(r'```json.*```', metadata_filters_output, re.DOTALL)
+        if json_markdown_output:
+            metadata_filters_output = json_markdown_output[-1]
+            # Clean the json tags.
+            metadata_filters_output = metadata_filters_output[7:]
+            metadata_filters_output = metadata_filters_output[:-3]
+        metadata_filters = parser.invoke(metadata_filters_output)
+        return metadata_filters.filters
+
+    def _prepare_and_execute_query(self, query: str, embeddings_str: str) -> HandlerResponse:
+        try:
+            metadata_filters = self._generate_metadata_filters(query)
+            checked_sql_query = self._prepare_pgvector_query(metadata_filters)
+            checked_sql_query_with_embeddings = checked_sql_query.format(embeddings=embeddings_str)
+            return self.vector_store_handler.native_query(checked_sql_query_with_embeddings)
+        except OutputParserException as e:
+            logger.warning(f'LLM failed to generate structured metadata filters: {str(e)}')
+            return HandlerResponse(RESPONSE_TYPE.ERROR, error_message=str(e))
+        except Exception as e:
+            logger.warning(f'Failed to prepare and execute SQL query from structured metadata: {str(e)}')
+            return HandlerResponse(RESPONSE_TYPE.ERROR, error_message=str(e))
 
     def _get_relevant_documents(
         self, query: str, *, run_manager: CallbackManagerForRetrieverRun
     ) -> List[Document]:
         # Rewrite query to be suitable for retrieval.
         retrieval_query = self._prepare_retrieval_query(query)
-
-        # Generate & check the query to be executed
-        checked_sql_query = self._prepare_pgvector_query(query, run_manager)
-
         # Embed the rewritten retrieval query & include it in the similarity search pgvector query.
         embedded_query = self.embeddings_model.embed_query(retrieval_query)
-        checked_sql_query_with_embeddings = checked_sql_query.format(embeddings=str(embedded_query))
-        # Handle LLM output that has the ```sql delimiter possibly.
-        checked_sql_query_with_embeddings = checked_sql_query_with_embeddings.replace('```sql', '')
-        checked_sql_query_with_embeddings = checked_sql_query_with_embeddings.replace('```', '')
         # Actually execute the similarity search with metadata filters.
-        document_response = self.vector_store_handler.native_query(checked_sql_query_with_embeddings)
-        if document_response.resp_type == RESPONSE_TYPE.ERROR:
-            raise ValueError(f'Retrieving documents failed with error {document_response.error_message}')
+        document_response = self._prepare_and_execute_query(retrieval_query, str(embedded_query))
+        num_retries = 0
+        while num_retries < self.num_retries:
+            if document_response.resp_type != RESPONSE_TYPE.ERROR and len(document_response.data_frame) > 0:
+                # Successfully retrieved documents.
+                break
+            if document_response.resp_type == RESPONSE_TYPE.ERROR:
+                # LLMs won't always generate structured metadata so we should have a fallback after retrying.
+                logger.info(f'SQL Retriever query failed with error {document_response.error_message}')
+            elif len(document_response.data_frame) == 0:
+                logger.info('No documents retrieved from SQL Retriever query')
+
+            document_response = self._prepare_and_execute_query(retrieval_query, str(embedded_query))
+            num_retries += 1
+            if num_retries >= self.num_retries:
+                logger.info('Using fallback retriever in SQL retriever.')
+                return self.fallback_retriever._get_relevant_documents(retrieval_query, run_manager=run_manager)
+
         document_df = document_response.data_frame
         retrieved_documents = []
         for _, document_row in document_df.iterrows():
@@ -146,4 +211,8 @@ Output:
                 document_row.get('content', ''),
                 metadata=document_row.get('metadata', {})
             ))
-        return retrieved_documents
+        if retrieved_documents:
+            return retrieved_documents
+        # If the SQL query constructed did not return any documents, fallback.
+        logger.info('No documents returned from SQL retriever. using fallback retriever.')
+        return self.fallback_retriever._get_relevant_documents(retrieval_query, run_manager=run_manager)

mindsdb/integrations/utilities/rag/settings.py

@@ -3,7 +3,6 @@ from typing import List, Union, Any, Optional, Dict
 
 from langchain_community.vectorstores.chroma import Chroma
 from langchain_community.vectorstores.pgvector import PGVector
-from langchain_community.tools.sql_database.prompt import QUERY_CHECKER as DEFAULT_QUERY_CHECKER_PROMPT_TEMPLATE
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from langchain_core.language_models import BaseChatModel
@@ -94,6 +93,25 @@ Output only a single better search query and nothing else like in the example.
 Here is the user input: {input}
 '''
 
+DEFAULT_METADATA_FILTERS_PROMPT_TEMPLATE = '''Construct a list of PostgreSQL metadata filters to filter documents in the database based on the user input.
+
+<< INSTRUCTIONS >>
+{format_instructions}
+
+RETURN ONLY THE FINAL JSON. DO NOT EXPLAIN, JUST RETURN THE FINAL JSON.
+
+<< TABLES YOU HAVE ACCESS TO >>
+
+{schema}
+
+<< EXAMPLES >>
+
+{examples}
+
+Here is the user input:
+{input}
+'''
+
 DEFAULT_SQL_PROMPT_TEMPLATE = '''
 Construct a valid {dialect} SQL query to select documents relevant to the user input.
 Source documents are found in the {source_table} table. You may need to join with other tables to get additional document metadata.
@@ -136,7 +154,6 @@ Columns:
         "description": "Metadata for the document chunk. Always select metadata and always join with the {source_table} table on the string metadata field 'original_row_id'"
     }}
 }}
-```
 
 {schema}
 
@@ -150,11 +167,118 @@ Here is the user input:
 {input}
 '''
 
+DEFAULT_QUESTION_REFORMULATION_TEMPLATE = """Given the original question and the retrieved context,
+analyze what additional information is needed for a complete, accurate answer.
+
+Original Question: {question}
+
+Retrieved Context:
+{context}
+
+Analysis Instructions:
+1. Evaluate Context Coverage:
+   - Identify key entities and concepts from the question
+   - Check for temporal information (dates, periods, sequences)
+   - Verify causal relationships are explained
+   - Confirm presence of requested quantitative data
+   - Assess if geographic or spatial context is sufficient
+
+2. Quality Assessment:
+   If the retrieved context is:
+   - Irrelevant or tangential
+   - Too general or vague
+   - Potentially contradictory
+   - Missing key perspectives
+   - Lacking proper evidence
+   Generate questions to address these specific gaps.
+
+3. Follow-up Question Requirements:
+   - Questions must directly contribute to answering the original query
+   - Break complex relationships into simpler, sequential steps
+   - Maintain specificity rather than broad inquiries
+   - Avoid questions answerable from existing context
+   - Ensure questions build on each other logically
+   - Limit questions to 150 characters each
+   - Each question must be self-contained
+   - Questions must end with a question mark
+
+4. Response Format:
+   - Return a JSON array of strings
+   - Use square brackets and double quotes
+   - Questions must be unique (no duplicates)
+   - If context is sufficient, return empty array []
+   - Maximum 3 follow-up questions
+   - Minimum length per question: 30 characters
+   - No null values or empty strings
+
+Example:
+Original: "How did the development of antibiotics affect military casualties in WWII?"
+
+Invalid responses:
+{'questions': ['What are antibiotics?']}  // Wrong format
+['What is WWII?']  // Too basic
+['How did it impact things?']  // Too vague
+['', 'Question 2']  // Contains empty string
+['Same question?', 'Same question?']  // Duplicate
+
+Valid response:
+["What were military casualty rates from infections before widespread antibiotic use in 1942?",
+ "How did penicillin availability change throughout different stages of WWII?",
+ "What were the primary battlefield infections treated with antibiotics during WWII?"]
+
+or [] if context fully answers the original question.
+
+Your task: Based on the analysis of the original question and context,
+output ONLY a JSON array of follow-up questions needed to provide a complete answer.
+If no additional information is needed, output an empty array [].
+
+Follow-up Questions:"""
+
+DEFAULT_QUERY_RETRY_PROMPT_TEMPLATE = '''
+{query}
+
+The {dialect} query above failed with the error message: {error}.
+
+<< TABLES YOU HAVE ACCESS TO >>
+1. {embeddings_table} - Contains document chunks, vector embeddings, and metadata for documents.
+
+Columns:
+```json
+{{
+    "id": {{
+        "type": "string",
+        "description": "Unique ID for this document chunk"
+    }},
+    "content": {{
+        "type": "string",
+        "description": "A document chunk (subset of the original document)"
+    }},
+    "embeddings": {{
+        "type": "vector",
+        "description": "Vector embeddings for the document chunk."
+    }},
+    "metadata": {{
+        "type": "jsonb",
+        "description": "Metadata for the document chunk."
+    }}
+}}
+
+{schema}
+
+Rewrite the query so it works.
+
+Output the final SQL query only.
+
+SQL Query:
+'''
+
+DEFAULT_NUM_QUERY_RETRIES = 2
+
 
 class LLMConfig(BaseModel):
     model_name: str = Field(default=DEFAULT_LLM_MODEL, description='LLM model to use for generation')
     provider: str = Field(default=DEFAULT_LLM_MODEL_PROVIDER, description='LLM model provider to use for generation')
-    params: Dict[str, Any] = {}
+    params: Dict[str, Any] = Field(default_factory=dict)
 
 
 class MultiVectorRetrieverMode(Enum):
@@ -183,17 +307,21 @@ class VectorStoreConfig(BaseModel):
     collection_name: str = DEFAULT_COLLECTION_NAME
     connection_string: str = None
     kb_table: Any = None
+    is_sparse: bool = False
+    vector_size: Optional[int] = None
 
     class Config:
         arbitrary_types_allowed = True
         extra = "forbid"
 
 
-class RetrieverType(Enum):
-    VECTOR_STORE = 'vector_store'
-    AUTO = 'auto'
-    MULTI = 'multi'
-    SQL = 'sql'
+class RetrieverType(str, Enum):
+    """Retriever type for RAG pipeline"""
+    VECTOR_STORE = "vector_store"
+    AUTO = "auto"
+    MULTI = "multi"
+    SQL = "sql"
+    MULTI_HOP = "multi_hop"
 
 
 class SearchType(Enum):
@@ -267,6 +395,13 @@ class MetadataSchema(BaseModel):
     columns: List[ColumnSchema] = Field(
         description="List of column schemas describing the metadata columns available for the table"
     )
+    join: str = Field(
+        description="SQL join string to join this table with source documents table",
+        default=''
+    )
+
+    class Config:
+        frozen = True
 
 
 class LLMExample(BaseModel):
@@ -283,15 +418,13 @@ class SQLRetrieverConfig(BaseModel):
         default_factory=LLMConfig,
         description="LLM configuration to use for generating the final SQL query for retrieval"
     )
-    sql_prompt_template: str = Field(
-        default=DEFAULT_SQL_PROMPT_TEMPLATE,
-        description="""Prompt template to generate the SQL query to execute against the vector database. Currently only pgvector is supported.
-        Has 'dialect', 'input', 'embeddings_table', 'source_table', 'embeddings', 'distance_function', 'schema', and 'examples' input variables.
-        """
+    metadata_filters_prompt_template: str = Field(
+        default=DEFAULT_METADATA_FILTERS_PROMPT_TEMPLATE,
+        description="Prompt template to generate PostgreSQL metadata filters. Has 'format_instructions', 'schema', 'examples', and 'input' input variables"
     )
-    query_checker_template: str = Field(
-        default=DEFAULT_QUERY_CHECKER_PROMPT_TEMPLATE,
-        description="Prompt template to use for double checking SQL queries before execution. Has 'query' and 'dialect' input variables."
+    num_retries: int = Field(
+        default=DEFAULT_NUM_QUERY_RETRIES,
+        description="How many times for an LLM to try rewriting a failed SQL query before using the fallback retriever."
     )
     rewrite_prompt_template: str = Field(
         default=DEFAULT_SEMANTIC_PROMPT_TEMPLATE,
@@ -332,8 +465,34 @@ class SummarizationConfig(BaseModel):
 class RerankerConfig(BaseModel):
     model: str = DEFAULT_RERANKING_MODEL
     base_url: str = DEFAULT_LLM_ENDPOINT
-    filtering_threshold: float = 0.99
+    filtering_threshold: float = 0.5
     num_docs_to_keep: Optional[int] = None
+    max_concurrent_requests: int = 20
+    max_retries: int = 3
+    retry_delay: float = 1.0
+    early_stop: bool = True  # Whether to enable early stopping
+    early_stop_threshold: float = 0.8  # Confidence threshold for early stopping
+
+
+class MultiHopRetrieverConfig(BaseModel):
+    """Configuration for multi-hop retrieval"""
+    base_retriever_type: RetrieverType = Field(
+        default=RetrieverType.VECTOR_STORE,
+        description="Type of base retriever to use for multi-hop retrieval"
+    )
+    max_hops: int = Field(
+        default=3,
+        description="Maximum number of follow-up questions to generate",
+        ge=1
+    )
+    reformulation_template: str = Field(
+        default=DEFAULT_QUESTION_REFORMULATION_TEMPLATE,
+        description="Template for reformulating questions"
+    )
+    llm_config: LLMConfig = Field(
+        default_factory=LLMConfig,
+        description="LLM configuration to use for generating follow-up questions"
+    )
 
 
 class RAGPipelineModel(BaseModel):
@@ -462,6 +621,20 @@ class RAGPipelineModel(BaseModel):
         description="Reranker configuration"
     )
 
+    multi_hop_config: Optional[MultiHopRetrieverConfig] = Field(
+        default=None,
+        description="Configuration for multi-hop retrieval. Required when retriever_type is MULTI_HOP."
+    )
+
+    @field_validator("multi_hop_config")
+    @classmethod
+    def validate_multi_hop_config(cls, v: Optional[MultiHopRetrieverConfig], info):
+        """Validate that multi_hop_config is set when using multi-hop retrieval."""
+        values = info.data
+        if values.get("retriever_type") == RetrieverType.MULTI_HOP and v is None:
+            raise ValueError("multi_hop_config must be set when using multi-hop retrieval")
+        return v
+
     class Config:
         arbitrary_types_allowed = True
         extra = "forbid"

mindsdb/integrations/utilities/sql_utils.py

@@ -178,7 +178,7 @@ def project_dataframe(df, targets, table_columns):
 
     # adapt column names to projection
     if len(df_col_rename) > 0:
-        df = df.rename(columns=df_col_rename)
+        df.rename(columns=df_col_rename, inplace=True)
     return df
 
 

mindsdb/interfaces/agents/agents_controller.py

@@ -1,5 +1,5 @@
 import datetime
-from typing import Dict, Iterator, List, Union, Tuple
+from typing import Dict, Iterator, List, Union, Tuple, Optional
 
 from langchain_core.tools import BaseTool
 from sqlalchemy.orm.attributes import flag_modified
@@ -70,7 +70,7 @@ class AgentsController:
 
         return model, provider
 
-    def get_agent(self, agent_name: str, project_name: str = 'mindsdb') -> db.Agents:
+    def get_agent(self, agent_name: str, project_name: str = 'mindsdb') -> Optional[db.Agents]:
         '''
         Gets an agent by name.
 
@@ -79,7 +79,7 @@ class AgentsController:
             project_name (str): The name of the containing project - must exist
 
         Returns:
-            agent (db.Agents): The database agent object
+            agent (Optional[db.Agents]): The database agent object
         '''
 
         project = self.project_controller.get(name=project_name)
@@ -252,6 +252,16 @@ class AgentsController:
         existing_agent = self.get_agent(agent_name, project_name=project_name)
         if existing_agent is None:
             raise EntityNotExistsError(f'Agent with name not found: {agent_name}')
+        is_demo = (existing_agent.params or {}).get('is_demo', False)
+        if (
+            is_demo and (
+                (name is not None and name != agent_name)
+                or (model_name or provider)
+                or (len(skills_to_add) > 0 or len(skills_to_remove) > 0 or len(skills_to_rewrite) > 0)
+                or (isinstance(params, dict) and len(params) > 1 and 'prompt_template' not in params)
+            )
+        ):
+            raise ValueError("It is forbidden to change properties of the demo object")
 
         if name is not None and name != agent_name:
             # Check to see if updated name already exists
@@ -352,6 +362,8 @@ class AgentsController:
         agent = self.get_agent(agent_name, project_name)
         if agent is None:
             raise ValueError(f'Agent with name does not exist: {agent_name}')
+        if isinstance(agent.params, dict) and agent.params.get('is_demo') is True:
+            raise ValueError('Unable to delete demo object')
         agent.deleted_at = datetime.datetime.now()
         db.session.commit()
 
@@ -362,24 +374,22 @@ class AgentsController:
             project_name: str = 'mindsdb',
             tools: List[BaseTool] = None,
             stream: bool = False) -> Union[Iterator[object], pd.DataFrame]:
-        '''
+        """
         Queries an agent to get a completion.
 
         Parameters:
             agent (db.Agents): Existing agent to get completion from
             messages (List[Dict[str, str]]): Chat history to send to the agent
-            trace_id (str): ID of Langfuse trace to use
-            observation_id (str): ID of parent Langfuse observation to use
             project_name (str): Project the agent belongs to (default mindsdb)
             tools (List[BaseTool]): Tools to use while getting the completion
-            stream (bool): Whether or not to stream the response
+            stream (bool): Whether to stream the response
 
         Returns:
             response (Union[Iterator[object], pd.DataFrame]): Completion as a DataFrame or iterator of completion chunks
 
         Raises:
             ValueError: Agent's model does not exist.
-        '''
+        """
         if stream:
             return self._get_completion_stream(
                 agent,

mindsdb/interfaces/agents/constants.py

@@ -165,6 +165,7 @@ PROVIDER_TO_MODELS = MappingProxyType(
 
 ASSISTANT_COLUMN = "answer"
 CONTEXT_COLUMN = "context"
+TRACE_ID_COLUMN = "trace_id"
 DEFAULT_AGENT_TIMEOUT_SECONDS = 300
 # These should require no additional arguments.
 DEFAULT_AGENT_TOOLS = []