massgen 0.0.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

massgen/backend/azure_openai.py

@@ -0,0 +1,385 @@
+# -*- coding: utf-8 -*-
+"""
+Azure OpenAI backend implementation.
+Uses the official Azure OpenAI client for proper Azure integration.
+"""
+# -*- coding: utf-8 -*-
+from __future__ import annotations
+
+import os
+from typing import Any, AsyncGenerator, Dict, List, Optional
+
+from ..logger_config import (
+    log_backend_activity,
+    log_backend_agent_message,
+    log_stream_chunk,
+)
+from .base import LLMBackend, StreamChunk
+
+
+class AzureOpenAIBackend(LLMBackend):
+    """Azure OpenAI backend using the official Azure OpenAI client.
+
+    Supports Azure OpenAI deployments with proper Azure authentication and configuration.
+
+    Environment Variables:
+        AZURE_OPENAI_API_KEY: Azure OpenAI API key
+        AZURE_OPENAI_ENDPOINT: Azure OpenAI endpoint URL
+        AZURE_OPENAI_API_VERSION: Azure OpenAI API version (optional, defaults to 2024-12-01-preview)
+    """
+
+    def __init__(self, api_key: Optional[str] = None, **kwargs):
+        super().__init__(api_key, **kwargs)
+
+        # Get Azure configuration from parameters or environment variables
+        self.api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
+
+        if not self.api_key:
+            raise ValueError("Azure OpenAI API key is required. Set AZURE_OPENAI_API_KEY environment variable or pass api_key parameter.")
+
+    def get_provider_name(self) -> str:
+        """Get the name of this provider."""
+        return "Azure OpenAI"
+
+    async def stream_with_tools(self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], **kwargs) -> AsyncGenerator[StreamChunk, None]:
+        """
+        Stream a response with tool calling support using Azure OpenAI.
+
+        Args:
+            messages: Conversation messages
+            tools: Available tools schema
+            **kwargs: Additional parameters including model (deployment name)
+        """
+        # Extract agent_id for logging
+        agent_id = kwargs.get("agent_id", None)
+
+        log_backend_activity(
+            self.get_provider_name(),
+            "Starting stream_with_tools",
+            {"num_messages": len(messages), "num_tools": len(tools) if tools else 0},
+            agent_id=agent_id,
+        )
+
+        try:
+            # Merge constructor config with stream kwargs (stream kwargs take priority)
+            all_params = {**self.config, **kwargs}
+
+            # Import Azure OpenAI client
+            from openai import AsyncAzureOpenAI
+
+            azure_endpoint = all_params.get("azure_endpoint") or all_params.get("base_url") or os.getenv("AZURE_OPENAI_ENDPOINT")
+            api_version = all_params.get("api_version") or os.getenv("AZURE_OPENAI_API_VERSION", "2024-12-01-preview")
+
+            # Validate required configuration
+            if not azure_endpoint:
+                raise ValueError("Azure OpenAI endpoint URL is required. Set AZURE_OPENAI_ENDPOINT environment variable or pass azure_endpoint/base_url parameter.")
+
+            if not api_version:
+                raise ValueError("Azure OpenAI API version is required. Set AZURE_OPENAI_API_VERSION environment variable or pass api_version parameter.")
+
+            # Clean up endpoint URL
+            if azure_endpoint.endswith("/"):
+                azure_endpoint = azure_endpoint[:-1]
+
+            # Initialize Azure OpenAI client
+            self.client = AsyncAzureOpenAI(
+                api_version=api_version,
+                azure_endpoint=azure_endpoint,
+                api_key=self.api_key,
+            )
+
+            # Get deployment name from model parameter
+            deployment_name = all_params.get("model")
+            if not deployment_name:
+                raise ValueError("Azure OpenAI requires a deployment name. Pass it as the 'model' parameter.")
+
+            # Check if workflow tools are present
+            workflow_tools = [t for t in tools if t.get("function", {}).get("name") in ["new_answer", "vote"]] if tools else []
+            has_workflow_tools = len(workflow_tools) > 0
+
+            # Modify messages to include workflow tool instructions if needed
+            modified_messages = self._prepare_messages_with_workflow_tools(messages, workflow_tools) if has_workflow_tools else messages
+
+            # Log messages being sent
+            log_backend_agent_message(
+                agent_id or "default",
+                "SEND",
+                {"messages": modified_messages, "tools": len(tools) if tools else 0},
+                backend_name=self.get_provider_name(),
+            )
+
+            # Prepare API parameters
+            api_params = {
+                "messages": modified_messages,
+                "model": deployment_name,  # Use deployment name directly
+                "stream": True,
+            }
+
+            # Only add tools if explicitly provided and not empty
+            if tools and len(tools) > 0:
+                # Convert tools to Azure OpenAI format if needed
+                converted_tools = self._convert_tools_format(tools)
+                api_params["tools"] = converted_tools
+            else:
+                # Disable tool calling for simple queries
+                api_params["tool_choice"] = "none"
+
+            # Add other parameters (excluding model since we already set it)
+            # Filter out unsupported Azure OpenAI parameters
+            excluded_params = self.get_base_excluded_config_params() | {
+                # Azure OpenAI specific exclusions
+                "model",
+                "messages",
+                "stream",
+                "tools",
+            }
+            for key, value in kwargs.items():
+                if key not in excluded_params and value is not None:
+                    api_params[key] = value
+
+            # Create streaming response (now properly async)
+            stream = await self.client.chat.completions.create(**api_params)
+
+            # Process streaming response with content accumulation
+            accumulated_content = ""
+            complete_response = ""  # Keep track of the complete response
+            last_yield_type = None
+
+            async for chunk in stream:
+                converted = self._convert_chunk_to_stream_chunk(chunk)
+
+                # Accumulate content chunks
+                if converted.type == "content" and converted.content:
+                    accumulated_content += converted.content
+                    complete_response += converted.content  # Add to complete response
+                    # Only yield content when we have meaningful chunks (words, not single characters)
+                    if len(accumulated_content) >= 10 or " " in accumulated_content:
+                        log_backend_agent_message(
+                            agent_id or "default",
+                            "RECV",
+                            {"content": accumulated_content},
+                            backend_name=self.get_provider_name(),
+                        )
+                        log_stream_chunk(
+                            "backend.azure_openai",
+                            "content",
+                            accumulated_content,
+                            agent_id,
+                        )
+                        yield StreamChunk(type="content", content=accumulated_content)
+                        accumulated_content = ""
+                elif converted.type != "content":
+                    # Log non-content chunks
+                    if converted.type == "error":
+                        log_stream_chunk("backend.azure_openai", "error", converted.error, agent_id)
+                    elif converted.type == "done":
+                        log_stream_chunk("backend.azure_openai", "done", None, agent_id)
+                    # Yield non-content chunks immediately
+                    last_yield_type = converted.type
+                    yield converted
+
+            # Yield any remaining accumulated content
+            if accumulated_content:
+                log_backend_agent_message(
+                    agent_id or "default",
+                    "RECV",
+                    {"content": accumulated_content},
+                    backend_name=self.get_provider_name(),
+                )
+                log_stream_chunk("backend.azure_openai", "content", accumulated_content, agent_id)
+                yield StreamChunk(type="content", content=accumulated_content)
+
+            # After streaming is complete, check if we have workflow tool calls
+            if has_workflow_tools:
+                workflow_tool_calls = self._extract_workflow_tool_calls(complete_response)
+                if workflow_tool_calls:
+                    log_stream_chunk(
+                        "backend.azure_openai",
+                        "tool_calls",
+                        workflow_tool_calls,
+                        agent_id,
+                    )
+                    yield StreamChunk(type="tool_calls", tool_calls=workflow_tool_calls)
+                    last_yield_type = "tool_calls"
+
+            # Ensure stream termination is signaled
+            if last_yield_type != "done":
+                log_stream_chunk("backend.azure_openai", "done", None, agent_id)
+                yield StreamChunk(type="done")
+
+        except Exception as e:
+            error_msg = f"Azure OpenAI API error: {str(e)}"
+            log_stream_chunk("backend.azure_openai", "error", error_msg, agent_id)
+            yield StreamChunk(type="error", error=error_msg)
+
+    def _prepare_messages_with_workflow_tools(self, messages: List[Dict[str, Any]], workflow_tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Prepare messages with workflow tool instructions."""
+        if not workflow_tools:
+            return messages
+
+        # Find the system message
+        system_message = None
+        for msg in messages:
+            if msg.get("role") == "system":
+                system_message = msg
+                break
+
+        # Create enhanced system message with workflow tool instructions
+        enhanced_system = self._build_workflow_tools_system_prompt(system_message.get("content", "") if system_message else "", workflow_tools)
+
+        # Create new messages list with enhanced system message
+        new_messages = []
+        for msg in messages:
+            if msg.get("role") == "system":
+                new_messages.append({"role": "system", "content": enhanced_system})
+            else:
+                new_messages.append(msg)
+
+        return new_messages
+
+    def _build_workflow_tools_system_prompt(self, base_system: str, workflow_tools: List[Dict[str, Any]]) -> str:
+        """Build system prompt with workflow tool instructions."""
+        system_parts = []
+
+        if base_system:
+            system_parts.append(base_system)
+
+        # Add workflow tools information
+        if workflow_tools:
+            system_parts.append("\n--- Available Tools ---")
+            for tool in workflow_tools:
+                name = tool.get("function", {}).get("name", "unknown")
+                description = tool.get("function", {}).get("description", "No description")
+                system_parts.append(f"- {name}: {description}")
+
+                # Add usage examples for workflow tools
+                if name == "new_answer":
+                    system_parts.append('    Usage: {"tool_name": "new_answer", ' '"arguments": {"content": "your answer"}}')
+                elif name == "vote":
+                    # Extract valid agent IDs from enum if available
+                    agent_id_enum = None
+                    for t in workflow_tools:
+                        if t.get("function", {}).get("name") == "vote":
+                            agent_id_param = t.get("function", {}).get("parameters", {}).get("properties", {}).get("agent_id", {})
+                            if "enum" in agent_id_param:
+                                agent_id_enum = agent_id_param["enum"]
+                            break
+
+                    if agent_id_enum:
+                        agent_list = ", ".join(agent_id_enum)
+                        system_parts.append(f'    Usage: {{"tool_name": "vote", ' f'"arguments": {{"agent_id": "agent1", ' f'"reason": "explanation"}}}} // Choose agent_id from: {agent_list}')
+                    else:
+                        system_parts.append('    Usage: {"tool_name": "vote", ' '"arguments": {"agent_id": "agent1", ' '"reason": "explanation"}}')
+
+            system_parts.append("\n--- MassGen Workflow Instructions ---")
+            system_parts.append("IMPORTANT: You must respond with a structured JSON decision at the end of your response.")
+            system_parts.append("You must use the coordination tools (new_answer, vote) " "to participate in multi-agent workflows.")
+            system_parts.append("The JSON MUST be formatted as a strict JSON code block:")
+            system_parts.append("1. Start with ```json on one line")
+            system_parts.append("2. Include your JSON content (properly formatted)")
+            system_parts.append("3. End with ``` on one line")
+            system_parts.append('Example format:\n```json\n{"tool_name": "vote", "arguments": {"agent_id": "agent1", "reason": "explanation"}}\n```')
+            system_parts.append("The JSON block should be placed at the very end of your response, after your analysis.")
+
+        return "\n".join(system_parts)
+
+    def _extract_workflow_tool_calls(self, content: str) -> List[Dict[str, Any]]:
+        """Extract workflow tool calls from content."""
+        try:
+            import json
+            import re
+
+            # Look for JSON inside markdown code blocks first
+            markdown_json_pattern = r"```json\s*(\{.*?\})\s*```"
+            markdown_matches = re.findall(markdown_json_pattern, content, re.DOTALL)
+
+            for match in reversed(markdown_matches):
+                try:
+                    parsed = json.loads(match.strip())
+                    if isinstance(parsed, dict) and "tool_name" in parsed:
+                        # Convert to MassGen tool call format
+                        tool_call = {
+                            "id": f"call_{hash(match) % 10000}",  # Generate a unique ID
+                            "type": "function",
+                            "function": {
+                                "name": parsed["tool_name"],
+                                "arguments": json.dumps(parsed["arguments"]),
+                            },
+                        }
+                        return [tool_call]
+                except json.JSONDecodeError:
+                    continue
+
+            # Also look for JSON without markdown blocks
+            json_pattern = r'\{[^{}]*"tool_name"[^{}]*\}'
+            json_matches = re.findall(json_pattern, content, re.DOTALL)
+
+            for match in json_matches:
+                try:
+                    parsed = json.loads(match.strip())
+                    if isinstance(parsed, dict) and "tool_name" in parsed:
+                        # Convert to MassGen tool call format
+                        tool_call = {
+                            "id": f"call_{hash(match) % 10000}",  # Generate a unique ID
+                            "type": "function",
+                            "function": {
+                                "name": parsed["tool_name"],
+                                "arguments": json.dumps(parsed["arguments"]),
+                            },
+                        }
+                        return [tool_call]
+                except json.JSONDecodeError:
+                    continue
+
+            return []
+
+        except Exception:
+            return []
+
+    def _convert_tools_format(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Convert tools to Azure OpenAI format if needed."""
+        # Azure OpenAI uses the same tool format as OpenAI
+        return tools
+
+    def _convert_chunk_to_stream_chunk(self, chunk) -> StreamChunk:
+        """Convert Azure OpenAI chunk to MassGen StreamChunk format."""
+        try:
+            if hasattr(chunk, "choices") and chunk.choices:
+                choice = chunk.choices[0]
+
+                if hasattr(choice, "delta") and choice.delta:
+                    delta = choice.delta
+
+                    # Handle content - this should be the main response
+                    if hasattr(delta, "content") and delta.content:
+                        return StreamChunk(type="content", content=delta.content)
+
+                    # Handle tool calls - but only if we actually want them
+                    if hasattr(delta, "tool_calls") and delta.tool_calls:
+                        # For now, let's ignore tool calls and treat them as content
+                        # This prevents the empty response issue
+                        tool_call_text = ""
+                        for tool_call in delta.tool_calls:
+                            if hasattr(tool_call, "function") and tool_call.function:
+                                if hasattr(tool_call.function, "arguments") and tool_call.function.arguments:
+                                    tool_call_text += tool_call.function.arguments
+
+                        if tool_call_text:
+                            return StreamChunk(type="content", content=tool_call_text)
+
+                    # Handle finish reason
+                    if hasattr(choice, "finish_reason") and choice.finish_reason:
+                        if choice.finish_reason == "stop":
+                            return StreamChunk(type="done")
+                        elif choice.finish_reason == "tool_calls":
+                            return StreamChunk(type="done")  # Treat as done
+
+            # Default chunk - this should not happen for valid responses
+            return StreamChunk(type="content", content="")
+
+        except Exception as e:
+            return StreamChunk(type="error", error=f"Error processing chunk: {str(e)}")
+
+    def extract_tool_call_id(self, tool_call: Dict[str, Any]) -> str:
+        """Extract tool call id from Chat Completions-style tool call."""
+        return tool_call.get("id", "")