diff --git a/mcp_server_snowflake/cortex_services/tools.py b/mcp_server_snowflake/cortex_services/tools.py index bdbaa8b..c9d98a9 100644 --- a/mcp_server_snowflake/cortex_services/tools.py +++ b/mcp_server_snowflake/cortex_services/tools.py @@ -76,7 +76,6 @@ async def query_cortex_agent( payload = { "messages": [{"role": "user", "content": [{"type": "text", "text": query}]}], "tool_choice": {"type": "auto"}, - "stream": False, # Ignored by Agent API } try: response = requests.post( diff --git a/mcp_server_snowflake/utils.py b/mcp_server_snowflake/utils.py index 297a0ee..94d7a53 100644 --- a/mcp_server_snowflake/utils.py +++ b/mcp_server_snowflake/utils.py @@ -309,11 +309,11 @@ def parse_agent_response(self, response_stream: requests.Response) -> str: # If the next line is 'data:', it contains the JSON payload we need json_data = line[len("data: ") :] try: - final_text = ( - json.loads(json_data) - .get("content", [{}])[-1] - .get("text", "No final response found.") - ) + content = json.loads(json_data).get("content", []) + final_text = "No final response found." + for item in content: + if item.get("type") == "text" and "text" in item: + final_text = item["text"] # Return formatted AgentResponse for consistency with other parsers response = AgentResponse(results=final_text) return response.model_dump_json()