agents.py

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)