Spaces:

chrizefan
/

digichat

Runtime error

digichat / agents.py

Upload folder using huggingface_hub

fe52ef9 verified over 1 year ago

24.8 kB

from typing import List, Dict, Any, Optional, Generator, Union
from client_utils import get_client, get_model
from agent_config import AgentConfig
import time
import logging
import json

# Update import for Azure Search
from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient
# Remove the problematic import
# from azure.search.documents.models import Vector
from embedding_client import EmbeddingClient

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,  # Changed from INFO to DEBUG
    format='%(asctime%s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class BaseAgent:
    """Base class for all agents"""
    def __init__(self, config: AgentConfig):
        self.config = config
        # Get model info from config
        self.model_config = config.get_model_config()
        
        # Extract the engine name from model_config (for client lookup)
        self.engine_name = self.model_config.get("model", "grok-2-latest")
        
        # Get the actual model name for API calls
        self.model_name = get_model(self.engine_name)
        
        # Set temperature from config
        self.temperature = self.model_config.get("temperature", 0.0)
        
        # Get the appropriate client for this engine
        self.client = get_client(self.engine_name)
        
        logger.info(f"Initialized {self.__class__.__name__} with engine {self.engine_name}, model {self.model_name}")

    def process(self, messages: List[Dict[str, Any]]) -> str:
        """Process messages and return a response"""
        logger.debug(f"Processing messages: {json.dumps(messages, indent=2)}")
        # Use system prompt from config if defined
        try:
            system_prompt = self.config.get_system_prompt(self.__class__.__name__.lower())
        except Exception:
            try:
                system_prompt = self.config.get_system_prompt("default")
            except Exception:
                system_prompt = None

        # Only prepend the system prompt if not already present
        if system_prompt and (not messages or messages[0].get("role") != "system"):
            messages = [{"role": "system", "content": system_prompt}] + messages

        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=self.temperature
        )
        logger.debug(f"Received completion: {completion}")
        return completion

    def process_stream(self, messages: List[Dict[str, Any]]) -> Generator[Any, None, None]:
        """Process messages and stream the response"""
        logger.debug(f"Streaming messages: {json.dumps(messages, indent=2)}")
        # Use system prompt from config if defined
        try:
            system_prompt = self.config.get_system_prompt(self.__class__.__name__.lower())
        except Exception:
            try:
                system_prompt = self.config.get_system_prompt("default")
            except Exception:
                system_prompt = None

        # Only prepend the system prompt if not already present
        if system_prompt and (not messages or messages[0].get("role") != "system"):
            messages = [{"role": "system", "content": system_prompt}] + messages

        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=self.temperature,
            stream=True
        )
        for chunk in completion:
            logger.debug(f"Streaming chunk: {chunk}")
            return chunk

class DataAgent(BaseAgent):
    """Agent specialized in processing queries with different data sources"""
    def __init__(self, config: AgentConfig):
        super().__init__(config)
        self.tool_name = None  # Track which tool is being executed
        self.data_sources = config.get_data_sources()  # Get all data sources

    def get_extra_body(self) -> Optional[Dict[str, Any]]:
        """Get the extra body configuration for the current search type"""
        # Dynamically match the data source based on the current tool name
        for data_source_name, data_source_config in self.data_sources.items():
            if data_source_name in self.tool_name:
                return {"data_sources": [data_source_config]}
        
        # If no matching data source is found, return None
        return None

    def process(self, messages: List[Dict[str, Any]], tool_name: str = None, **kwargs) -> str:
        """Process messages and return a response with the specified tool"""
        self.tool_name = tool_name
        logger.info(f"Processing data agent for tool {tool_name}")
        extra_body = self.get_extra_body()
        # Pass num_results or other kwargs into extra_body if needed
        if extra_body is not None and kwargs:
            extra_body.update({k: v for k, v in kwargs.items() if k in extra_body})
        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=self.temperature,
            extra_body=extra_body
        )
        logger.debug(f"Data source completion: {completion}")
        return completion

    def process_stream(self, messages: List[Dict[str, Any]], tool_name: str = None, **kwargs) -> Generator[Any, None, None]:
        """Process messages and stream the response with the specified tool"""
        self.tool_name = tool_name
        logger.info(f"Streaming data agent for tool {tool_name}")
        extra_body = self.get_extra_body()
        # Pass num_results or other kwargs into extra_body if needed
        if extra_body is not None and kwargs:
            extra_body.update({k: v for k, v in kwargs.items() if k in extra_body})
        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=self.temperature,
            stream=True,
            extra_body=extra_body
        )
        for chunk in completion:
            yield chunk

