orchestration/orchestrator.py

import json
import logging
import os
import re
import time
import uuid
from abc import ABC, abstractmethod
from datetime import datetime

from .constants import Strategy, OutputFormat, OutputMode
from autogen_agentchat.teams import SelectorGroupChat
from connectors import CosmosDBClient
from .agent_strategy_factory import AgentStrategyFactory


# ---------- Configuration & Dependency Classes ----------

class OrchestratorConfig:
    def __init__(
        self,
        conversation_container: str = None,
        storage_account: str = None,
        orchestration_strategy: Strategy = None,
    ):
        self.conversation_container = conversation_container or os.environ.get('CONVERSATION_CONTAINER', 'conversations')
        self.storage_account = storage_account or os.environ.get('AZURE_STORAGE_ACCOUNT', 'your_storage_account')
        strategy_from_env = os.getenv('AUTOGEN_ORCHESTRATION_STRATEGY', 'classic_rag').replace('-', '_')
        self.orchestration_strategy = (orchestration_strategy or Strategy(strategy_from_env))


class ConversationManager:
    def __init__(self, cosmosdb_client: CosmosDBClient, config: OrchestratorConfig,
                 client_principal: dict, conversation_id: str):
        self.cosmosdb = cosmosdb_client
        self.config = config
        self.client_principal = client_principal
        self.conversation_id = self._use_or_create_conversation_id(conversation_id)
        self.short_id = self.conversation_id[:8]

    def _use_or_create_conversation_id(self, conversation_id: str) -> str:
        if not conversation_id:
            conversation_id = str(uuid.uuid4())
            logging.info(f"[orchestrator] Creating new conversation_id: {conversation_id}")
        return conversation_id

    async def get_or_create_conversation(self) -> dict:
        conversation = await self.cosmosdb.get_document(self.config.conversation_container, self.conversation_id)
        if not conversation:
            new_conversation = {
                "id": self.conversation_id,
                "start_date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "user_id": self.client_principal.get("id", "unknown") if self.client_principal else "unknown",
                "user_name": self.client_principal.get("name", "anonymous") if self.client_principal else "anonymous",
                "conversation_id": self.conversation_id,
                "history": []  # initialize an empty chat log
            }
            conversation = await self.cosmosdb.create_document(self.config.conversation_container, self.conversation_id, new_conversation)
            logging.info(f"[orchestrator] {self.short_id} Created new conversation.")
        else:
            logging.info(f"[orchestrator] {self.short_id} Retrieved existing conversation.")
        return conversation

    async def update_conversation(self, conversation: dict, ask: str, answer: str):
        logging.info(f"[orchestrator] {self.short_id} Updating conversation.")
        history = conversation.get("history", [])
        history.extend([
            {"speaker": "user", "content": ask},
            {"speaker": "assistant", "content": answer}
        ])
        simplified_conversation = {
            "id": self.conversation_id,
            "start_date": conversation.get("start_date", datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
            "user_id": self.client_principal.get("id", "unknown") if self.client_principal else "unknown",
            "user_name": self.client_principal.get("name", "anonymous") if self.client_principal else "anonymous",
            "conversation_id": self.conversation_id,
            "history": history
        }
        await self.cosmosdb.update_document(self.config.conversation_container, simplified_conversation)
        logging.info(f"[orchestrator] {self.short_id} Finished updating conversation.")


# ---------- Parsing Helpers ----------

class ChatLogParser:
    @staticmethod
    def get_chat_log(messages) -> list:
        def make_serializable(obj):
            if isinstance(obj, (list, tuple)):
                return [make_serializable(item) for item in obj]
            elif isinstance(obj, dict):
                return {key: make_serializable(value) for key, value in obj.items()}
            elif isinstance(obj, str):
                return obj
            else:
                return repr(obj)

        chat_log = []
        for msg in messages:
            safe_content = make_serializable(msg.content)
            chat_log.append({
                "speaker": msg.source,
                "message_type": msg.type,
                "content": safe_content
            })
        return chat_log

    @staticmethod
    def extract_data_points(chat_log) -> list:
        data_points = []
        call_id_map = {}
        allowed_extensions = {'vtt', 'xlsx', 'xls', 'pdf', 'png', 'jpeg', 'jpg', 'bmp', 'tiff', 'docx', 'pptx'}
        extension_pattern = "|".join(allowed_extensions)
        pattern = rf'[\w\-.]+\.(?:{extension_pattern}): .*?(?=(?:[\w\-.]+\.(?:{extension_pattern})\:)|$)'
        if chat_log:
            for msg in chat_log:
                try:
                    if "message_type" in msg and "content" in msg and isinstance(msg["content"], list) and msg["content"]:
                        if msg["message_type"] == "ToolCallRequestEvent":
                            content = msg["content"][0]
                            if "id='" in content:
                                call_id = content.split("id='")[1].split("',")[0]
                                call_id_map[call_id] = None
                        elif msg["message_type"] == "ToolCallExecutionEvent":
                            content = msg["content"][0]
                            if "call_id='" in content:
                                call_id = content.split("call_id='")[1].split("')")[0]
                                if call_id in call_id_map:
                                    if "content='" in content:
                                        data = content.split("content='")[1].rsplit("',", 1)[0]
                                        entries = re.findall(pattern, data, re.DOTALL | re.IGNORECASE)
                                        data_points.extend(entries)
                except Exception as e:
                    logging.warning(f"[orchestrator] Error processing message: {e}.")
        else:
            logging.warning("[orchestrator] Chat log is empty or not provided.")
        return data_points


class MessageParser:
    @staticmethod
    def parse_message(message_str: str) -> dict:
        pattern = r"(\w+)=((?:'(?:\\'|[^'])*'|\"(?:\\\"|[^\"])*\"|\[[^\]]+\]|\([^\)]+\)|[^\s]+))"
        pairs = re.findall(pattern, message_str)
        result = {}
        for key, value in pairs:
            if (value.startswith("'") and value.endswith("'")) or (value.startswith("\"") and value.endswith("\"")):
                value = value[1:-1]
            result[key] = value
        if 'models_usage' in result and result['models_usage'].startswith('RequestUsage'):
            start_index = message_str.find('models_usage=')
            if start_index != -1:
                usage_str = message_str[start_index:].split(' ', 1)[1]
                usage_match = re.match(r"RequestUsage\((.*?)\)", usage_str)
                if usage_match:
                    result['models_usage'] = usage_match.group(1)
        return result


# ---------- Orchestrators ----------

class BaseOrchestrator(ABC):
    def __init__(
        self,
        conversation_id: str,
        config: OrchestratorConfig,
        client_principal: dict = None,
        access_token: str = None,
        agent_strategy=None
    ):
        self.client_principal = client_principal
        self.access_token = access_token
        self.cosmosdb =  CosmosDBClient()
        self.conversation_manager = ConversationManager(self.cosmosdb, config, client_principal, conversation_id)
        self.conversation_id = self.conversation_manager.conversation_id
        self.short_id = self.conversation_id[:8]
        self.config = config

        # Agent strategy is injected if provided; otherwise, use the factory.
        if agent_strategy:
            self.agent_strategy = agent_strategy
        else:
            strategy_key = config.orchestration_strategy
            self.agent_strategy = AgentStrategyFactory.get_strategy(strategy_key)

    @abstractmethod
    async def answer(self, ask: str):
        pass

    @abstractmethod
    async def _create_agents_with_strategy(self, history: list[dict]) -> dict:
        pass

    async def _update_conversation(self, conversation: dict, ask: str, answer: str):
        await self.conversation_manager.update_conversation(conversation, ask, answer)


class RequestResponseOrchestrator(BaseOrchestrator):
    async def answer(self, ask: str) -> dict:
        start_time = time.time()
        conversation = await self.conversation_manager.get_or_create_conversation()
        history = conversation.get("history", [])
        agent_configuration = await self._create_agents_with_strategy(history)
        answer_dict = await self._initiate_group_chat(agent_configuration, ask)
        await self._update_conversation(conversation, ask, answer_dict['answer'])
        response_time = time.time() - start_time
        logging.info(f"[orchestrator] {self.short_id} Generated response in {response_time:.3f} sec.")
        return answer_dict

    async def _initiate_group_chat(self, agent_configuration: dict, ask: str) -> dict:
        try:
            logging.info(f"[orchestrator] {self.short_id} Creating group chat via SelectorGroupChat.")
            group_chat = SelectorGroupChat(
                participants=agent_configuration["agents"],
                model_client=agent_configuration["model_client"],
                termination_condition=agent_configuration["termination_condition"],
                selector_func=agent_configuration["selector_func"],
            )
            result = await group_chat.run(task=ask)

            # Parse the last message for answer and reasoning.
            reasoning = ""
            last_message = result.messages[-1].content if result.messages else '{"answer":"No answer provided."}'
            try:
                if isinstance(last_message, list):
                    message_data = last_message
                elif isinstance(last_message, (str, bytes, bytearray)):
                    message_data = json.loads(last_message)
                else:
                    raise TypeError(f"Unexpected message type: {type(last_message)}")
                answer = message_data.get("answer", "Oops! The agent team did not generate a response for the user.")
                reasoning = message_data.get("reasoning", "")
            except json.JSONDecodeError:
                answer = "Oops! The agent team did not generate a response for the user."
                logging.warning(f"[orchestrator] {self.short_id} Error: Malformed JSON. Using default values for answer and reasoning.")

            if answer.endswith(agent_configuration['terminate_message']):
                answer = answer[:-len(agent_configuration['terminate_message'])].strip()

            chat_log = ChatLogParser.get_chat_log(result.messages)
            data_points = ChatLogParser.extract_data_points(chat_log)

            answer_dict = {
                "conversation_id": self.conversation_id,
                "answer": answer,
                "reasoning": reasoning,
                "thoughts": chat_log,
                "data_points": data_points
            }
            return answer_dict

        except Exception as e:
            logging.error(f"[orchestrator] {self.short_id} An error occurred: {str(e)}", exc_info=True)
            return {
                "conversation_id": self.conversation_id,
                "answer": f"We encountered an issue processing your request. Please try again later. Error: {str(e)}",
                "reasoning": "",
                "thoughts": [],
                "data_points": []
            }

    async def _create_agents_with_strategy(self, history: list[dict]) -> dict:
        logging.info(f"[orchestrator] {self.short_id} Creating agents using {self.agent_strategy.strategy_type} strategy.")
        return await self.agent_strategy.create_agents(history, self.client_principal, self.access_token, OutputMode.REQUEST_RESPONSE, OutputFormat.JSON)


class StreamingOrchestrator(BaseOrchestrator):
    def __init__(
        self,
        conversation_id: str,
        config: OrchestratorConfig,
        client_principal: dict = None,
        access_token: str = None,
        agent_strategy=None
    ):
        super().__init__(conversation_id, config, client_principal, access_token, agent_strategy)
        self.optimize_for_audio = False

    def set_optimize_for_audio(self, optimize_for_audio: bool):
        self.optimize_for_audio = optimize_for_audio

    async def answer(self, ask: str):
        conversation = await self.conversation_manager.get_or_create_conversation()
        history = conversation.get("history", [])
        agent_config = await self._create_agents_with_strategy(history)
        group_chat = SelectorGroupChat(
            participants=agent_config["agents"],
            model_client=agent_config["model_client"],
            termination_condition=agent_config["termination_condition"],
            selector_func=agent_config["selector_func"],
        )
        stream = group_chat.run_stream(task=ask)
        streamed_conversation_id = False
        final_answer = ""
        try:
            async for response in stream:
                msg_str = str(response)
                try:
                    msg = MessageParser.parse_message(msg_str)
                except Exception as e:
                    logging.debug(f"Exception while parsing message: {e}")
                    continue
                msg_type = msg.get("type", "")
                msg_content = msg.get("content", "")
                if msg_type == 'ModelClientStreamingChunkEvent':
                    # stream conversation_id in the first chunk
                    if not streamed_conversation_id:
                        yield f"{conversation['id']} "
                        streamed_conversation_id = True

                    yield msg_content
                    # logging.info(f"Yielding chunk: {msg_content}")
                    final_answer += msg_content
        except Exception as error:
            logging.error(f"Error in streaming response: {error}", exc_info=True)
            error_message = f"\nWe encountered an issue processing your request. Please contact our support team for assistance. \n\n *{error}*"
            yield error_message
            final_answer += error_message

        await self._update_conversation(conversation, ask, final_answer)

    async def _create_agents_with_strategy(self, history: list[dict]) -> dict:
        output_format = OutputFormat.TEXT_TTS if self.optimize_for_audio else OutputFormat.TEXT
        logging.info(f"[orchestrator] {self.short_id} Creating agents using {self.agent_strategy.strategy_type} strategy. In Streaming mode and {output_format} output format.")        
        return await self.agent_strategy.create_agents(history, self.client_principal, self.access_token, OutputMode.STREAMING, output_format)