Skip to content

Commit

Permalink
Update SQL query logging method
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
dusvyat committed Sep 10, 2024
1 parent fbd9228 commit 5a82929
Showing 1 changed file with 3 additions and 10 deletions.
13 changes: 3 additions & 10 deletions mindsdb_sdk/utils/agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand Down

0 comments on commit 5a82929

Please sign in to comment.