class SearchAgent(BaseAgent):
    """Agent specialized in querying Azure AI Search and Azure Vector Index"""
    def __init__(self, config: AgentConfig):
        super().__init__(config)
        # Extract Azure-specific configurations from config
        self.azure_index = config.get_azure_index()
        self.search_endpoint = self.azure_index.get("search_endpoint")
        self.search_key = self.azure_index.get("search_key")
        self.search_index_name = self.azure_index.get("search_index_name")
        self.vector_index_name = self.azure_index.get("vector_index_name")
        self.embeddings_deployment = self.azure_index.get("embeddings_deployment")
        
        # Initialize the embedding client
        self.embedding_client = EmbeddingClient(
            azure_endpoint=self.azure_index.get("azure_openai_endpoint"),
            api_key=self.azure_index.get("azure_openai_key"),
            deployment=self.embeddings_deployment
        )
        
        # Initialize search clients
        self.search_client = self._create_search_client(self.search_index_name)
        self.vector_client = self._create_search_client(self.vector_index_name)
        
        logger.info(f"Initialized Search Agent with endpoint {self.azure_index.get('search_endpoint')}")
    
    def _create_search_client(self, index_name: str) -> Optional[SearchClient]:
        """Create a search client for the specified index"""
        if not self.search_endpoint or not self.search_key or not index_name:
            logger.warning(f"Missing configuration for search client: endpoint={bool(self.search_endpoint)}, key={bool(self.search_key)}, index={bool(index_name)}")
            return None
            
        try:
            credential = AzureKeyCredential(self.search_key)
            client = SearchClient(
                endpoint=self.search_endpoint,
                index_name=index_name,
                credential=credential
            )
            logger.info(f"Successfully created search client for index {index_name}")
            return client
        except Exception as e:
            logger.error(f"Failed to create search client for index {index_name}: {e}")
            return None
    
    def _generate_embedding(self, text: str) -> List[float]:
        """Generate embedding for vector search using Azure OpenAI"""
        try:
            # Use the embedding client to generate embeddings
            return self.embedding_client.get_embedding(text)
        except Exception as e:
            logger.error(f"Error generating embedding: {e}")
            # Return a fallback embedding (zeros would be neutral in vector space)
            import numpy as np
            embedding = np.zeros(1536)  # Standard dimension for embeddings
            return embedding.tolist()
    
    def query_index(self, query: str, search_type: str, **kwargs) -> Dict[str, Any]:
        """Query the Azure Search index based on search type"""
        
        search_text = kwargs.get("search_text", query)
        logger.info(f"Executing {search_type} search: {search_text}")
        
        if search_type == "vector":
            if not self.vector_client:
                return {"error": "Vector search client not configured"}
            client = self.vector_client
            #embedding = self._generate_embedding(query)
            try:
                results = list(client.search(
                    search_text=search_text,
                    top=kwargs.get("num_results"),
                    include_total_count=True,
                    filter=kwargs.get("filter"),
                    order_by=kwargs.get("orderby"),
                    select=kwargs.get("select"),
                    vector_queries=[
                        {
                            "text": search_text,
                            "fields": kwargs.get("vector_field", "text_vector"),
                            "k": kwargs.get("num_results"),
                            "kind": "text", 
                        }
                    ],
                ))
                
            except Exception as e:
                results = [f"Vector search failed: {e}"]
                logger.info(results)
        else:
            if not self.search_client:
                return {"error": "Index search client not configured"}
            client = self.search_client

            fields = [
                "Id", "Owner", "Name", "CreatedBy", "LastModifiedBy", "ParentId",
                "CreatedDataTime", "LastModifiedDate",
                "Discriminator", "SourceLocationType", "metadata_storage_content_type",
                "metadata_storage_size", "metadata_storage_last_modified",
                "metadata_storage_content_md5", "metadata_storage_name",
                "metadata_storage_path", "metadata_storage_file_extension",
            ]
            
            try:
                raw_results = list(client.search(
                    search_text=search_text,
                    query_type="full",
                    filter=kwargs.get("filter"),
                    order_by=kwargs.get("orderby"),
                    top=kwargs.get("num_results"),
                    include_total_count=True,
                    select=kwargs.get("select"),
                ))
                # Remove key-value pairs with null values from each result dict
                results = [
                    {k: v for k, v in doc.items() if v is not None}
                    for doc in raw_results
                ]
            except Exception as e:
                results = [f"Index search failed: {e}"]
                logger.info(results)
                
        logger.info(f"Search returned {len(results)} results")
        return results
    
    def format_results(self, results: Dict[str, Any]) -> str:
        """Format the search results into readable text"""
        if "error" in results:
            return f"Error: {results['error']}"
        
        if not results or not results.get("value") or len(results["value"]) == 0:
            return "No results found for your query."
            
        docs = results["value"]
        response_lines = [f"Found {len(docs)} relevant documents:"]
        
        for i, doc in enumerate(docs, 1):  # Limit to top 5 for readability
            title = doc.get("metadata_title", doc.get("title", f"Document {i}"))
            content = doc.get("content", "")
            author = doc.get("metadata_author", "Unknown")
            score = doc.get("@search.score", 0)
            url = doc.get("url", "")
            
            response_lines.append(f"\n### {i}. {title}")
            response_lines.append(f"Author: {author} | Score: {score:.2f}")
            
            if url:
                response_lines.append(f"URL: {url}")
                
            # Get a snippet from content (first 150 chars)
            if content:
                snippet = content
                response_lines.append(f"\nPreview: {snippet}")
                
            response_lines.append("-" * 40)
            
        return "\n".join(response_lines)

    def process(self, messages: List[Dict[str, Any]], **kwargs) -> List[Dict[str, Any]]:
        """Process messages by directly querying Azure Search index"""
        # Extract the query from the last user message
        query = ""
        for message in reversed(messages):
            if message["role"] == "user":
                query = message["content"]
                break

        if not query:
            return "No query provided"

        # Determine search type based on tool name or kwargs
        tool_name = kwargs.get("tool_name", "")
        # Use explicit search_mode if provided, else fallback to tool_name logic
        search_mode = kwargs.get("search_mode")
        if search_mode in ("vector", "index"):
            search_type = search_mode
        elif "vector" in tool_name or "content" in tool_name:
            search_type = "vector"
        else:
            search_type = "index"

        results = self.query_index(query, search_type, **kwargs)
        return results

    def process_stream(self, messages: List[Dict[str, Any]], **kwargs) -> Generator[Any, None, None]:
        for chunk in self.process(messages, **kwargs):
            yield chunk

