/aisuite/providers/google_provider.py

"""The interface to Google's Vertex AI."""

import os
import json
from typing import List, Dict, Any, Optional, Union, BinaryIO, AsyncGenerator

import vertexai
from vertexai.generative_models import (
    GenerativeModel,
    GenerationConfig,
    Content,
    Part,
    Tool,
    FunctionDeclaration,
)
import pprint

from aisuite.framework import ChatCompletionResponse, Message
from aisuite.framework.message import (
    TranscriptionResult,
    Word,
    Segment,
    Alternative,
    StreamingTranscriptionChunk,
)
from aisuite.provider import Provider, ASRError, Audio

DEFAULT_TEMPERATURE = 0.7
ENABLE_DEBUG_MESSAGES = False

# Links.
# https://codelabs.developers.google.com/codelabs/gemini-function-calling#6
# https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling#chat-samples


class GoogleMessageConverter:
    @staticmethod
    def convert_user_role_message(message: Dict[str, Any]) -> Content:
        """Convert user or system messages to Google Vertex AI format."""
        parts = [Part.from_text(message["content"])]
        return Content(role="user", parts=parts)

    @staticmethod
    def convert_assistant_role_message(message: Dict[str, Any]) -> Content:
        """Convert assistant messages to Google Vertex AI format."""
        if "tool_calls" in message and message["tool_calls"]:
            # Handle function calls
            tool_call = message["tool_calls"][
                0
            ]  # Assuming single function call for now
            function_call = tool_call["function"]

            # Create a Part from the function call
            parts = [
                Part.from_dict(
                    {
                        "function_call": {
                            "name": function_call["name"],
                            # "arguments": json.loads(function_call["arguments"])
                        }
                    }
                )
            ]
            # return Content(role="function", parts=parts)
        else:
            # Handle regular text messages
            parts = [Part.from_text(message["content"])]
            # return Content(role="model", parts=parts)

        return Content(role="model", parts=parts)

    @staticmethod
    def convert_tool_role_message(message: Dict[str, Any]) -> Part:
        """Convert tool messages to Google Vertex AI format."""
        if "content" not in message:
            raise ValueError("Tool result message must have a content field")

        try:
            content_json = json.loads(message["content"])
            part = Part.from_function_response(
                name=message["name"], response=content_json
            )
            # Wrap tool response in Content so history is a list of Content objects.
            return Content(role="function", parts=[part])
        except json.JSONDecodeError:
            raise ValueError("Tool result message must be valid JSON")

    @staticmethod
    def convert_request(messages: List[Dict[str, Any]]) -> List[Content]:
        """Convert messages to Google Vertex AI format."""
        # Convert all messages to dicts if they're Message objects
        messages = [
            message.model_dump() if hasattr(message, "model_dump") else message
            for message in messages
        ]

        formatted_messages = []
        for message in messages:
            if message["role"] == "tool":
                vertex_message = GoogleMessageConverter.convert_tool_role_message(
                    message
                )
                if vertex_message:
                    formatted_messages.append(vertex_message)
            elif message["role"] == "assistant":
                formatted_messages.append(
                    GoogleMessageConverter.convert_assistant_role_message(message)
                )
            else:  # user or system role
                formatted_messages.append(
                    GoogleMessageConverter.convert_user_role_message(message)
                )

        return formatted_messages

    @staticmethod
    def convert_response(response) -> ChatCompletionResponse:
        """Normalize the response from Vertex AI to match OpenAI's response format."""
        openai_response = ChatCompletionResponse()

        if ENABLE_DEBUG_MESSAGES:
            print("Dumping the response")
            pprint.pprint(response)

        # TODO: We need to go through each part, because function call may not be the first part.
        #       Currently, we are only handling the first part, but this is not enough.
        #
        # This is a valid response:
        # candidates {
        #   content {
        #     role: "model"
        #     parts {
        #       text: "The current temperature in San Francisco is 72 degrees Celsius. \n\n"
        #     }
        #     parts {
        #       function_call {
        #         name: "is_it_raining"
        #         args {
        #           fields {
        #             key: "location"
        #             value {
        #               string_value: "San Francisco"
        #             }
        #           }
        #         }
        #       }
        #     }
        #   }
        #   finish_reason: STOP

        # Check if the response contains function calls
        # Note: Just checking if the function_call attribute exists is not enough,
        #       it is important to check if the function_call is not None.
        if (
            hasattr(response.candidates[0].content.parts[0], "function_call")
            and response.candidates[0].content.parts[0].function_call
        ):
            function_call = response.candidates[0].content.parts[0].function_call

            # args is a MapComposite.
            # Convert the MapComposite to a dictionary
            args_dict = {}
            # Another way to try is: args_dict = dict(function_call.args)
            for key, value in function_call.args.items():
                args_dict[key] = value
            if ENABLE_DEBUG_MESSAGES:
                print("Dumping the args_dict")
                pprint.pprint(args_dict)

            openai_response.choices[0].message = {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "type": "function",
                        "id": f"call_{hash(function_call.name)}",  # Generate a unique ID
                        "function": {
                            "name": function_call.name,
                            "arguments": json.dumps(args_dict),
                        },
                    }
                ],
                "refusal": None,
            }
            openai_response.choices[0].message = Message(
                **openai_response.choices[0].message
            )
            openai_response.choices[0].finish_reason = "tool_calls"
        else:
            # Handle regular text response
            openai_response.choices[0].message.content = (
                response.candidates[0].content.parts[0].text
            )
            openai_response.choices[0].finish_reason = "stop"

        return openai_response


