solana-agent 27.4.2__py3-none-any.whl → 27.5.0__py3-none-any.whl

solana_agent/services/agent.py

@@ -9,6 +9,7 @@ import asyncio
 import datetime as main_datetime
 from datetime import datetime
 import json
+import logging  # Add logging
 from typing import AsyncGenerator, Dict, List, Literal, Optional, Any, Union
 
 from solana_agent.interfaces.services.agent import AgentService as AgentServiceInterface
@@ -16,6 +17,11 @@ from solana_agent.interfaces.providers.llm import LLMProvider
 from solana_agent.plugins.manager import PluginManager
 from solana_agent.plugins.registry import ToolRegistry
 from solana_agent.domains.agent import AIAgent, BusinessMission
+from solana_agent.interfaces.guardrails.guardrails import (
+    OutputGuardrail,
+)
+
+logger = logging.getLogger(__name__)  # Add logger
 
 
 class AgentService(AgentServiceInterface):
@@ -29,6 +35,9 @@ class AgentService(AgentServiceInterface):
         api_key: Optional[str] = None,
         base_url: Optional[str] = None,
         model: Optional[str] = None,
+        output_guardrails: List[
+            OutputGuardrail
+        ] = None,  # <-- Add output_guardrails parameter
     ):
         """Initialize the agent service.
 
@@ -36,6 +45,10 @@ class AgentService(AgentServiceInterface):
             llm_provider: Provider for language model interactions
             business_mission: Optional business mission and values
             config: Optional service configuration
+            api_key: API key for the LLM provider
+            base_url: Base URL for the LLM provider
+            model: Model name for the LLM provider
+            output_guardrails: List of output guardrail instances
         """
         self.llm_provider = llm_provider
         self.business_mission = business_mission
@@ -46,6 +59,7 @@ class AgentService(AgentServiceInterface):
         self.api_key = api_key
         self.base_url = base_url
         self.model = model
+        self.output_guardrails = output_guardrails or []  # <-- Store guardrails
 
         self.plugin_manager = PluginManager(
             config=self.config,
@@ -71,6 +85,7 @@ class AgentService(AgentServiceInterface):
             specialization=specialization,
         )
         self.agents.append(agent)