class OrchestratorAgent(BaseAgent):
    """Agent that coordinates between specialized agents using them as tools"""
    def __init__(self, config: AgentConfig):
        # Initialize the base agent with the given config
        super().__init__(config)
        
        self.tools = []
        self.data_agent = None
        self.search_agent = None
        
        # Set up specialized agents if tools are available
        self.tools = config.get_tools()
        self.has_tools = bool(self.tools)
    
        # Initialize the data agent for data source queries
        self.data_agent = DataAgent(config)
        
        # Initialize the search agent for Azure search queries
        self.search_agent = SearchAgent(config)
        
        logger.debug(f"Available tools: {json.dumps(self.tools, indent=2)}")


    def _execute_tool(self, tool_name: str, tool_args: Dict[str, Any], stream: bool = False, messages: List[Dict[str, Any]] = None) -> Union[str, Generator[Any, None, None]]:
        """Execute a tool and return its result"""
        if not self.has_tools:
            error_msg = "No tools available for this engine"
            logger.error(error_msg)
            raise ValueError(error_msg)

        logger.info(f"Executing tool: {tool_name} with arguments: {json.dumps(tool_args, indent=2)}")

        # Use the full message history, appending the tool call as the last user message
        if messages is None:
            messages = []
        if self.config.get_system_prompt(tool_name):
            messages = [{"role": "system", "content": self.config.get_system_prompt(tool_name)}] + messages
            
        tool_agent_map = self.config.get_tool_agent_map()
        agent_type = tool_agent_map.get(tool_name)

        if agent_type == "search_agent":
            if stream:
                return self.search_agent.process_stream(messages, tool_name=tool_name, **tool_args)
            else:
                return self.search_agent.process(messages, tool_name=tool_name, **tool_args)
        elif agent_type == "data_agent":
            if stream:
                return self.data_agent.process_stream(messages, tool_name=tool_name, **tool_args)
            else:
                return self.data_agent.process(messages, tool_name=tool_name, **tool_args)
        elif agent_type == "base_agent":
            if stream:
                return self.process_stream(messages)
            else:
                return self.process(messages)
        else:
            logger.error(f"Unknown tool name: {tool_name}")
            return f"Error: Unknown tool name: {tool_name}"

    def _handle_tool_calls_recursive(self, messages: List[Dict[str, Any]], max_iterations: int = 10) -> Generator[Any, None, None]:
        """Recursively handle tool calls until we get a response without tools"""
        if max_iterations <= 0 or not self.has_tools:
            logger.warning("Max iterations reached or no tools available, stopping recursive tool calls")
            return

        logger.info("Getting model's response with tool calls")
        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            tools=self.tools,
            tool_choice="auto",
            temperature=self.temperature,
            stream=True
        )

        response_text = ""
        tool_calls = []

        for chunk in completion:
            try:
                if chunk.choices and chunk.choices[0].delta:
                    delta = chunk.choices[0].delta
                    if hasattr(delta, 'content') and delta.content:
                        response_text += delta.content
                        logger.debug(f"Assistant content chunk: {delta.content}")
                    if hasattr(delta, 'tool_calls') and delta.tool_calls:
                        for tool_call in delta.tool_calls:
                            if tool_call.index is not None:
                                while len(tool_calls) <= tool_call.index:
                                    tool_calls.append({})
                                if tool_calls[tool_call.index] == {}:
                                    tool_calls[tool_call.index]["index"] = tool_call.index
                                    tool_calls[tool_call.index]["function"] = {"name": "", "arguments": ""}
                                if tool_call.id:
                                    tool_calls[tool_call.index]["id"] = tool_call.id
                                if tool_call.type:
                                    tool_calls[tool_call.index]["type"] = tool_call.type
                                if tool_call.function.name:
                                    tool_calls[tool_call.index]["function"]["name"] = tool_call.function.name
                                if tool_call.function.arguments:
                                    tool_calls[tool_call.index]["function"]["arguments"] += tool_call.function.arguments
                                    logger.debug(f"Tool call arguments chunk: {tool_call.function.arguments}")
                    yield {
                        "role": "assistant",
                        "content": response_text,
                    }
            except Exception as e:
                logger.error(f"Error processing chunk in handle_tool_calls: {e}")
                continue

        if not tool_calls:
            logger.info("No tool calls made, ending process")
            return
        else:
            logger.info(f"Tool calls detected: {json.dumps(tool_calls, indent=2)}")
            for tool_call in tool_calls:
                yield {
                    "role": "assistant",
                    "content": "",
                    "metadata": {
                        "title": tool_call["function"]["name"],
                        "id": tool_call["id"],
                    }
                }
                
        assistant_message = {
            "role": "assistant",
            "content": response_text,
            "tool_calls": tool_calls,
        }

        # Now execute tools and collect tool messages
        tool_messages = []
        for tool_call in tool_calls:
            tool_name = tool_call["function"]["name"]
            try:
                tool_args = eval(tool_call["function"]["arguments"])
                if not isinstance(tool_args, dict):
                    raise ValueError("Tool arguments must be a dictionary")
                logger.debug(f"Parsed tool arguments for {tool_name}: {json.dumps(tool_args, indent=2)}")
            except (SyntaxError, ValueError) as e:
                logger.error(f"Error parsing tool arguments: {e}")
                continue

            response_stream = self._execute_tool(tool_name, tool_args, stream=True, messages=messages)
            response_list = []
            response_text = ""
            
            tool_message = {
                "role": "tool",
                "tool_call_id": tool_call["id"],
                "content": "",
            }
            # Capitalize tool_name and replace underscores with spaces for title
            title = tool_name.replace("_", " ").title()
            for chunk in response_stream:
                try:
                    if hasattr(chunk, 'choices') and chunk.choices and hasattr(chunk.choices[0], 'delta') and hasattr(chunk.choices[0].delta, 'content'):
                        response_text += chunk.choices[0].delta.content or ""
                    elif isinstance(chunk, dict) or isinstance(chunk, list):
                        response_list.append(chunk)
                        response_text += json.dumps(chunk, indent=2)
                    elif isinstance(chunk, str):
                        response_text += chunk
                        
                    tool_message = {
                            "role": "tool",
                            "tool_call_id": tool_call["id"],
                            "content": response_text
                        }
                    yield {
                        "role": "assistant",
                        "content": json.dumps(response_list[:10], indent=2) if response_list else response_text,
                        "metadata": {
                            "title": title,
                            "id": tool_call["id"],
                            "status": "pending",
                        }
                    }
                    
                except Exception as e:
                    response_text = f"Error processing tool response: {e}"
                    logger.error(f"Error processing tool response chunk: {e}")
                    break
                
            yield {
                "role": "assistant",
                "content": json.dumps(response_list[:10], indent=2) if response_list else response_text if response_text else "No results found.",
                "metadata": {
                    "title": title,
                    "id": tool_call["id"],
                    "log": f"(results: {len(response_list)})" if response_list else "",
                    "status": "done",
                }
            }
            
            tool_messages.append(tool_message)
            logger.info(f"Completed tool execution: {tool_name}")

        messages.append(assistant_message)
        messages.extend(tool_messages)

        yield from self._handle_tool_calls_recursive(messages, max_iterations - 1)

    def process_stream(self, messages: List[Dict[str, Any]]) -> Generator[Any, None, None]:
        """Process messages and stream the response"""
        logger.info("Starting orchestration process")
        logger.debug(f"Input messages: {json.dumps(messages, indent=2)}")
        
        # If no tools are available, fall back to the base agent behavior
        if not self.has_tools:
            logger.info("No tools available, falling back to standard model completion")
            yield from super().process_stream(messages)
            return
        
        # Add the system prompt to the messages only if not already present
        system_prompt = self.config.get_system_prompt("orchestrator")
        if not messages or messages[0].get("role") != "system":
            messages = [{"role": "system", "content": system_prompt}, *messages]
        
        # Start the recursive tool call handling
        yield from self._handle_tool_calls_recursive(messages)

	from typing import List, Dict, Any, Optional, Generator, Union
	from client_utils import get_client, get_model
	from agent_config import AgentConfig
	import time
	import logging
	import json

	# Update import for Azure Search
	from azure.core.credentials import AzureKeyCredential
	from azure.search.documents import SearchClient
	# Remove the problematic import
	# from azure.search.documents.models import Vector
	from embedding_client import EmbeddingClient

	# Configure logging
	logging.basicConfig(
	level=logging.DEBUG, # Changed from INFO to DEBUG
	format='%(asctime%s - %(name)s - %(levelname)s - %(message)s'
	)
	logger = logging.getLogger(__name__)

	class BaseAgent:
	"""Base class for all agents"""
	def __init__(self, config: AgentConfig):
	self.config = config
	# Get model info from config
	self.model_config = config.get_model_config()

	# Extract the engine name from model_config (for client lookup)
	self.engine_name = self.model_config.get("model", "grok-2-latest")

	# Get the actual model name for API calls
	self.model_name = get_model(self.engine_name)

	# Set temperature from config
	self.temperature = self.model_config.get("temperature", 0.0)

	# Get the appropriate client for this engine
	self.client = get_client(self.engine_name)

	logger.info(f"Initialized {self.__class__.__name__} with engine {self.engine_name}, model {self.model_name}")

	def process(self, messages: List[Dict[str, Any]]) -> str:
	"""Process messages and return a response"""
	logger.debug(f"Processing messages: {json.dumps(messages, indent=2)}")
	# Use system prompt from config if defined
	try:
	system_prompt = self.config.get_system_prompt(self.__class__.__name__.lower())
	except Exception:
	try:
	system_prompt = self.config.get_system_prompt("default")
	except Exception:
	system_prompt = None

	# Only prepend the system prompt if not already present
	if system_prompt and (not messages or messages[0].get("role") != "system"):
	messages = [{"role": "system", "content": system_prompt}] + messages

	completion = self.client.chat.completions.create(
	model=self.model_name,
	messages=messages,
	temperature=self.temperature
	)
	logger.debug(f"Received completion: {completion}")
	return completion

	def process_stream(self, messages: List[Dict[str, Any]]) -> Generator[Any, None, None]:
	"""Process messages and stream the response"""
	logger.debug(f"Streaming messages: {json.dumps(messages, indent=2)}")
	# Use system prompt from config if defined
	try:
	system_prompt = self.config.get_system_prompt(self.__class__.__name__.lower())
	except Exception:
	try:
	system_prompt = self.config.get_system_prompt("default")
	except Exception:
	system_prompt = None

	# Only prepend the system prompt if not already present
	if system_prompt and (not messages or messages[0].get("role") != "system"):
	messages = [{"role": "system", "content": system_prompt}] + messages

	completion = self.client.chat.completions.create(
	model=self.model_name,
	messages=messages,
	temperature=self.temperature,
	stream=True
	)
	for chunk in completion:
	logger.debug(f"Streaming chunk: {chunk}")
	return chunk

	class DataAgent(BaseAgent):
	"""Agent specialized in processing queries with different data sources"""
	def __init__(self, config: AgentConfig):
	super().__init__(config)
	self.tool_name = None # Track which tool is being executed
	self.data_sources = config.get_data_sources() # Get all data sources

	def get_extra_body(self) -> Optional[Dict[str, Any]]:
	"""Get the extra body configuration for the current search type"""
	# Dynamically match the data source based on the current tool name
	for data_source_name, data_source_config in self.data_sources.items():
	if data_source_name in self.tool_name:
	return {"data_sources": [data_source_config]}

	# If no matching data source is found, return None
	return None

	def process(self, messages: List[Dict[str, Any]], tool_name: str = None, **kwargs) -> str:
	"""Process messages and return a response with the specified tool"""
	self.tool_name = tool_name
	logger.info(f"Processing data agent for tool {tool_name}")
	extra_body = self.get_extra_body()
	# Pass num_results or other kwargs into extra_body if needed
	if extra_body is not None and kwargs:
	extra_body.update({k: v for k, v in kwargs.items() if k in extra_body})
	completion = self.client.chat.completions.create(
	model=self.model_name,
	messages=messages,
	temperature=self.temperature,
	extra_body=extra_body
	)
	logger.debug(f"Data source completion: {completion}")
	return completion

	def process_stream(self, messages: List[Dict[str, Any]], tool_name: str = None, **kwargs) -> Generator[Any, None, None]:
	"""Process messages and stream the response with the specified tool"""
	self.tool_name = tool_name
	logger.info(f"Streaming data agent for tool {tool_name}")
	extra_body = self.get_extra_body()
	# Pass num_results or other kwargs into extra_body if needed
	if extra_body is not None and kwargs:
	extra_body.update({k: v for k, v in kwargs.items() if k in extra_body})
	completion = self.client.chat.completions.create(
	model=self.model_name,
	messages=messages,
	temperature=self.temperature,
	stream=True,
	extra_body=extra_body
	)
	for chunk in completion:
	yield chunk

	class SearchAgent(BaseAgent):
	"""Agent specialized in querying Azure AI Search and Azure Vector Index"""
	def __init__(self, config: AgentConfig):
	super().__init__(config)
	# Extract Azure-specific configurations from config
	self.azure_index = config.get_azure_index()
	self.search_endpoint = self.azure_index.get("search_endpoint")
	self.search_key = self.azure_index.get("search_key")
	self.search_index_name = self.azure_index.get("search_index_name")
	self.vector_index_name = self.azure_index.get("vector_index_name")
	self.embeddings_deployment = self.azure_index.get("embeddings_deployment")

	# Initialize the embedding client
	self.embedding_client = EmbeddingClient(
	azure_endpoint=self.azure_index.get("azure_openai_endpoint"),
	api_key=self.azure_index.get("azure_openai_key"),
	deployment=self.embeddings_deployment
	)

	# Initialize search clients
	self.search_client = self._create_search_client(self.search_index_name)
	self.vector_client = self._create_search_client(self.vector_index_name)

	logger.info(f"Initialized Search Agent with endpoint {self.azure_index.get('search_endpoint')}")

	def _create_search_client(self, index_name: str) -> Optional[SearchClient]:
	"""Create a search client for the specified index"""
	if not self.search_endpoint or not self.search_key or not index_name:
	logger.warning(f"Missing configuration for search client: endpoint={bool(self.search_endpoint)}, key={bool(self.search_key)}, index={bool(index_name)}")
	return None

	try:
	credential = AzureKeyCredential(self.search_key)
	client = SearchClient(
	endpoint=self.search_endpoint,
	index_name=index_name,
	credential=credential
	)
	logger.info(f"Successfully created search client for index {index_name}")
	return client
	except Exception as e:
	logger.error(f"Failed to create search client for index {index_name}: {e}")
	return None

	def _generate_embedding(self, text: str) -> List[float]:
	"""Generate embedding for vector search using Azure OpenAI"""
	try:
	# Use the embedding client to generate embeddings
	return self.embedding_client.get_embedding(text)
	except Exception as e:
	logger.error(f"Error generating embedding: {e}")
	# Return a fallback embedding (zeros would be neutral in vector space)
	import numpy as np
	embedding = np.zeros(1536) # Standard dimension for embeddings
	return embedding.tolist()

	def query_index(self, query: str, search_type: str, **kwargs) -> Dict[str, Any]:
	"""Query the Azure Search index based on search type"""

	search_text = kwargs.get("search_text", query)
	logger.info(f"Executing {search_type} search: {search_text}")

	if search_type == "vector":
	if not self.vector_client:
	return {"error": "Vector search client not configured"}
	client = self.vector_client
	#embedding = self._generate_embedding(query)
	try:
	results = list(client.search(
	search_text=search_text,
	top=kwargs.get("num_results"),
	include_total_count=True,
	filter=kwargs.get("filter"),
	order_by=kwargs.get("orderby"),
	select=kwargs.get("select"),
	vector_queries=[
	{
	"text": search_text,
	"fields": kwargs.get("vector_field", "text_vector"),
	"k": kwargs.get("num_results"),
	"kind": "text",
	}
	],
	))

	except Exception as e:
	results = [f"Vector search failed: {e}"]
	logger.info(results)
	else:
	if not self.search_client:
	return {"error": "Index search client not configured"}
	client = self.search_client

	fields = [
	"Id", "Owner", "Name", "CreatedBy", "LastModifiedBy", "ParentId",
	"CreatedDataTime", "LastModifiedDate",
	"Discriminator", "SourceLocationType", "metadata_storage_content_type",
	"metadata_storage_size", "metadata_storage_last_modified",
	"metadata_storage_content_md5", "metadata_storage_name",
	"metadata_storage_path", "metadata_storage_file_extension",
	]

	try:
	raw_results = list(client.search(
	search_text=search_text,
	query_type="full",
	filter=kwargs.get("filter"),
	order_by=kwargs.get("orderby"),
	top=kwargs.get("num_results"),
	include_total_count=True,
	select=kwargs.get("select"),
	))
	# Remove key-value pairs with null values from each result dict
	results = [
	{k: v for k, v in doc.items() if v is not None}
	for doc in raw_results
	]
	except Exception as e:
	results = [f"Index search failed: {e}"]
	logger.info(results)

	logger.info(f"Search returned {len(results)} results")
	return results

	def format_results(self, results: Dict[str, Any]) -> str:
	"""Format the search results into readable text"""
	if "error" in results:
	return f"Error: {results['error']}"

	if not results or not results.get("value") or len(results["value"]) == 0:
	return "No results found for your query."

	docs = results["value"]
	response_lines = [f"Found {len(docs)} relevant documents:"]

	for i, doc in enumerate(docs, 1): # Limit to top 5 for readability
	title = doc.get("metadata_title", doc.get("title", f"Document {i}"))
	content = doc.get("content", "")
	author = doc.get("metadata_author", "Unknown")
	score = doc.get("@search.score", 0)
	url = doc.get("url", "")

	response_lines.append(f"\n### {i}. {title}")
	response_lines.append(f"Author: {author} \| Score: {score:.2f}")

	if url:
	response_lines.append(f"URL: {url}")

	# Get a snippet from content (first 150 chars)
	if content:
	snippet = content
	response_lines.append(f"\nPreview: {snippet}")

	response_lines.append("-" * 40)

	return "\n".join(response_lines)

	def process(self, messages: List[Dict[str, Any]], **kwargs) -> List[Dict[str, Any]]:
	"""Process messages by directly querying Azure Search index"""
	# Extract the query from the last user message
	query = ""
	for message in reversed(messages):
	if message["role"] == "user":
	query = message["content"]
	break

	if not query:
	return "No query provided"

	# Determine search type based on tool name or kwargs
	tool_name = kwargs.get("tool_name", "")
	# Use explicit search_mode if provided, else fallback to tool_name logic
	search_mode = kwargs.get("search_mode")
	if search_mode in ("vector", "index"):
	search_type = search_mode
	elif "vector" in tool_name or "content" in tool_name:
	search_type = "vector"
	else:
	search_type = "index"

	results = self.query_index(query, search_type, **kwargs)
	return results

	def process_stream(self, messages: List[Dict[str, Any]], **kwargs) -> Generator[Any, None, None]:
	for chunk in self.process(messages, **kwargs):
	yield chunk

	class OrchestratorAgent(BaseAgent):
	"""Agent that coordinates between specialized agents using them as tools"""
	def __init__(self, config: AgentConfig):
	# Initialize the base agent with the given config
	super().__init__(config)

	self.tools = []
	self.data_agent = None
	self.search_agent = None

	# Set up specialized agents if tools are available
	self.tools = config.get_tools()
	self.has_tools = bool(self.tools)

	# Initialize the data agent for data source queries
	self.data_agent = DataAgent(config)

	# Initialize the search agent for Azure search queries
	self.search_agent = SearchAgent(config)

	logger.debug(f"Available tools: {json.dumps(self.tools, indent=2)}")


	def _execute_tool(self, tool_name: str, tool_args: Dict[str, Any], stream: bool = False, messages: List[Dict[str, Any]] = None) -> Union[str, Generator[Any, None, None]]:
	"""Execute a tool and return its result"""
	if not self.has_tools:
	error_msg = "No tools available for this engine"
	logger.error(error_msg)
	raise ValueError(error_msg)

	logger.info(f"Executing tool: {tool_name} with arguments: {json.dumps(tool_args, indent=2)}")

	# Use the full message history, appending the tool call as the last user message
	if messages is None:
	messages = []
	if self.config.get_system_prompt(tool_name):
	messages = [{"role": "system", "content": self.config.get_system_prompt(tool_name)}] + messages

	tool_agent_map = self.config.get_tool_agent_map()
	agent_type = tool_agent_map.get(tool_name)

	if agent_type == "search_agent":
	if stream:
	return self.search_agent.process_stream(messages, tool_name=tool_name, **tool_args)
	else:
	return self.search_agent.process(messages, tool_name=tool_name, **tool_args)
	elif agent_type == "data_agent":
	if stream:
	return self.data_agent.process_stream(messages, tool_name=tool_name, **tool_args)
	else:
	return self.data_agent.process(messages, tool_name=tool_name, **tool_args)
	elif agent_type == "base_agent":
	if stream:
	return self.process_stream(messages)
	else:
	return self.process(messages)
	else:
	logger.error(f"Unknown tool name: {tool_name}")
	return f"Error: Unknown tool name: {tool_name}"

	def _handle_tool_calls_recursive(self, messages: List[Dict[str, Any]], max_iterations: int = 10) -> Generator[Any, None, None]:
	"""Recursively handle tool calls until we get a response without tools"""
	if max_iterations <= 0 or not self.has_tools:
	logger.warning("Max iterations reached or no tools available, stopping recursive tool calls")
	return

	logger.info("Getting model's response with tool calls")
	completion = self.client.chat.completions.create(
	model=self.model_name,
	messages=messages,
	tools=self.tools,
	tool_choice="auto",
	temperature=self.temperature,
	stream=True
	)

	response_text = ""
	tool_calls = []

	for chunk in completion:
	try:
	if chunk.choices and chunk.choices[0].delta:
	delta = chunk.choices[0].delta
	if hasattr(delta, 'content') and delta.content:
	response_text += delta.content
	logger.debug(f"Assistant content chunk: {delta.content}")
	if hasattr(delta, 'tool_calls') and delta.tool_calls:
	for tool_call in delta.tool_calls:
	if tool_call.index is not None:
	while len(tool_calls) <= tool_call.index:
	tool_calls.append({})
	if tool_calls[tool_call.index] == {}:
	tool_calls[tool_call.index]["index"] = tool_call.index
	tool_calls[tool_call.index]["function"] = {"name": "", "arguments": ""}
	if tool_call.id:
	tool_calls[tool_call.index]["id"] = tool_call.id
	if tool_call.type:
	tool_calls[tool_call.index]["type"] = tool_call.type
	if tool_call.function.name:
	tool_calls[tool_call.index]["function"]["name"] = tool_call.function.name
	if tool_call.function.arguments:
	tool_calls[tool_call.index]["function"]["arguments"] += tool_call.function.arguments
	logger.debug(f"Tool call arguments chunk: {tool_call.function.arguments}")
	yield {
	"role": "assistant",
	"content": response_text,
	}
	except Exception as e:
	logger.error(f"Error processing chunk in handle_tool_calls: {e}")
	continue

	if not tool_calls:
	logger.info("No tool calls made, ending process")
	return
	else:
	logger.info(f"Tool calls detected: {json.dumps(tool_calls, indent=2)}")
	for tool_call in tool_calls:
	yield {
	"role": "assistant",
	"content": "",
	"metadata": {
	"title": tool_call["function"]["name"],
	"id": tool_call["id"],
	}
	}

	assistant_message = {
	"role": "assistant",
	"content": response_text,
	"tool_calls": tool_calls,
	}

	# Now execute tools and collect tool messages
	tool_messages = []
	for tool_call in tool_calls:
	tool_name = tool_call["function"]["name"]
	try:
	tool_args = eval(tool_call["function"]["arguments"])
	if not isinstance(tool_args, dict):
	raise ValueError("Tool arguments must be a dictionary")
	logger.debug(f"Parsed tool arguments for {tool_name}: {json.dumps(tool_args, indent=2)}")
	except (SyntaxError, ValueError) as e:
	logger.error(f"Error parsing tool arguments: {e}")
	continue

	response_stream = self._execute_tool(tool_name, tool_args, stream=True, messages=messages)
	response_list = []
	response_text = ""

	tool_message = {
	"role": "tool",
	"tool_call_id": tool_call["id"],
	"content": "",
	}
	# Capitalize tool_name and replace underscores with spaces for title
	title = tool_name.replace("_", " ").title()
	for chunk in response_stream:
	try:
	if hasattr(chunk, 'choices') and chunk.choices and hasattr(chunk.choices[0], 'delta') and hasattr(chunk.choices[0].delta, 'content'):
	response_text += chunk.choices[0].delta.content or ""
	elif isinstance(chunk, dict) or isinstance(chunk, list):
	response_list.append(chunk)
	response_text += json.dumps(chunk, indent=2)
	elif isinstance(chunk, str):
	response_text += chunk

	tool_message = {
	"role": "tool",
	"tool_call_id": tool_call["id"],
	"content": response_text
	}
	yield {
	"role": "assistant",
	"content": json.dumps(response_list[:10], indent=2) if response_list else response_text,
	"metadata": {
	"title": title,
	"id": tool_call["id"],
	"status": "pending",
	}
	}

	except Exception as e:
	response_text = f"Error processing tool response: {e}"
	logger.error(f"Error processing tool response chunk: {e}")
	break

	yield {
	"role": "assistant",
	"content": json.dumps(response_list[:10], indent=2) if response_list else response_text if response_text else "No results found.",
	"metadata": {
	"title": title,
	"id": tool_call["id"],
	"log": f"(results: {len(response_list)})" if response_list else "",
	"status": "done",
	}
	}

	tool_messages.append(tool_message)
	logger.info(f"Completed tool execution: {tool_name}")

	messages.append(assistant_message)
	messages.extend(tool_messages)

	yield from self._handle_tool_calls_recursive(messages, max_iterations - 1)

	def process_stream(self, messages: List[Dict[str, Any]]) -> Generator[Any, None, None]:
	"""Process messages and stream the response"""
	logger.info("Starting orchestration process")
	logger.debug(f"Input messages: {json.dumps(messages, indent=2)}")

	# If no tools are available, fall back to the base agent behavior
	if not self.has_tools:
	logger.info("No tools available, falling back to standard model completion")
	yield from super().process_stream(messages)
	return

	# Add the system prompt to the messages only if not already present
	system_prompt = self.config.get_system_prompt("orchestrator")
	if not messages or messages[0].get("role") != "system":
	messages = [{"role": "system", "content": system_prompt}, *messages]

	# Start the recursive tool call handling
	yield from self._handle_tool_calls_recursive(messages)