class GoogleProvider(Provider):
    """Implements the ProviderInterface for interacting with Google's Vertex AI."""

    def __init__(self, **config):
        """Set up the Google AI client with a project ID."""
        super().__init__()

        self.project_id = config.get("project_id") or os.getenv("GOOGLE_PROJECT_ID")
        self.location = config.get("region") or os.getenv("GOOGLE_REGION")
        self.app_creds_path = config.get("application_credentials") or os.getenv(
            "GOOGLE_APPLICATION_CREDENTIALS"
        )

        if not self.project_id or not self.location or not self.app_creds_path:
            raise EnvironmentError(
                "Missing one or more required Google environment variables: "
                "GOOGLE_PROJECT_ID, GOOGLE_REGION, GOOGLE_APPLICATION_CREDENTIALS. "
                "Please refer to the setup guide: /guides/google.md."
            )

        vertexai.init(project=self.project_id, location=self.location)

        self.transformer = GoogleMessageConverter()

        # Initialize Speech client lazily
        self._speech_client = None

        # Initialize audio functionality
        self.audio = GoogleAudio(self)

    def chat_completions_create(self, model, messages, **kwargs):
        """Request chat completions from the Google AI API.

        Args:
        ----
            model (str): Identifies the specific provider/model to use.
            messages (list of dict): A list of message objects in chat history.
            kwargs (dict): Optional arguments for the Google AI API.

        Returns:
        -------
            The ChatCompletionResponse with the completion result.

        """

        # Set the temperature if provided, otherwise use the default
        temperature = kwargs.get("temperature", DEFAULT_TEMPERATURE)

        # Convert messages to Vertex AI format
        message_history = self.transformer.convert_request(messages)

        # Handle tools if provided
        tools = None
        if "tools" in kwargs:
            tools = [
                Tool(
                    function_declarations=[
                        FunctionDeclaration(
                            name=tool["function"]["name"],
                            description=tool["function"].get("description", ""),
                            parameters={
                                "type": "object",
                                "properties": {
                                    param_name: {
                                        "type": param_info.get("type", "string"),
                                        "description": param_info.get(
                                            "description", ""
                                        ),
                                        **(
                                            {"enum": param_info["enum"]}
                                            if "enum" in param_info
                                            else {}
                                        ),
                                    }
                                    for param_name, param_info in tool["function"][
                                        "parameters"
                                    ]["properties"].items()
                                },
                                "required": tool["function"]["parameters"].get(
                                    "required", []
                                ),
                            },
                        )
                        for tool in kwargs["tools"]
                    ]
                )
            ]

        # Create the GenerativeModel
        model = GenerativeModel(
            model,
            generation_config=GenerationConfig(temperature=temperature),
            tools=tools,
        )

        if ENABLE_DEBUG_MESSAGES:
            print("Dumping the message_history")
            pprint.pprint(message_history)

        # Start chat and get response
        chat = model.start_chat(history=message_history[:-1])
        last_message = message_history[-1]

        # If the last message is a function response, send the Part object directly
        # Otherwise, send just the text content
        message_to_send = (
            last_message
            if isinstance(last_message, Content) and last_message.role == "function"
            else last_message.parts[0].text
        )
        # response = chat.send_message(message_to_send)
        response = chat.send_message(message_to_send)

        # Convert and return the response
        return self.transformer.convert_response(response)

    @property
    def speech_client(self):
        """Lazy initialization of Google Cloud Speech client."""
        if self._speech_client is None:
            try:
                from google.cloud import speech

                self._speech_client = speech.SpeechClient()
            except ImportError:
                raise ImportError(
                    "google-cloud-speech is required for ASR functionality. "
                    "Install it with: pip install google-cloud-speech"
                )
        return self._speech_client


