bitcloud2/Final_Assignment_Template

1import os2import time3 4import httpx5import warnings6from typing import List, Dict, Optional7from smolagents import ApiModel, ChatMessage8 9 10class GeminiApiModel(ApiModel):11    """12    ApiModel implementation using the Google Gemini API via direct HTTP requests.13    """14 15    def __init__(16        self,17        model_id: str = "gemini-pro",18        api_key: Optional[str] = None,19        **kwargs,20    ):21        """22        Initializes the GeminiApiModel.23 24        Args:25            model_id (str): The Gemini model ID to use (e.g., "gemini-pro").26            api_key (str, optional): Google AI Studio API key. Defaults to GEMINI_API_KEY environment variable.27            **kwargs: Additional keyword arguments passed to the parent ApiModel.28        """29        self.model_id = model_id30        # Prefer explicitly passed key, fallback to environment variable31        self.api_key = api_key if api_key else os.environ.get("GEMINI_API_KEY")32        if not self.api_key:33            warnings.warn(34                "GEMINI_API_KEY not provided via argument or environment variable. API calls will likely fail.",35                UserWarning,36            )37        # Gemini API doesn't inherently support complex role structures or function calling like OpenAI.38        # We'll flatten messages for simplicity.39        super().__init__(40            model_id=model_id,41            flatten_messages_as_text=True,  # Flatten messages to a single text prompt42            **kwargs,43        )44 45    def create_client(self):46        """No dedicated client needed as we use httpx directly."""47        return None  # Or potentially return httpx client if reused48 49    def __call__(50        self,51        messages: List[Dict[str, str]],52        stop_sequences: Optional[53            List[str]54        ] = None,  # Note: Gemini API might not support stop sequences directly here55        grammar: Optional[56            str57        ] = None,  # Note: Gemini API doesn't support grammar directly58        tools_to_call_from: Optional[59            List["Tool"]60        ] = None,  # Note: Basic Gemini API doesn't support tools61        **kwargs,62    ) -> ChatMessage:63        """64        Calls the Google Gemini API with the provided messages.65 66        Args:67            messages: A list of message dictionaries (e.g., [{'role': 'user', 'content': '...'}]).68            stop_sequences: Optional stop sequences (may not be supported).69            grammar: Optional grammar constraint (not supported).70            tools_to_call_from: Optional list of tools (not supported).71            **kwargs: Additional keyword arguments.72 73        Returns:74            A ChatMessage object containing the response.75        """76        if not self.api_key:77            raise ValueError("GEMINI_API_KEY is not set.")78 79        # Prepare the prompt by concatenating message content80        # The Gemini Pro basic API expects a simple text prompt.81        prompt = self._messages_to_prompt(messages)82        prompt += (83            "\n\n"84            + "If you have a result from a web search that looks helpful, please use httpx to get the HTML from the URL listed."85            + "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."86        )87        # print(f"--- Gemini API prompt: ---\n{prompt}\n--- End of prompt ---")88 89        url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model_id}:generateContent?key={self.api_key}"90        headers = {"Content-Type": "application/json"}91        # Construct the payload according to Gemini API requirements92        data = {"contents": [{"parts": [{"text": prompt}]}]}93 94        # Add generation config if provided via kwargs (optional)95        generation_config = {}96        if "temperature" in kwargs:97            generation_config["temperature"] = kwargs["temperature"]98        if "max_output_tokens" in kwargs:99            generation_config["maxOutputTokens"] = kwargs["max_output_tokens"]100        # Add other relevant config parameters here if needed101 102        if generation_config:103            data["generationConfig"] = generation_config104 105        # Handle stop sequences if provided (basic support)106        # Note: This is a best-effort addition, check Gemini API docs for formal support107        if stop_sequences:108            if "generationConfig" not in data:109                data["generationConfig"] = {}110            # Assuming Gemini API might support 'stopSequences' in generationConfig111            data["generationConfig"]["stopSequences"] = stop_sequences112 113        raw_response = None114        try:115            response = httpx.post(116                url, headers=headers, json=data, timeout=120.0117            )  # Increased timeout118            time.sleep(6)  # Add delay to respect rate limits119            response.raise_for_status()120            response_json = response.json()121            raw_response = response_json  # Store raw response122 123            # Parse the response - adjust based on actual Gemini API structure124            if "candidates" in response_json and response_json["candidates"]:125                part = response_json["candidates"][0]["content"]["parts"][0]126                if "text" in part:127                    content = part["text"]128                    # Check for "FINAL ANSWER: " and extract the rest of the string129                    final_answer_marker = "FINAL ANSWER: "130                    if final_answer_marker in content:131                        content = content.split(final_answer_marker)[-1].strip()132 133                    # Simulate token counts if available, otherwise default to 0134                    # The basic generateContent API might not return usage directly in the main response135                    # It might be in safetyRatings or other metadata if enabled/available.136                    # Setting to 0 for now as it's not reliably present in the simplest call.137                    self.last_input_token_count = 0138                    self.last_output_token_count = 0139                    # If usage data becomes available in response_json, parse it here:140                    # e.g., if response_json.get("usageMetadata"):141                    #    self.last_input_token_count = response_json["usageMetadata"].get("promptTokenCount", 0)142                    #    self.last_output_token_count = response_json["usageMetadata"].get("candidatesTokenCount", 0)143 144                    return ChatMessage(145                        role="assistant", content=content, raw=raw_response146                    )147 148            # Handle cases where the expected response structure isn't found149            error_content = f"Error or unexpected response format: {response_json}"150            return ChatMessage(151                role="assistant", content=error_content, raw=raw_response152            )153 154        except httpx.RequestError as exc:155            error_content = (156                f"An error occurred while requesting {exc.request.url!r}: {exc}"157            )158            return ChatMessage(159                role="assistant", content=error_content, raw={"error": str(exc)}160            )161        except httpx.HTTPStatusError as exc:162            error_content = f"Error response {exc.response.status_code} while requesting {exc.request.url!r}: {exc.response.text}"163            return ChatMessage(164                role="assistant",165                content=error_content,166                raw={167                    "error": str(exc),168                    "status_code": exc.response.status_code,169                    "response_text": exc.response.text,170                },171            )172        except Exception as e:173            error_content = f"An unexpected error occurred: {e}"174            return ChatMessage(175                role="assistant", content=error_content, raw={"error": str(e)}176            )177 178    def _messages_to_prompt(self, messages: List[Dict[str, str]]) -> str:179        """Converts a list of messages into a single string prompt."""180        # Simple concatenation, could be more sophisticated based on roles if needed181        # Ensure we handle cases where 'content' might not be a string (though it should be)182        return "\n".join([str(msg.get("content", "")) for msg in messages])183