MindsDB 25.7.3.0__py3-none-any.whl → 25.8.2.0__py3-none-any.whl

mindsdb/interfaces/agents/agents_controller.py

@@ -20,7 +20,7 @@ from mindsdb.utilities import log
 
 from mindsdb.utilities.exception import EntityExistsError, EntityNotExistsError
 
-from .constants import ASSISTANT_COLUMN, SUPPORTED_PROVIDERS, PROVIDER_TO_MODELS, DEFAULT_TEXT2SQL_DATABASE
+from .constants import ASSISTANT_COLUMN, SUPPORTED_PROVIDERS, PROVIDER_TO_MODELS
 from .langchain_agent import get_llm_provider
 
 logger = log.getLogger(__name__)
@@ -145,11 +145,60 @@ class AgentsController:
 
         return all_agents.all()
 
+    def _create_default_sql_skill(
+        self,
+        name,
+        project_name,
+        include_tables: List[str] = None,
+        include_knowledge_bases: List[str] = None,
+    ):
+        # Create a default SQL skill
+        skill_name = f"{name}_sql_skill"
+        skill_params = {
+            "type": "sql",
+            "description": f"Auto-generated SQL skill for agent {name}",
+        }
+
+        # Add restrictions provided
+        if include_tables:
+            skill_params["include_tables"] = include_tables
+        if include_knowledge_bases:
+            skill_params["include_knowledge_bases"] = include_knowledge_bases
+
+        try:
+            # Check if skill already exists
+            existing_skill = self.skills_controller.get_skill(skill_name, project_name)
+            if existing_skill is None:
+                # Create the skill
+                skill_type = skill_params.pop("type")
+                self.skills_controller.add_skill(
+                    name=skill_name, project_name=project_name, type=skill_type, params=skill_params
+                )
+            else:
+                # Update the skill if parameters have changed
+                params_changed = False
+
+                # Check if skill parameters need to be updated
+                for param_key, param_value in skill_params.items():
+                    if existing_skill.params.get(param_key) != param_value:
+                        existing_skill.params[param_key] = param_value
+                        params_changed = True
+
+                # Update the skill if needed
+                if params_changed:
+                    flag_modified(existing_skill, "params")
+                    db.session.commit()
+
+        except Exception as e:
+            raise ValueError(f"Failed to auto-create or update SQL skill: {str(e)}")
+
+        return skill_name
+
     def add_agent(
         self,
         name: str,
         project_name: str = None,
-        model_name: str = None,
+        model_name: Union[str, dict] = None,
         skills: List[Union[str, dict]] = None,
         provider: str = None,
         params: Dict[str, Any] = None,
@@ -165,16 +214,21 @@ class AgentsController:
                  with one of keys is "name", and other is additional parameters for relationship agent<>skill
             provider (str): The provider of the model
             params (Dict[str, str]): Parameters to use when running the agent
+                data: Dict, data sources for an agent, keys:
+                  - knowledge_bases: List of KBs to use
+                  - tables: list of tables to use
+                model: Dict, parameters for the model to use
+                  - provider: The provider of the model (e.g., 'openai', 'google')
+                  - Other model-specific parameters like 'api_key', 'model_name', etc.
+                <provider>_api_key: API key for the provider (e.g., openai_api_key)
+
+                # Deprecated parameters:
                 database: The database to use for text2sql skills (default is 'mindsdb')
                 knowledge_base_database: The database to use for knowledge base queries (default is 'mindsdb')
                 include_tables: List of tables to include for text2sql skills
                 ignore_tables: List of tables to ignore for text2sql skills
                 include_knowledge_bases: List of knowledge bases to include for text2sql skills
                 ignore_knowledge_bases: List of knowledge bases to ignore for text2sql skills
-                <provider>_api_key: API key for the provider (e.g., openai_api_key)
-                data: Dict, data sources for an agent, keys:
-                  - knowledge_bases: List of KBs to use (alternative to `include_knowledge_bases`)
-                  - tables: list of tables to use (alternative to `include_tables`)
 
         Returns:
             agent (db.Agents): The created agent
@@ -223,99 +277,41 @@ class AgentsController:
             # It will be picked up by get_api_key() in handler_utils.py
             pass
 
-        # Extract table and knowledge base parameters from params
-        database = params.pop("database", None)
-        knowledge_base_database = params.pop("knowledge_base_database", DEFAULT_TEXT2SQL_DATABASE)
-        include_tables = params.pop("include_tables", None)
-        ignore_tables = params.pop("ignore_tables", None)
-        include_knowledge_bases = params.pop("include_knowledge_bases", None)
-        ignore_knowledge_bases = params.pop("ignore_knowledge_bases", None)
-
-        # Save the extracted parameters back to params for API responses only if they were explicitly provided
-        # or if they're needed for skills
-        need_params = include_tables or ignore_tables or include_knowledge_bases or ignore_knowledge_bases
-
-        if "database" in params or need_params:
-            params["database"] = database
+        depreciated_params = [
+            "database",
+            "knowledge_base_database",
+            "include_tables",
+            "ignore_tables",
+            "include_knowledge_bases",
+            "ignore_knowledge_bases",
+        ]
+        if any(param in params for param in depreciated_params):
+            raise ValueError(
+                f"Parameters {', '.join(depreciated_params)} are deprecated. "
+                "Use 'data' parameter with 'tables' and 'knowledge_bases' keys instead."
+            )
 
+        include_tables = None
+        include_knowledge_bases = None
         if "data" in params:
-            if include_knowledge_bases is None:
-                include_knowledge_bases = params["data"].get("knowledge_bases")
-            if include_tables is None:
-                include_tables = params["data"].get("tables")
-
-        if "knowledge_base_database" in params or include_knowledge_bases or ignore_knowledge_bases:
-            params["knowledge_base_database"] = knowledge_base_database
-
-        if include_tables is not None:
-            params["include_tables"] = include_tables
-        if ignore_tables is not None:
-            params["ignore_tables"] = ignore_tables
-        if include_knowledge_bases is not None:
-            params["include_knowledge_bases"] = include_knowledge_bases
-        if ignore_knowledge_bases is not None:
-            params["ignore_knowledge_bases"] = ignore_knowledge_bases
+            include_knowledge_bases = params["data"].get("knowledge_bases")
+            include_tables = params["data"].get("tables")
 
         # Convert string parameters to lists if needed
         if isinstance(include_tables, str):
             include_tables = [t.strip() for t in include_tables.split(",")]
-        if isinstance(ignore_tables, str):
-            ignore_tables = [t.strip() for t in ignore_tables.split(",")]
         if isinstance(include_knowledge_bases, str):
             include_knowledge_bases = [kb.strip() for kb in include_knowledge_bases.split(",")]
-        if isinstance(ignore_knowledge_bases, str):
-            ignore_knowledge_bases = [kb.strip() for kb in ignore_knowledge_bases.split(",")]
 
         # Auto-create SQL skill if no skills are provided but include_tables or include_knowledge_bases params are provided
         if not skills and (include_tables or include_knowledge_bases):
-            # Create a default SQL skill
-            skill_name = f"{name}_sql_skill"
-            skill_params = {
-                "type": "sql",
-                "database": database,
-                "description": f"Auto-generated SQL skill for agent {name}",
-            }
-
-            # Add table restrictions if provided
-            if include_tables:
-                skill_params["include_tables"] = include_tables
-            if ignore_tables:
-                skill_params["ignore_tables"] = ignore_tables
-
-            # Add knowledge base parameters if provided
-            if knowledge_base_database:
-                skill_params["knowledge_base_database"] = knowledge_base_database
-            if include_knowledge_bases:
-                skill_params["include_knowledge_bases"] = include_knowledge_bases
-            if ignore_knowledge_bases:
-                skill_params["ignore_knowledge_bases"] = ignore_knowledge_bases
-            try:
-                # Check if skill already exists
-                existing_skill = self.skills_controller.get_skill(skill_name, project_name)
-                if existing_skill is None:
-                    # Create the skill
-                    skill_type = skill_params.pop("type")
-                    self.skills_controller.add_skill(
-                        name=skill_name, project_name=project_name, type=skill_type, params=skill_params
-                    )
-                else:
-                    # Update the skill if parameters have changed
-                    params_changed = False
-
-                    # Check if skill parameters need to be updated
-                    for param_key, param_value in skill_params.items():
-                        if existing_skill.params.get(param_key) != param_value:
-                            existing_skill.params[param_key] = param_value
-                            params_changed = True
-
-                    # Update the skill if needed
-                    if params_changed:
-                        flag_modified(existing_skill, "params")
-                        db.session.commit()
-
-                skills = [skill_name]
-            except Exception as e:
-                raise ValueError(f"Failed to auto-create or update SQL skill: {str(e)}")
+            skill = self._create_default_sql_skill(
+                name,
+                project_name,
+                include_tables=include_tables,
+                include_knowledge_bases=include_knowledge_bases,
+            )
+            skills = [skill]
 
         agent = db.Agents(
             name=name,
@@ -340,44 +336,16 @@ class AgentsController:
                 db.session.rollback()
                 raise ValueError(f"Skill with name does not exist: {skill_name}")
 
-            if existing_skill.type == "sql":
-                # Run Data Catalog loader if enabled
-                if config.get("data_catalog", {}).get("enabled", False):
-                    if include_tables:
-                        database_table_map = {}
-                        for table in include_tables:
-                            parts = table.split(".", 1)
-                            database_table_map[parts[0]] = database_table_map.get(parts[0], []) + [parts[1]]
-
-                        for database_name, table_names in database_table_map.items():
-                            data_catalog_loader = DataCatalogLoader(
-                                database_name=database_name, table_names=table_names
-                            )
-                            data_catalog_loader.load_metadata()
-
-                    elif "database" in existing_skill.params:
-                        data_catalog_loader = DataCatalogLoader(
-                            database_name=existing_skill.params["database"],
-                            table_names=parameters["tables"] if "tables" in parameters else None,
-                        )
-                        data_catalog_loader.load_metadata()
-
-                    else:
-                        raise ValueError(
-                            "Data Catalog loading is enabled, but the provided parameters are insufficient to load metadata. "
-                        )
+            # Run Data Catalog loader if enabled.
+            if include_tables:
+                self._run_data_catalog_loader_for_table_entries(include_tables, project_name, skill=existing_skill)
+            else:
+                self._run_data_catalog_loader_for_skill(existing_skill, project_name, tables=parameters.get("tables"))
 
+            if existing_skill.type == "sql":
                 # Add table restrictions if this is a text2sql skill
-                if include_tables or ignore_tables:
-                    parameters["tables"] = include_tables or ignore_tables
-
-                # Pass database parameter if provided
-                if database and "database" not in parameters:
-                    parameters["database"] = database
-
-                # Pass knowledge base database parameter if provided
-                if knowledge_base_database and "knowledge_base_database" not in parameters:
-                    parameters["knowledge_base_database"] = knowledge_base_database
+                if include_tables:
+                    parameters["tables"] = include_tables
 
                 # Add knowledge base parameters to both the skill and the association parameters
                 if include_knowledge_bases:
@@ -385,16 +353,6 @@ class AgentsController:
                     if "include_knowledge_bases" not in existing_skill.params:
                         existing_skill.params["include_knowledge_bases"] = include_knowledge_bases
                         flag_modified(existing_skill, "params")
-                elif ignore_knowledge_bases:
-                    parameters["ignore_knowledge_bases"] = ignore_knowledge_bases
-                    if "ignore_knowledge_bases" not in existing_skill.params:
-                        existing_skill.params["ignore_knowledge_bases"] = ignore_knowledge_bases
-                        flag_modified(existing_skill, "params")
-
-                # Ensure knowledge_base_database is set in the skill's params
-                if knowledge_base_database and "knowledge_base_database" not in existing_skill.params:
-                    existing_skill.params["knowledge_base_database"] = knowledge_base_database
-                    flag_modified(existing_skill, "params")
 
             association = db.AgentSkillsAssociation(parameters=parameters, agent=agent, skill=existing_skill)
             db.session.add(association)
@@ -409,7 +367,7 @@ class AgentsController:
         agent_name: str,
         project_name: str = default_project,
         name: str = None,
-        model_name: str = None,
+        model_name: Union[str, dict] = None,
         skills_to_add: List[Union[str, dict]] = None,
         skills_to_remove: List[str] = None,
         skills_to_rewrite: List[Union[str, dict]] = None,
@@ -423,7 +381,7 @@ class AgentsController:
             agent_name (str): The name of the new agent, or existing agent to update
             project_name (str): The containing project
             name (str): The updated name of the agent
-            model_name (str): The name of the existing ML model the agent will use
+            model_name (str | dict): The name of the existing ML model the agent will use
             skills_to_add (List[Union[str, dict]]): List of skill names to add to the agent, or list of dicts
                  with one of keys is "name", and other is additional parameters for relationship agent<>skill
             skills_to_remove (List[str]): List of skill names to remove from the agent
@@ -452,6 +410,8 @@ 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}")
+        existing_params = existing_agent.params or {}
+
         is_demo = (existing_agent.params or {}).get("is_demo", False)
         if is_demo and (
             (name is not None and name != agent_name)
@@ -462,6 +422,8 @@ class AgentsController:
             raise ValueError("It is forbidden to change properties of the demo object")
 
         if name is not None and name != agent_name:
+            if not name.islower():
+                raise ValueError(f"The name must be in lower case: {name}")
             # Check to see if updated name already exists
             agent_with_new_name = self.get_agent(name, project_name=project_name)
             if agent_with_new_name is not None:
@@ -469,12 +431,34 @@ class AgentsController:
             existing_agent.name = name
 
         if model_name or provider:
+            if isinstance(model_name, dict):
+                # move into params
+                existing_params["model"] = model_name
+                model_name = None
+
             # check model and provider
             model, provider = self.check_model_provider(model_name, provider)
             # Update model and provider
             existing_agent.model_name = model_name
             existing_agent.provider = provider
 
+        if "data" in params:
+            if len(skills_to_add) > 0 or len(skills_to_remove) > 0:
+                raise ValueError(
+                    "'data' parameter cannot be used with 'skills_to_remove' or 'skills_to_add' parameters"
+                )
+
+            include_knowledge_bases = params["data"].get("knowledge_bases")
+            include_tables = params["data"].get("tables")
+
+            skill = self._create_default_sql_skill(
+                agent_name,
+                project_name,
+                include_tables=include_tables,
+                include_knowledge_bases=include_knowledge_bases,
+            )
+            skills_to_rewrite = [{"name": skill}]
+
         # check that all skills exist
         skill_name_to_record_map = {}
         for skill_meta in skills_to_add + skills_to_remove + skills_to_rewrite:
@@ -503,12 +487,20 @@ class AgentsController:
 
             # add skills
             for skill_name in set(skills_to_add_names) - set(existing_agent_skills_names):
+                # Run Data Catalog loader if enabled for the new skill
+                self._run_data_catalog_loader_for_skill(
+                    skill_name,
+                    project_name,
+                    tables=next((x for x in skills_to_add if x["name"] == skill_name), {}).get("tables"),
+                )
+
                 skill_parameters = next(x for x in skills_to_add if x["name"] == skill_name).copy()
                 del skill_parameters["name"]
                 association = db.AgentSkillsAssociation(
                     parameters=skill_parameters, agent=existing_agent, skill=skill_name_to_record_map[skill_name]
                 )
                 db.session.add(association)
+
         elif len(skills_to_rewrite) > 0:
             skill_name_to_parameters = {
                 x["name"]: {k: v for k, v in x.items() if k != "name"} for x in skills_to_rewrite
@@ -519,9 +511,23 @@ class AgentsController:
                     db.session.delete(rel)
                 else:
                     existing_skill_names.add(rel.skill.name)
-                    rel.parameters = skill_name_to_parameters[rel.skill.name]
+                    skill_parameters = skill_name_to_parameters[rel.skill.name]
+
+                    # Run Data Catalog loader if enabled for the updated skill
+                    self._run_data_catalog_loader_for_skill(
+                        rel.skill.name, project_name, tables=skill_parameters.get("tables")
+                    )
+
+                    rel.parameters = skill_parameters
                     flag_modified(rel, "parameters")
             for new_skill_name in set(skill_name_to_parameters) - existing_skill_names:
+                # Run Data Catalog loader if enabled for the new skill
+                self._run_data_catalog_loader_for_skill(
+                    new_skill_name,
+                    project_name,
+                    tables=skill_name_to_parameters[new_skill_name].get("tables"),
+                )
+
                 association = db.AgentSkillsAssociation(
                     parameters=skill_name_to_parameters[new_skill_name],
                     agent=existing_agent,
@@ -530,8 +536,15 @@ class AgentsController:
                 db.session.add(association)
 
         if params is not None:
+            if params.get("data", {}).get("tables"):
+                new_table_entries = set(params["data"]["tables"]) - set(
+                    existing_params.get("data", {}).get("tables", [])
+                )
+                if new_table_entries:
+                    # Run Data Catalog loader for new table entries if enabled.
+                    self._run_data_catalog_loader_for_table_entries(new_table_entries, project_name)
+
             # Merge params on update
-            existing_params = existing_agent.params or {}
             existing_params.update(params)
             # Remove None values entirely.
             params = {k: v for k, v in existing_params.items() if v is not None}
@@ -543,6 +556,86 @@ class AgentsController:
 
         return existing_agent
 
+    def _run_data_catalog_loader_for_skill(
+        self,
+        skill: Union[str, db.Skills],
+        project_name: str,
+        tables: List[str] = None,
+    ):
+        """
+        Runs Data Catalog loader for a skill if enabled in the config.
+        This is used to load metadata for SQL skills when they are added or updated.
+        """
+        if not config.get("data_catalog", {}).get("enabled", False):
+            return
+
+        skill = skill if isinstance(skill, db.Skills) else self.skills_controller.get_skill(skill, project_name)
+        if skill.type == "sql":
+            if "database" in skill.params:
+                valid_table_names = skill.params.get("tables") if skill.params.get("tables") else tables
+                data_catalog_loader = DataCatalogLoader(
+                    database_name=skill.params["database"], table_names=valid_table_names
+                )
+                data_catalog_loader.load_metadata()
+            else:
+                raise ValueError(
+                    "Data Catalog loading is enabled, but the provided parameters for the new skills are insufficient to load metadata. "
+                )
+
+    def _run_data_catalog_loader_for_table_entries(
+        self,
+        table_entries: List[str],
+        project_name: str,
+        skill: Union[str, db.Skills] = None,
+    ):
+        """
+        Runs Data Catalog loader for a list of table entries if enabled in the config.
+        This is used to load metadata for SQL skills when they are added or updated.
+        """
+        if not config.get("data_catalog", {}).get("enabled", False):
+            return
+
+        skill = skill if isinstance(skill, db.Skills) else self.skills_controller.get_skill(skill, project_name)
+        if not skill or skill.type == "sql":
+            database_table_map = {}
+            for table_entry in table_entries:
+                parts = table_entry.split(".", 1)
+
+                # Ensure the table name is in 'database.table' format.
+                if len(parts) != 2:
+                    logger.warning(
+                        f"Invalid table name format: {table_entry}. Expected 'database.table' format."
+                        "Metadata will not be loaded for this entry."
+                    )
+                    continue
+
+                database, table = parts[0], parts[1]
+
+                # Wildcards in database names are not supported at the moment by data catalog loader.
+                if "*" in database:
+                    logger.warning(
+                        f"Invalid database name format: {database}. Wildcards are not supported."
+                        "Metadata will not be loaded for this entry."
+                    )
+                    continue
+
+                # Wildcards in table names are supported either.
+                # However, the table name itself can be a wildcard representing all tables.
+                if table == "*":
+                    table = None
+                elif "*" in table:
+                    logger.warning(
+                        f"Invalid table name format: {table}. Wildcards are not supported."
+                        "Metadata will not be loaded for this entry."
+                    )
+                    continue
+
+                database_table_map[database] = database_table_map.get(database, []) + [table]
+
+            for database_name, table_names in database_table_map.items():
+                data_catalog_loader = DataCatalogLoader(database_name=database_name, table_names=table_names)
+                data_catalog_loader.load_metadata()
+
     def delete_agent(self, agent_name: str, project_name: str = default_project):
         """
         Deletes an agent by name.
@@ -584,20 +677,22 @@ class AgentsController:
     def get_completion(
         self,
         agent: db.Agents,
-        messages: List[Dict[str, str]],
+        messages: list[Dict[str, str]],
         project_name: str = default_project,
-        tools: List[BaseTool] = None,
+        tools: list[BaseTool] = None,
         stream: bool = False,
+        params: dict | None = None,
     ) -> 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
+            messages (list[Dict[str, str]]): Chat history to send to the agent
             project_name (str): Project the agent belongs to (default mindsdb)
-            tools (List[BaseTool]): Tools to use while getting the completion
+            tools (list[BaseTool]): Tools to use while getting the completion
             stream (bool): Whether to stream the response
+            params (dict | None): params to redefine agent params
 
         Returns:
             response (Union[Iterator[object], pd.DataFrame]): Completion as a DataFrame or iterator of completion chunks
@@ -606,7 +701,7 @@ class AgentsController:
             ValueError: Agent's model does not exist.
         """
         if stream:
-            return self._get_completion_stream(agent, messages, project_name=project_name, tools=tools)
+            return self._get_completion_stream(agent, messages, project_name=project_name, tools=tools, params=params)
         from .langchain_agent import LangchainAgent
 
         model, provider = self.check_model_provider(agent.model_name, agent.provider)
@@ -619,25 +714,27 @@ class AgentsController:
         llm_params = self.get_agent_llm_params(agent.params)
 
         lang_agent = LangchainAgent(agent, model, llm_params=llm_params)
-        return lang_agent.get_completion(messages)
+        return lang_agent.get_completion(messages, params=params)
 
     def _get_completion_stream(
         self,
         agent: db.Agents,
-        messages: List[Dict[str, str]],
+        messages: list[Dict[str, str]],
         project_name: str = default_project,
-        tools: List[BaseTool] = None,
+        tools: list[BaseTool] = None,
+        params: dict | None = None,
     ) -> Iterator[object]:
         """
         Queries an agent to get a stream of completion chunks.
 
         Parameters:
             agent (db.Agents): Existing agent to get completion from
-            messages (List[Dict[str, str]]): Chat history to send to the agent
+            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
+            tools (list[BaseTool]): Tools to use while getting the completion
+            params (dict | None): params to redefine agent params
 
         Returns:
             chunks (Iterator[object]): Completion chunks as an iterator
@@ -659,4 +756,4 @@ class AgentsController:
         llm_params = self.get_agent_llm_params(agent.params)
 
         lang_agent = LangchainAgent(agent, model=model, llm_params=llm_params)
-        return lang_agent.get_completion(messages, stream=True)
+        return lang_agent.get_completion(messages, stream=True, params=params)

mindsdb/interfaces/agents/constants.py

@@ -26,7 +26,6 @@ OPEN_AI_CHAT_MODELS = (
 SUPPORTED_PROVIDERS = {
     "openai",
     "anthropic",
-    "anyscale",
     "litellm",
     "ollama",
     "nvidia_nim",

mindsdb/interfaces/agents/langchain_agent.py

@@ -11,7 +11,7 @@ import pandas as pd
 from langchain.agents import AgentExecutor
 from langchain.agents.initialize import initialize_agent
 from langchain.chains.conversation.memory import ConversationSummaryBufferMemory
-from langchain_community.chat_models import ChatAnyscale, ChatLiteLLM, ChatOllama
+from langchain_community.chat_models import ChatLiteLLM, ChatOllama
 from langchain_writer import ChatWriter
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.agents import AgentAction, AgentStep
@@ -165,8 +165,6 @@ def create_chat_model(args: Dict):
         except NotImplementedError:
             chat_open_ai.tiktoken_model_name = DEFAULT_TIKTOKEN_MODEL_NAME
         return chat_open_ai
-    if args["provider"] == "anyscale":
-        return ChatAnyscale(**model_kwargs)
     if args["provider"] == "litellm":
         return ChatLiteLLM(**model_kwargs)
     if args["provider"] == "ollama":
@@ -321,7 +319,7 @@ class LangchainAgent:
             self.provider,
         ]
 
-    def get_completion(self, messages, stream: bool = False):
+    def get_completion(self, messages, stream: bool = False, params: dict | None = None):
         # Get metadata and tags to be used in the trace
         metadata = self.get_metadata()
         tags = self.get_tags()
@@ -342,7 +340,9 @@ class LangchainAgent:
         if stream:
             return self._get_completion_stream(messages)
 
-        args = self.args
+        args = {}
+        args.update(self.args)
+        args.update(params or {})
 
         df = pd.DataFrame(messages)
 
@@ -598,7 +598,12 @@ AI: {response}"""
                     completions.append(result[ASSISTANT_COLUMN])
                     contexts.append(result[CONTEXT_COLUMN])
             except TimeoutError:
-                timeout_message = "I'm sorry! I couldn't come up with a response in time. Please try again."
+                timeout_message = (
+                    f"I'm sorry! I couldn't generate a response within the allotted time ({agent_timeout_seconds} seconds). "
+                    "If you need more time for processing, you can adjust the timeout settings. "
+                    "Please refer to the documentation for instructions on how to change the timeout value. "
+                    "Feel free to try your request again."
+                )
                 logger.warning(f"Agent execution timed out after {agent_timeout_seconds} seconds")
                 for _ in range(len(futures) - len(completions)):
                     completions.append(timeout_message)