+        logger.info(f"Registered AI agent: {name}")
 
     def get_agent_system_prompt(self, agent_name: str) -> str:
         """Get the system prompt for an agent.
@@ -152,28 +167,41 @@ class AgentService(AgentServiceInterface):
         """Execute a tool on behalf of an agent."""
 
         if not self.tool_registry:
+            logger.error("Tool registry not available during tool execution.")
             return {"status": "error", "message": "Tool registry not available"}
 
         tool = self.tool_registry.get_tool(tool_name)
         if not tool:
+            logger.warning(f"Tool '{tool_name}' not found for execution.")
             return {"status": "error", "message": f"Tool '{tool_name}' not found"}
 
         # Check if agent has access to this tool
         agent_tools = self.tool_registry.get_agent_tools(agent_name)
 
         if not any(t.get("name") == tool_name for t in agent_tools):
+            logger.warning(
+                f"Agent '{agent_name}' attempted to use unassigned tool '{tool_name}'."
+            )
             return {
                 "status": "error",
                 "message": f"Agent '{agent_name}' doesn't have access to tool '{tool_name}'",
             }
 
         try:
+            logger.info(
+                f"Executing tool '{tool_name}' for agent '{agent_name}' with params: {parameters}"
+            )
             result = await tool.execute(**parameters)
+            logger.info(
+                f"Tool '{tool_name}' execution result status: {result.get('status')}"
+            )
             return result
         except Exception as e:
             import traceback
 
-            print(traceback.format_exc())
+            logger.error(
+                f"Error executing tool '{tool_name}': {e}\n{traceback.format_exc()}"
+            )
             return {"status": "error", "message": f"Error executing tool: {str(e)}"}
 
     async def generate_response(
@@ -201,10 +229,17 @@ class AgentService(AgentServiceInterface):
         ] = "aac",
         prompt: Optional[str] = None,
     ) -> AsyncGenerator[Union[str, bytes], None]:  # pragma: no cover
-        """Generate a response with support for text/audio input/output."""
+        """Generate a response with support for text/audio input/output and guardrails.
+
+        If output_format is 'text' and output_guardrails are present, the response
+        will be buffered entirely before applying guardrails and yielding a single result.
+        Otherwise, text responses stream chunk-by-chunk. Audio responses always buffer.
+        """
         agent = next((a for a in self.agents if a.name == agent_name), None)
         if not agent:
             error_msg = f"Agent '{agent_name}' not found."
+            logger.warning(error_msg)
+            # Handle error output (unchanged)
             if output_format == "audio":
                 async for chunk in self.llm_provider.tts(
                     error_msg,
@@ -217,261 +252,372 @@ class AgentService(AgentServiceInterface):
                 yield error_msg
             return
 
+        # --- Determine Buffering Strategy ---
+        # Buffer text ONLY if format is text AND guardrails are present
+        should_buffer_text = bool(self.output_guardrails) and output_format == "text"
+        logger.debug(
+            f"Text buffering strategy: {'Buffer full response' if should_buffer_text else 'Stream chunks'}"
+        )
+
         try:
-            # --- 1. Get Base System Prompt ---
+            # --- System Prompt Assembly ---
             system_prompt_parts = [self.get_agent_system_prompt(agent_name)]
 
-            # --- 2. Add Tool Usage Instructions EARLY ---
-            tool_usage_prompt_text = ""
-            if self.tool_registry:
-                tool_usage_prompt_text = self._get_tool_usage_prompt(agent_name)
-                if tool_usage_prompt_text:
-                    system_prompt_parts.append(
-                        f"\n\n--- TOOL USAGE INSTRUCTIONS ---{tool_usage_prompt_text}"
-                    )
-                    print(
-                        f"Tools available to agent {agent_name}: {[t.get('name') for t in self.get_agent_tools(agent_name)]}"
-                    )
+            # Add tool usage instructions if tools are available for the agent
+            tool_instructions = self._get_tool_usage_prompt(agent_name)
+            if tool_instructions:
+                system_prompt_parts.append(tool_instructions)
 
-            # --- 3. Add User ID ---
-            system_prompt_parts.append("\n\n--- USER & SESSION INFO ---")
-            system_prompt_parts.append(f"User ID: {user_id}")
+            # Add user ID context
+            system_prompt_parts.append(f"USER IDENTIFIER: {user_id}")
 
-            # --- 4. Add Memory Context ---
+            # Add memory context if provided
             if memory_context:
-                # Make the header clearly separate it
-                system_prompt_parts.append(
-                    "\n\n--- CONVERSATION HISTORY (Memory Context) ---"
-                )
-                system_prompt_parts.append(memory_context)
+                system_prompt_parts.append(f"\nCONVERSATION HISTORY:\n{memory_context}")
 
-            # --- 5. Add Additional Prompt (if provided) ---
+            # Add optional prompt if provided
             if prompt:
-                # Make the header clearly separate it
-                system_prompt_parts.append(
-                    "\n\n--- ADDITIONAL INSTRUCTIONS FOR THIS TURN ---"
-                )
-                system_prompt_parts.append(prompt)
+                system_prompt_parts.append(f"\nADDITIONAL PROMPT:\n{prompt}")
 
-            # --- Assemble the final system prompt ---
-            final_system_prompt = "\n".join(system_prompt_parts)
+            final_system_prompt = "\n\n".join(
+                filter(None, system_prompt_parts)
+            )  # Join non-empty parts
+            # --- End System Prompt Assembly ---
 
-            # Variables for tracking the complete response
-            complete_text_response = ""
-            full_response_buffer = ""
+            # --- Response Generation ---
+            complete_text_response = (
+                ""  # Always used for final storage and potentially for buffering
+            )
+            full_response_buffer = ""  # Used ONLY for audio buffering
 
-            # Variables for robust handling of tool call markers that may be split across chunks
+            # Tool call handling variables (unchanged)
             tool_buffer = ""
-            pending_chunk = ""  # To hold text that might contain partial markers
+            pending_chunk = ""
             is_tool_call = False
-
-            # Define start and end markers
             start_marker = "[TOOL]"
             end_marker = "[/TOOL]"
 
-            # Generate and stream response (ALWAYS use non-realtime for text generation)
-            print(f"Generating response with {len(query)} characters of query text")
+            logger.info(
+                f"Generating response for agent '{agent_name}' with query length {len(str(query))}"
+            )
             async for chunk in self.llm_provider.generate_text(
-                prompt=query,
+                prompt=str(query),
                 system_prompt=final_system_prompt,
                 api_key=self.api_key,
                 base_url=self.base_url,
                 model=self.model,
             ):
-                # If we have pending text from the previous chunk, combine it with this chunk
+                # --- Chunk Processing & Tool Call Logic (Modified Yielding) ---
                 if pending_chunk:
                     combined_chunk = pending_chunk + chunk
-                    pending_chunk = ""  # Reset pending chunk
+                    pending_chunk = ""
                 else:
                     combined_chunk = chunk
 
                 # STEP 1: Check for tool call start marker
                 if start_marker in combined_chunk and not is_tool_call:
-                    print(
-                        f"Found tool start marker in chunk of length {len(combined_chunk)}"
-                    )
                     is_tool_call = True
-
-                    # Extract text before the marker and the marker itself with everything after
                     start_pos = combined_chunk.find(start_marker)
                     before_marker = combined_chunk[:start_pos]
                     after_marker = combined_chunk[start_pos:]
 
-                    # Yield text that appeared before the marker
-                    if before_marker and output_format == "text":
-                        yield before_marker
+                    if before_marker:
+                        processed_before_marker = before_marker
+                        # Apply guardrails ONLY if NOT buffering text
+                        if not should_buffer_text:
+                            for guardrail in self.output_guardrails:
+                                try:
+                                    processed_before_marker = await guardrail.process(
+                                        processed_before_marker
+                                    )
+                                except Exception as e:
+                                    logger.error(
+                                        f"Error applying output guardrail {guardrail.__class__.__name__} to pre-tool text: {e}"
+                                    )
+
+                        # Yield ONLY if NOT buffering text
+                        if (
+                            processed_before_marker
+                            and not should_buffer_text
+                            and output_format == "text"
+                        ):
+                            yield processed_before_marker
+
+                        # Always accumulate for final response / audio buffer
+                        if processed_before_marker:
+                            complete_text_response += processed_before_marker
+                            if output_format == "audio":
+                                full_response_buffer += processed_before_marker
 
-                    # Start collecting the tool call
                     tool_buffer = after_marker
-                    continue  # Skip to next chunk
+                    continue
 
                 # STEP 2: Handle ongoing tool call collection
                 if is_tool_call:
                     tool_buffer += combined_chunk
-
-                    # Check if the tool call is complete
                     if end_marker in tool_buffer:
-                        print(f"Tool call complete, buffer size: {len(tool_buffer)}")
-
-                        # Process the tool call
                         response_text = await self._handle_tool_call(
                             agent_name=agent_name, tool_text=tool_buffer
                         )
-
-                        # Clean the response to remove any markers or formatting
                         response_text = self._clean_tool_response(response_text)
-                        print(
-                            f"Tool execution complete, result size: {len(response_text)}"
-                        )
-
-                        # Create new prompt with search/tool results
-                        # Ensure query is string
                         user_prompt = f"{str(query)}\n\nTOOL RESULT: {response_text}"
 
-                        # --- REBUILD the system prompt for the follow-up call ---
-                        # Start with base prompt again
+                        # --- Rebuild system prompt for follow-up ---
                         follow_up_system_prompt_parts = [
                             self.get_agent_system_prompt(agent_name)
                         ]
-                        # Add the instruction NOT to use tools again
+                        # Re-add tool instructions if needed for follow-up context
+                        if tool_instructions:
+                            follow_up_system_prompt_parts.append(tool_instructions)
                         follow_up_system_prompt_parts.append(
-                            "\n\nCRITICAL: You have received the results from a tool. Base your response on the 'Search Result' provided in the user prompt. DO NOT use the tool calling format again for this turn."
+                            f"USER IDENTIFIER: {user_id}"
                         )
-                        follow_up_system_prompt_parts.append(
-                            "\n\n--- USER & SESSION INFO ---"
-                        )
-                        follow_up_system_prompt_parts.append(f"User ID: {user_id}")
+                        # Include original memory + original query + tool result context
                         if memory_context:
-                            # Make the header clearly separate it
                             follow_up_system_prompt_parts.append(
-                                "\n\n--- CONVERSATION HISTORY (Memory Context) ---"
+                                f"\nORIGINAL CONVERSATION HISTORY:\n{memory_context}"
                             )
-                            follow_up_system_prompt_parts.append(memory_context)
+                        # Add the original prompt if it was provided
                         if prompt:
-                            # Make the header clearly separate it
                             follow_up_system_prompt_parts.append(
-                                "\n\n--- ADDITIONAL INSTRUCTIONS FOR THIS TURN ---"
+                                f"\nORIGINAL ADDITIONAL PROMPT:\n{prompt}"
                             )
-                            follow_up_system_prompt_parts.append(prompt)
+                        # Add context about the tool call that just happened
+                        follow_up_system_prompt_parts.append(
+                            f"\nPREVIOUS TOOL CALL CONTEXT:\nOriginal Query: {str(query)}\nTool Used: (Inferred from result)\nTool Result: {response_text}"
+                        )
 
-                        # --- Assemble the final follow_up prompt ---
-                        final_follow_up_system_prompt = "\n".join(
-                            follow_up_system_prompt_parts
+                        final_follow_up_system_prompt = "\n\n".join(
+                            filter(None, follow_up_system_prompt_parts)
                         )
-                        # --- End Rebuild ---"
-
-                        # Generate a new response with the tool results
-                        print("Generating new response with tool results")
-                        if output_format == "text":
-                            # Stream the follow-up response for text output
-                            async for (
-                                processed_chunk
-                            ) in self.llm_provider.generate_text(
-                                prompt=user_prompt,
-                                system_prompt=final_follow_up_system_prompt,
-                                api_key=self.api_key,
-                                base_url=self.base_url,
-                                model=self.model,
-                            ):
-                                complete_text_response += processed_chunk
-                                yield processed_chunk
-                        else:
-                            # For audio output, collect the full response first
-                            tool_response = ""
-                            async for (
-                                processed_chunk
-                            ) in self.llm_provider.generate_text(
-                                prompt=user_prompt,
-                                system_prompt=final_follow_up_system_prompt,
+                        # --- End Rebuild system prompt ---
+
+                        logger.info("Generating follow-up response with tool results")
+                        async for processed_chunk in self.llm_provider.generate_text(
+                            prompt=user_prompt,  # Use the prompt that includes the tool result
+                            system_prompt=final_follow_up_system_prompt,
+                            api_key=self.api_key,
+                            base_url=self.base_url,
+                            model=self.model,
+                        ):
+                            chunk_to_yield_followup = processed_chunk
+                            # Apply guardrails ONLY if NOT buffering text
+                            if not should_buffer_text:
+                                for guardrail in self.output_guardrails:
+                                    try:
+                                        chunk_to_yield_followup = (
+                                            await guardrail.process(
+                                                chunk_to_yield_followup
+                                            )
+                                        )
+                                    except Exception as e:
+                                        logger.error(
+                                            f"Error applying output guardrail {guardrail.__class__.__name__} to follow-up chunk: {e}"
+                                        )
+
+                            # Yield ONLY if NOT buffering text
+                            if (
+                                chunk_to_yield_followup
+                                and not should_buffer_text
+                                and output_format == "text"
                             ):
-                                tool_response += processed_chunk
+                                yield chunk_to_yield_followup
 
-                            # Clean and add to our complete text record and audio buffer
-                            tool_response = self._clean_for_audio(tool_response)
-                            complete_text_response += tool_response
-                            full_response_buffer += tool_response
+                            # Always accumulate
+                            if chunk_to_yield_followup:
+                                complete_text_response += chunk_to_yield_followup
+                                if output_format == "audio":
+                                    full_response_buffer += chunk_to_yield_followup
 
-                        # Reset tool handling state
                         is_tool_call = False
                         tool_buffer = ""
                         pending_chunk = ""
-                        break  # Exit the original generation loop after tool processing
+                        break  # Exit the original generation loop
 
-                    # Continue collecting tool call content without yielding
-                    continue
+                    continue  # Continue collecting tool call
 
-                # STEP 3: Check for possible partial start markers at the end of the chunk
-                # This helps detect markers split across chunks
+                # STEP 3: Check for possible partial start markers
                 potential_marker = False
+                chunk_to_yield = combined_chunk
                 for i in range(1, len(start_marker)):
                     if combined_chunk.endswith(start_marker[:i]):
-                        # Found a partial marker at the end
-                        # Save the partial marker
                         pending_chunk = combined_chunk[-i:]
-                        # Everything except the partial marker
                         chunk_to_yield = combined_chunk[:-i]
                         potential_marker = True
-                        print(f"Potential partial marker detected: '{pending_chunk}'")
                         break
 
                 if potential_marker:
-                    # Process the safe part of the chunk
-                    if chunk_to_yield and output_format == "text":
-                        yield chunk_to_yield
-                    if chunk_to_yield:
-                        complete_text_response += chunk_to_yield
+                    chunk_to_yield_safe = chunk_to_yield
+                    # Apply guardrails ONLY if NOT buffering text
+                    if not should_buffer_text:
+                        for guardrail in self.output_guardrails:
+                            try:
+                                chunk_to_yield_safe = await guardrail.process(
+                                    chunk_to_yield_safe
+                                )
+                            except Exception as e:
+                                logger.error(
+                                    f"Error applying output guardrail {guardrail.__class__.__name__} to safe chunk: {e}"
+                                )
+
+                    # Yield ONLY if NOT buffering text
+                    if (
+                        chunk_to_yield_safe
+                        and not should_buffer_text
+                        and output_format == "text"
+                    ):
+                        yield chunk_to_yield_safe
+
+                    # Always accumulate
+                    if chunk_to_yield_safe:
+                        complete_text_response += chunk_to_yield_safe
                         if output_format == "audio":
-                            full_response_buffer += chunk_to_yield
+                            full_response_buffer += chunk_to_yield_safe
                     continue
 
-                # STEP 4: Normal text processing for non-tool call content
-                if output_format == "text":
-                    yield combined_chunk
+                # STEP 4: Normal text processing
+                chunk_to_yield_normal = combined_chunk
+                # Apply guardrails ONLY if NOT buffering text
+                if not should_buffer_text:
+                    for guardrail in self.output_guardrails:
+                        try:
+                            chunk_to_yield_normal = await guardrail.process(
+                                chunk_to_yield_normal
+                            )
+                        except Exception as e:
+                            logger.error(
+                                f"Error applying output guardrail {guardrail.__class__.__name__} to normal chunk: {e}"
+                            )
+
+                # Yield ONLY if NOT buffering text
+                if (
+                    chunk_to_yield_normal
+                    and not should_buffer_text
+                    and output_format == "text"
+                ):
+                    yield chunk_to_yield_normal
+
+                # Always accumulate
+                if chunk_to_yield_normal:
+                    complete_text_response += chunk_to_yield_normal
+                    if output_format == "audio":
+                        full_response_buffer += chunk_to_yield_normal
 
-                complete_text_response += combined_chunk
-                if output_format == "audio":
-                    full_response_buffer += combined_chunk
+            # --- Post-Loop Processing ---
 
-            # Process any incomplete tool call as regular text
+            # Process any incomplete tool call
             if is_tool_call and tool_buffer:
-                print(
-                    f"Incomplete tool call detected, returning as regular text: {len(tool_buffer)} chars"
+                logger.warning(
+                    f"Incomplete tool call detected, processing as regular text: {len(tool_buffer)} chars"
                 )
-                if output_format == "text":
-                    yield tool_buffer
-
-                complete_text_response += tool_buffer
-                if output_format == "audio":
-                    full_response_buffer += tool_buffer
-
-            # For audio output, generate speech from the complete buffer
-            if output_format == "audio" and full_response_buffer:
-                # Clean text before TTS
-                print(
-                    f"Processing {len(full_response_buffer)} characters for audio output"
+                processed_tool_buffer = tool_buffer
+                # Apply guardrails ONLY if NOT buffering text
+                if not should_buffer_text:
+                    for guardrail in self.output_guardrails:
+                        try:
+                            processed_tool_buffer = await guardrail.process(
+                                processed_tool_buffer
+                            )
+                        except Exception as e:
+                            logger.error(
+                                f"Error applying output guardrail {guardrail.__class__.__name__} to incomplete tool buffer: {e}"
+                            )
+
+                # Yield ONLY if NOT buffering text
+                if (
+                    processed_tool_buffer
+                    and not should_buffer_text
+                    and output_format == "text"
+                ):
+                    yield processed_tool_buffer
+
+                # Always accumulate
+                if processed_tool_buffer:
+                    complete_text_response += processed_tool_buffer
+                    if output_format == "audio":
+                        full_response_buffer += processed_tool_buffer
+
+            # --- Final Output Generation ---
+
+            # Case 1: Text output WITH guardrails (apply to buffered response)
+            if should_buffer_text:
+                logger.info(
+                    f"Applying output guardrails to buffered text response (length: {len(complete_text_response)})"
                 )
-                full_response_buffer = self._clean_for_audio(full_response_buffer)
+                processed_full_text = complete_text_response
+                for guardrail in self.output_guardrails:
+                    try:
+                        processed_full_text = await guardrail.process(
+                            processed_full_text
+                        )
+                    except Exception as e:
+                        logger.error(
+                            f"Error applying output guardrail {guardrail.__class__.__name__} to full text buffer: {e}"
+                        )
 
-                # Process the entire response with TTS
+                if processed_full_text:
+                    yield processed_full_text
+                # Update last_text_response with the final processed text
+                self.last_text_response = processed_full_text
+
+            # Case 2: Audio output (apply guardrails to buffer before TTS) - Unchanged Logic
+            elif output_format == "audio" and full_response_buffer:
+                original_buffer = full_response_buffer
+                processed_audio_buffer = full_response_buffer
+                for (
+                    guardrail
+                ) in self.output_guardrails:  # Apply even if empty, for consistency
+                    try:
+                        processed_audio_buffer = await guardrail.process(
+                            processed_audio_buffer
+                        )
+                    except Exception as e:
+                        logger.error(
+                            f"Error applying output guardrail {guardrail.__class__.__name__} to audio buffer: {e}"
+                        )
+                if processed_audio_buffer != original_buffer:
+                    logger.info(
+                        f"Output guardrails modified audio buffer. Original length: {len(original_buffer)}, New length: {len(processed_audio_buffer)}"
+                    )
+
+                cleaned_audio_buffer = self._clean_for_audio(processed_audio_buffer)
+                logger.info(
+                    f"Processing {len(cleaned_audio_buffer)} characters for audio output"
+                )
                 async for audio_chunk in self.llm_provider.tts(
-                    text=full_response_buffer,
+                    text=cleaned_audio_buffer,
                     voice=audio_voice,
                     response_format=audio_output_format,
                     instructions=audio_instructions,
                 ):
                     yield audio_chunk
+                # Update last_text_response with the text *before* TTS cleaning
+                self.last_text_response = (
+                    processed_audio_buffer  # Store the guardrail-processed text
+                )
+
+            # Case 3: Text output WITHOUT guardrails (already streamed)
+            elif output_format == "text" and not should_buffer_text:
+                # Store the complete text response (accumulated from non-processed chunks)
+                self.last_text_response = complete_text_response
+                logger.info(
+                    "Text streaming complete (no guardrails applied post-stream)."
+                )
 
-            # Store the complete text response
-            self.last_text_response = complete_text_response
-            print(f"Response generation complete: {len(complete_text_response)} chars")
+            logger.info(
+                f"Response generation complete for agent '{agent_name}': {len(self.last_text_response)} final chars"
+            )
 
         except Exception as e:
-            error_msg = f"I apologize, but I encountered an error: {str(e)}"
-            print(f"Error in generate_response: {str(e)}")
+            # --- Error Handling (unchanged) ---
             import traceback
 
-            print(traceback.format_exc())
-
+            error_msg = (
+                "I apologize, but I encountered an error processing your request."
+            )
+            logger.error(
+                f"Error in generate_response for agent '{agent_name}': {e}\n{traceback.format_exc()}"
+            )
             if output_format == "audio":
                 async for chunk in self.llm_provider.tts(
                     error_msg,
@@ -484,35 +630,24 @@ class AgentService(AgentServiceInterface):
                 yield error_msg
 
     async def _bytes_to_generator(self, data: bytes) -> AsyncGenerator[bytes, None]:
-        """Convert bytes to an async generator for streaming.
-
-        Args:
-            data: Bytes of audio data
-
-        Yields:
-            Chunks of audio data
-        """
-        # Define a reasonable chunk size (adjust based on your needs)
+        """Convert bytes to an async generator for streaming."""
         chunk_size = 4096
-
         for i in range(0, len(data), chunk_size):
             yield data[i : i + chunk_size]
-            # Small delay to simulate streaming
             await asyncio.sleep(0.01)
 
     async def _handle_tool_call(self, agent_name: str, tool_text: str) -> str:
         """Handle marker-based tool calls."""
         try:
-            # Extract the content between markers
             start_marker = "[TOOL]"
             end_marker = "[/TOOL]"
-
             start_idx = tool_text.find(start_marker) + len(start_marker)
             end_idx = tool_text.find(end_marker)
+            if start_idx == -1 or end_idx == -1 or end_idx <= start_idx:
+                logger.error(f"Malformed tool call text received: {tool_text}")
+                return "Error: Malformed tool call format."
 
             tool_content = tool_text[start_idx:end_idx].strip()
-
-            # Parse the lines to extract name and parameters
             tool_name = None
             parameters = {}
 
@@ -520,43 +655,61 @@ class AgentService(AgentServiceInterface):
                 line = line.strip()
                 if not line:
                     continue
-
                 if line.startswith("name:"):
                     tool_name = line[5:].strip()
                 elif line.startswith("parameters:"):
                     params_text = line[11:].strip()
-                    # Parse comma-separated parameters
-                    param_pairs = params_text.split(",")
-                    for pair in param_pairs:
-                        if "=" in pair:
-                            k, v = pair.split("=", 1)
-                            parameters[k.strip()] = v.strip()
-
-            # Execute the tool
+                    try:
+                        # Attempt to parse as JSON first for robustness
+                        parameters = json.loads(params_text)
+                    except json.JSONDecodeError:
+                        # Fallback to comma-separated key=value pairs
+                        param_pairs = params_text.split(",")
+                        for pair in param_pairs:
+                            if "=" in pair:
+                                k, v = pair.split("=", 1)
+                                parameters[k.strip()] = v.strip()
+                        logger.warning(
+                            f"Parsed tool parameters using fallback method: {params_text}"
+                        )
+
+            if not tool_name:
+                logger.error(f"Tool name missing in tool call: {tool_content}")
+                return "Error: Tool name missing in call."
+
             result = await self.execute_tool(agent_name, tool_name, parameters)
 
-            # Return the result as string
             if result.get("status") == "success":
                 tool_result = str(result.get("result", ""))
                 return tool_result
             else:
                 error_msg = f"Error calling {tool_name}: {result.get('message', 'Unknown error')}"
+                logger.error(error_msg)
                 return error_msg
 
         except Exception as e:
             import traceback
 
-            print(traceback.format_exc())
+            logger.error(f"Error processing tool call: {e}\n{traceback.format_exc()}")
             return f"Error processing tool call: {str(e)}"
 
     def _get_tool_usage_prompt(self, agent_name: str) -> str:
         """Generate marker-based instructions for tool usage."""
-        # Get tools assigned to this agent
         tools = self.get_agent_tools(agent_name)
         if not tools:
             return ""
 
-        tools_json = json.dumps(tools, indent=2)
+        # Simplify tool representation for the prompt
+        simplified_tools = []
+        for tool in tools:
+            simplified_tool = {
+                "name": tool.get("name"),
+                "description": tool.get("description"),
+                "parameters": tool.get("parameters", {}).get("properties", {}),
+            }
+            simplified_tools.append(simplified_tool)
+
+        tools_json = json.dumps(simplified_tools, indent=2)
 
         return f"""
         AVAILABLE TOOLS:
