VertexAIServerModel Integration to Enable Openinference instrumentation with Phoenix

[dmbrundage]

I am not sure if this is a feature request or if I am just not understanding how to properly integrate vertex AI with smolagents and opentelemetry. I created a custom model class extending the base model and didn't have any issues until I tried to add openinference instrumentation with phoenix and SmolagentsInstrumentor I get an AttributeError because my custom model server isnt in smolagents.

AttributeError: module 'smolagents' has no attribute 'VertexAIServerModel'. Did you mean: 'OpenAIServerModel'?

I would like the ability to integrate VertexAIServerModel similar to the OpenAIServerModel to allow to select models from the model garden/registry in GCP. I added the custom model class to models.py and then I was able to use the model as expected and telemetry was successful in pheonix. Is this even a smolagents change? Sorry I am very new to this ecosystem.

code taken from this notebook: https://github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/use-cases/vertex_ai_deepseek_smolagents.ipynb

class VertexAIServerModel(Model):
    """This model connects to a Vertex AI-compatible API server."""

    def __init__(
        self, model_id: str, project_id: str, location: str, endpoint_id: str, **kwargs
    ):
        #  Try to import dependencies
        try:
            from google.auth import default
        except ModuleNotFoundError:
            raise ModuleNotFoundError(
                "Please install 'openai, google-auth and requests' extra to use VertexAIGeminiModel as described in the official documentation: https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/call-vertex-using-openai-library"
            ) from None

        # Initialize parent class with any additional keyword arguments
        super().__init__(**kwargs)
        self.model_id = model_id
        self.project_id = project_id
        self.location = location
        self.endpoint_id = endpoint_id
        self.kwargs = kwargs
        self._refresh_task = None

        # Initialize credentials and set up Google Cloud authentication with required permissions
        self.credentials, _ = default(
            scopes=["https://www.googleapis.com/auth/cloud-platform"]
        )
        self._refresh_token()
        self._setup_client()
        self._start_refresh_loop()

    def __call__(
        self,
        messages: list[dict[str, str]],
        **kwargs,
    ) -> ChatMessage:

        # Prepare the API call parameters
        completion_kwargs = self._prepare_completion_kwargs(
            messages=messages,
            model=self.model_id,
            **self.kwargs,
        )

        # Make the API call to Vertex AI
        response = self.client.chat.completions.create(**completion_kwargs)
        self.last_input_token_count = response.usage.prompt_tokens
        self.last_output_token_count = response.usage.completion_tokens

        # Convert API response to ChatMessage format
        message = ChatMessage.from_dict(
            response.choices[0].message.model_dump(
                include={"role", "content", "tool_calls"}
            )
        )
        return message

    def _refresh_token(self):
        """Refresh the Google Cloud token"""
        try:
            self.credentials.refresh(google.auth.transport.requests.Request())
            self._setup_client()
        except Exception as e:
            print(f"Token refresh failed: {e}")

    def _setup_client(self):
        """Setup OpenAI client with current credentials"""
        self.client = openai.OpenAI(
            base_url=f"https://{self.location}-aiplatform.googleapis.com/v1beta1/projects/{self.project_id}/locations/{self.location}/endpoints/{self.endpoint_id}",
            api_key=self.credentials.token,
        )

    def _start_refresh_loop(self):
        """Start the token refresh loop"""

        def refresh_loop():
            while True:
                time.sleep(3600)
                self._refresh_token()

        self._refresh_thread = threading.Thread(target=refresh_loop, daemon=True)
        self._refresh_thread.start()

[IlyaGusev]

Hey, you probably need to do a thing similar to the one I did in this PR: #1040; just add VertexAIServerModel to the exports of the module.