Create inference.py

inference.py

@@ -0,0 +1,100 @@
+from google import genai
+from google.genai import types
+import os
+
+# Global variable to hold the client
+client = None
+
+def initialize():
+    """
+    Initializes the Google Generative AI client.
+    """
+    global client
+    # It's a good practice to load the API key from an environment variable
+    api_key = os.environ.get("GEMINI_API_KEY")
+    if not api_key:
+        api_key = os.environ.get("GOOGLE_API_KEY")
+    
+    if not api_key:
+        raise ValueError("Neither GEMINI_API_KEY nor GOOGLE_API_KEY environment variable is set.")
+    
+    try:
+        client = genai.Client(api_key=api_key)
+        print("Google Generative AI client initialized.")
+    except Exception as e:
+        print(f"Error initializing Google Generative AI client: {e}")
+        raise
+
+def generate_content(prompt: str, model_name: str = None, allow_fallbacks: bool = True, generation_config: dict = None) -> str:
+    """
+    Generates content using the Google Generative AI model.
+
+    Args:
+        prompt: The prompt to send to the model.
+        model_name: The name of the model to use (e.g., "gemini-2.0-flash", "gemini-1.5-flash").
+                    If None, a default model will be used.
+        allow_fallbacks: (Currently not directly used by genai.Client.models.generate_content,
+                         but kept for compatibility with agent.py structure)
+        generation_config: A dictionary for generation parameters like temperature, max_output_tokens.
+
+    Returns:
+        The generated text content.
+    """
+    global client
+    if client is None:
+        # Attempt to initialize if not already done, though ideally initialize() should be called explicitly.
+        print("Client not initialized. Attempting to initialize now...")
+        initialize()
+        if client is None: # Check again after attempt
+             raise RuntimeError("Google Generative AI client is not initialized. Call initialize() first.")
+
+    # Default model if not specified
+    effective_model_name = model_name if model_name else "gemini-2.0-flash-lite"
+
+    # Prepare generation configuration for the API
+    config_obj = None
+    if generation_config:
+        config_params = {}
+        if 'temperature' in generation_config:
+            config_params['temperature'] = generation_config['temperature']
+        if 'max_output_tokens' in generation_config:
+            config_params['max_output_tokens'] = generation_config['max_output_tokens']
+        # Add other relevant parameters from generation_config as needed by the genai API
+        
+        if config_params:
+            config_obj = types.GenerateContentConfig(**config_params)
+
+    try:
+        response = client.models.generate_content(
+            model=effective_model_name,
+            contents=[prompt],  # Note: contents expects a list
+            config=config_obj
+        )
+        return response.text
+    except Exception as e:
+        print(f"Error during content generation: {e}")
+        # Depending on how agent.py handles errors, you might want to raise the exception
+        # or return a specific error message. For now, re-raising.
+        raise
+
+if __name__ == '__main__':
+    # Example usage (optional, for testing inference.py directly)
+    try:
+        # Make sure to set your GEMINI_API_KEY environment variable before running
+        # For example, in your terminal: $env:GEMINI_API_KEY="YOUR_API_KEY"
+        initialize()
+        if client:
+            sample_prompt = "Explain how AI works in a few words"
+            print(f"Sending prompt: '{sample_prompt}'")
+            config = {'temperature': 0.7, 'max_output_tokens': 50}
+            generated_text = generate_content(sample_prompt, generation_config=config)
+            print("\nGenerated text:")
+            print(generated_text)
+
+            sample_prompt_2 = "What is the capital of France?"
+            print(f"\nSending prompt: '{sample_prompt_2}'")
+            generated_text_2 = generate_content(sample_prompt_2, model_name="gemini-2.0-flash-lite") # Example with a different model
+            print("\nGenerated text:")
+            print(generated_text_2)
+    except Exception as e:
+        print(f"An error occurred: {e}")