@@ -569,21 +722,21 @@ class AgentService(AgentServiceInterface):
         TOOL USAGE FORMAT:
         [TOOL]
         name: tool_name
-        parameters: key1=value1, key2=value2
+        parameters: {{"key1": "value1", "key2": "value2"}}
         [/TOOL]
 
         EXAMPLES:
         ✅ CORRECT - ONLY the tool call with NOTHING else:
         [TOOL]
         name: search_internet
-        parameters: query=latest news on Solana
+        parameters: {{"query": "latest news on Solana"}}
         [/TOOL]
-        
+
         ❌ INCORRECT - Never add explanatory text like this:
         To get the latest news on Solana, I will search the internet.
         [TOOL]
         name: search_internet
-        parameters: query=latest news on Solana
+        parameters: {{"query": "latest news on Solana"}}
         [/TOOL]
 
         REMEMBER:
@@ -594,62 +747,33 @@ class AgentService(AgentServiceInterface):
         """
 
     def _clean_for_audio(self, text: str) -> str:
-        """Remove Markdown formatting, emojis, and non-pronounceable characters from text.
-
-        Args:
-            text: Input text with potential Markdown formatting and special characters
-
-        Returns:
-            Clean text without Markdown, emojis, and special characters
-        """
+        """Remove Markdown formatting, emojis, and non-pronounceable characters from text."""
         import re
 
         if not text:
             return ""
-
-        # Remove Markdown links - [text](url) -> text
         text = re.sub(r"\[([^\]]+)\]\([^\)]+\)", r"\1", text)
-
-        # Remove inline code with backticks
         text = re.sub(r"`([^`]+)`", r"\1", text)
-
-        # Remove bold formatting - **text** or __text__ -> text
         text = re.sub(r"(\*\*|__)(.*?)\1", r"\2", text)
-
-        # Remove italic formatting - *text* or _text_ -> text
         text = re.sub(r"(\*|_)(.*?)\1", r"\2", text)
-
-        # Remove headers - ## Header -> Header
         text = re.sub(r"^\s*#+\s*(.*?)$", r"\1", text, flags=re.MULTILINE)
-
-        # Remove blockquotes - > Text -> Text
         text = re.sub(r"^\s*>\s*(.*?)$", r"\1", text, flags=re.MULTILINE)
-
-        # Remove horizontal rules (---, ***, ___)
         text = re.sub(r"^\s*[-*_]{3,}\s*$", "", text, flags=re.MULTILINE)
-
-        # Remove list markers - * Item or - Item or 1. Item -> Item
         text = re.sub(r"^\s*[-*+]\s+(.*?)$", r"\1", text, flags=re.MULTILINE)
         text = re.sub(r"^\s*\d+\.\s+(.*?)$", r"\1", text, flags=re.MULTILINE)
-
-        # Remove multiple consecutive newlines (keep just one)
         text = re.sub(r"\n{3,}", "\n\n", text)
-
-        # Remove emojis and other non-pronounceable characters
-        # Common emoji Unicode ranges
         emoji_pattern = re.compile(
             "["
             "\U0001f600-\U0001f64f"  # emoticons
             "\U0001f300-\U0001f5ff"  # symbols & pictographs
             "\U0001f680-\U0001f6ff"  # transport & map symbols
             "\U0001f700-\U0001f77f"  # alchemical symbols
-            "\U0001f780-\U0001f7ff"  # Geometric Shapes
+            "\U0001f780-\U0001f7ff"  # Geometric Shapes Extended
             "\U0001f800-\U0001f8ff"  # Supplemental Arrows-C
             "\U0001f900-\U0001f9ff"  # Supplemental Symbols and Pictographs
-            "\U0001fa00-\U0001fa6f"  # Chess Symbols
             "\U0001fa70-\U0001faff"  # Symbols and Pictographs Extended-A
             "\U00002702-\U000027b0"  # Dingbats
-            "\U000024c2-\U0000257f"  # Enclosed characters
+            "\U000024c2-\U0001f251"
             "\U00002600-\U000026ff"  # Miscellaneous Symbols
             "\U00002700-\U000027bf"  # Dingbats
             "\U0000fe00-\U0000fe0f"  # Variation Selectors
@@ -658,26 +782,57 @@ class AgentService(AgentServiceInterface):
             flags=re.UNICODE,
         )
         text = emoji_pattern.sub(r" ", text)
-
-        # Replace special characters that can cause issues with TTS
-        text = re.sub(r"[^\w\s\.\,\;\:\?\!\'\"\-\(\)]", " ", text)
-
-        # Replace multiple spaces with a single space
+        text = re.sub(
+            r"[^\w\s\.\,\;\:\?\!\'\"\-\(\)]", " ", text
+        )  # Keep basic punctuation
         text = re.sub(r"\s+", " ", text)
-
         return text.strip()
 
     def _clean_tool_response(self, text: str) -> str:
         """Remove any tool markers or formatting that might have leaked into the response."""
         if not text:
             return ""
+        text = text.replace("[TOOL]", "").replace("[/TOOL]", "")
+        if text.lstrip().startswith("TOOL"):
+            text = text.lstrip()[4:].lstrip()  # Remove "TOOL" and leading space
+        return text.strip()
 
-        # Remove any tool markers that might be in the response
-        text = text.replace("[TOOL]", "")
-        text = text.replace("[/TOOL]", "")
+    # --- Add methods from factory logic ---
+    def load_and_register_plugins(self):
+        """Loads plugins using the PluginManager."""
+        try:
+            self.plugin_manager.load_plugins()
+            logger.info("Plugins loaded successfully via PluginManager.")
+        except Exception as e:
+            logger.error(f"Error loading plugins: {e}", exc_info=True)
 
-        # Remove the word TOOL from start if it appears
-        if text.lstrip().startswith("TOOL"):
-            text = text.lstrip().replace("TOOL", "", 1)
+    def register_agents_from_config(self):
+        """Registers agents defined in the main configuration."""
+        agents_config = self.config.get("agents", [])
+        if not agents_config:
+            logger.warning("No agents defined in the configuration.")
+            return
 
-        return text.strip()
+        for agent_config in agents_config:
+            name = agent_config.get("name")
+            instructions = agent_config.get("instructions")
+            specialization = agent_config.get("specialization")
+            tools = agent_config.get("tools", [])
+
+            if not name or not instructions or not specialization:
+                logger.warning(
+                    f"Skipping agent due to missing name, instructions, or specialization: {agent_config}"
+                )
+                continue
+
+            self.register_ai_agent(name, instructions, specialization)
+            # logger.info(f"Registered agent: {name}") # Logging done in register_ai_agent
+
+            # Assign tools to the agent
+            for tool_name in tools:
+                if self.assign_tool_for_agent(name, tool_name):
+                    logger.info(f"Assigned tool '{tool_name}' to agent '{name}'.")
+                else:
+                    logger.warning(
+                        f"Failed to assign tool '{tool_name}' to agent '{name}' (Tool might not be registered)."
+                    )