From 5a8292984dae283f8be27588feb09804b4d7f09a Mon Sep 17 00:00:00 2001 From: dusvyat Date: Tue, 10 Sep 2024 17:11:43 +0300 Subject: [PATCH] Update SQL query logging method Refactor the agent utils to log SQL queries directly from chunk content instead of extracting from actions. This simplifies the process and removes redundant code related to action handling. --- mindsdb_sdk/utils/agents.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/mindsdb_sdk/utils/agents.py b/mindsdb_sdk/utils/agents.py index e35fd39..7c64265 100644 --- a/mindsdb_sdk/utils/agents.py +++ b/mindsdb_sdk/utils/agents.py @@ -79,17 +79,10 @@ def stream_and_parse_sql_query(self, completion_stream: Generator[Dict[str, Any] for message in chunk['messages']: if message.get('role') == 'assistant': self.logger.info(f"Assistant message: {message.get('content', '')}") + if chunk.get('type') == 'sql': + sql_query = chunk['content'] + self.logger.info(f"Generated SQL: {sql_query}") - if 'actions' in chunk and chunk['actions']: - self.logger.info(f"Chunk contains {len(chunk['actions'])} actions") - for action in chunk['actions']: - if 'tool' in action and 'sql_db_query' in action['tool']: - match = re.search(r'tool_input="(.*?)"', action['tool']) - if match and not sql_query_found: - sql_query = match.group(1).replace("\\'", "'") - sql_query_found = True - self.logger.info(f"SQL query found: {sql_query}") - break elif isinstance(chunk, str): output = chunk self.logger.info(f"String chunk received: {chunk}")