# Audio Classes
class GoogleAudio(Audio):
    """Google Audio functionality container."""

    def __init__(self, provider):
        super().__init__()
        self.provider = provider
        self.transcriptions = self.Transcriptions(provider)

    class Transcriptions(Audio.Transcription):
        """Google Audio Transcriptions functionality."""

        def __init__(self, provider):
            self.provider = provider

        def create(
            self,
            model: str,
            file: Union[str, BinaryIO],
            **kwargs,
        ) -> TranscriptionResult:
            """
            Create audio transcription using Google Cloud Speech-to-Text API.

            All parameters are already validated and mapped by the Client layer.
            This is a simple pass-through to the Google API.
            """
            try:
                from google.cloud import speech

                # Set defaults
                kwargs["model"] = model if model != "default" else "latest_long"
                kwargs.setdefault("sample_rate_hertz", 16000)
                kwargs.setdefault("enable_automatic_punctuation", True)

                audio_data = self._read_audio_data(file)
                audio = speech.RecognitionAudio(content=audio_data)
                config = self._build_recognition_config(kwargs, speech, file)

                response = self.provider.speech_client.recognize(
                    config=config, audio=audio
                )
                return self._parse_google_response(response)

            except ImportError:
                raise ASRError(
                    "google-cloud-speech is required for ASR functionality. "
                    "Install it with: pip install google-cloud-speech"
                )
            except Exception as e:
                raise ASRError(f"Google Speech-to-Text error: {e}") from e

        async def create_stream_output(
            self,
            model: str,
            file: Union[str, BinaryIO],
            **kwargs,
        ) -> AsyncGenerator[StreamingTranscriptionChunk, None]:
            """
            Create streaming audio transcription using Google Cloud Speech-to-Text API.

            All parameters are already validated and mapped by the Client layer.
            This implementation handles streaming with Google's API.
            """
            try:
                from google.cloud import speech

                # Set defaults
                kwargs["model"] = model if model != "default" else "latest_long"
                kwargs.setdefault("sample_rate_hertz", 16000)
                kwargs.setdefault("enable_automatic_punctuation", True)

                config = self._build_recognition_config(kwargs, speech, file)
                streaming_config = speech.StreamingRecognitionConfig(
                    config=config, interim_results=True, single_utterance=False
                )

                audio_data = self._read_audio_data(file)
                request_generator = self._create_streaming_requests(
                    speech, streaming_config, audio_data
                )

                responses = self.provider.speech_client.streaming_recognize(
                    config=streaming_config, requests=request_generator
                )

                for response in responses:
                    for result in response.results:
                        if result.alternatives:
                            alternative = result.alternatives[0]
                            yield StreamingTranscriptionChunk(
                                text=alternative.transcript,
                                is_final=result.is_final,
                                confidence=getattr(alternative, "confidence", None),
                            )

            except ImportError:
                raise ASRError(
                    "google-cloud-speech is required for ASR functionality. "
                    "Install it with: pip install google-cloud-speech"
                )
            except Exception as e:
                raise ASRError(f"Google Speech-to-Text streaming error: {e}") from e

        def _read_audio_data(self, file: Union[str, BinaryIO]) -> bytes:
            """Read audio data from file or file-like object."""
            if isinstance(file, str):
                with open(file, "rb") as audio_file:
                    return audio_file.read()
            else:
                return file.read()

        def _detect_audio_encoding(self, file: Union[str, BinaryIO], speech):
            """Detect audio encoding based on file extension or content."""
            if isinstance(file, str):
                # File path - detect by extension
                file_lower = file.lower()
                if file_lower.endswith(".mp3"):
                    return speech.RecognitionConfig.AudioEncoding.MP3
                elif file_lower.endswith(".flac"):
                    return speech.RecognitionConfig.AudioEncoding.FLAC
                elif file_lower.endswith(".wav"):
                    return speech.RecognitionConfig.AudioEncoding.LINEAR16
                elif file_lower.endswith(".ogg"):
                    return speech.RecognitionConfig.AudioEncoding.OGG_OPUS
                elif file_lower.endswith(".webm"):
                    return speech.RecognitionConfig.AudioEncoding.WEBM_OPUS

            # Default to LINEAR16 for unknown formats
            return speech.RecognitionConfig.AudioEncoding.LINEAR16

        def _build_recognition_config(
            self, params: dict, speech, file: Union[str, BinaryIO]
        ):
            """Build Google Speech RecognitionConfig from parameters."""
            # Auto-detect encoding if not specified
            encoding = params.get("encoding")
            if encoding is None:
                encoding = self._detect_audio_encoding(file, speech)

            config_params = {
                "encoding": encoding,
                "sample_rate_hertz": params.get("sample_rate_hertz", 16000),
                "language_code": params.get("language_code", "en-US"),
                "enable_word_time_offsets": True,
                "enable_word_confidence": True,
                "enable_automatic_punctuation": params.get(
                    "enable_automatic_punctuation", True
                ),
                "model": params["model"],
            }

            for param in ["max_alternatives", "profanity_filter", "speech_contexts"]:
                if param in params:
                    config_params[param] = params[param]

            return speech.RecognitionConfig(**config_params)

        def _create_streaming_requests(
            self, speech, streaming_config, audio_data: bytes
        ):
            """Create streaming requests generator for Google Speech API."""

            def request_generator():
                chunk_size = 8192
                for i in range(0, len(audio_data), chunk_size):
                    chunk = audio_data[i : i + chunk_size]
                    yield speech.StreamingRecognizeRequest(audio_content=chunk)

            return request_generator()

        def _parse_google_response(self, response) -> TranscriptionResult:
            """Convert Google Speech-to-Text response to unified TranscriptionResult."""
            if not response.results or not response.results[0].alternatives:
                return TranscriptionResult(
                    text="", language=None, confidence=None, task="transcribe"
                )

            best_result = response.results[0]
            best_alternative = best_result.alternatives[0]
            text = best_alternative.transcript
            confidence = getattr(best_alternative, "confidence", None)

            words = []
            if hasattr(best_alternative, "words") and best_alternative.words:
                words = [
                    Word(
                        word=word.word,
                        start=(
                            word.start_time.total_seconds()
                            if hasattr(word, "start_time")
                            else 0.0
                        ),
                        end=(
                            word.end_time.total_seconds()
                            if hasattr(word, "end_time")
                            else 0.0
                        ),
                        confidence=getattr(word, "confidence", None),
                    )
                    for word in best_alternative.words
                ]

            alternatives = [
                Alternative(
                    transcript=alt.transcript,
                    confidence=getattr(alt, "confidence", None),
                )
                for alt in best_result.alternatives
            ]

            segments = []
            if words:
                segments = [
                    Segment(
                        id=0,
                        seek=0,
                        start=words[0].start,
                        end=words[-1].end,
                        text=text,
                        tokens=[],
                        temperature=0.0,
                        avg_logprob=0.0,
                        compression_ratio=0.0,
                        no_speech_prob=0.0,
                    )
                ]

            return TranscriptionResult(
                text=text,
                language=None,
                confidence=confidence,
                task="transcribe",
                words=words or None,
                alternatives=alternatives or None,
                segments=segments or None,
            )