from google import generativeai as genai # Renamed for clarity as per official docs
import agentops
from dotenv import load_dotenv
import os

load_dotenv()

os.environ["AGENTOPS_API_KEY"] = os.getenv("AGENTOPS_API_KEY", "your_agentops_api_key_here")
os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY", "your_gemini_api_key_here")

# Configure the GenAI client with the API key
genai.configure(api_key=os.environ["GEMINI_API_KEY"])

# Initialize AgentOps and Gemini client
agentops.init(tags=["gemini-example", "agentops-example"])
# The client is configured globally via genai.configure()

# Test synchronous generation
print("Testing synchronous generation:")
model = genai.GenerativeModel('gemini-1.5-flash') # Create model instance
response = model.generate_content("What are the three laws of robotics?")
print(response.text)

# Test streaming generation
print("\nTesting streaming generation:")
model = genai.GenerativeModel('gemini-1.5-flash') # Create model instance
response_stream = model.generate_content(
    "Explain the concept of machine learning in simple terms.",
    stream=True
)

for chunk in response_stream:
    print(chunk.text, end="")
print()  # Add newline after streaming output

# Test another synchronous generation
print("\nTesting another synchronous generation:")
model = genai.GenerativeModel('gemini-1.5-flash') # Create model instance
response = model.generate_content(
    "What is the difference between supervised and unsupervised learning?"
)
print(response.text)

# Example of token counting
print("\nTesting token counting:")
model = genai.GenerativeModel('gemini-1.5-flash') # Create model instance
token_response = model.count_tokens("This is a test sentence to count tokens.")
print(f"Token count: {token_response.total_tokens}")

agentops.end_session("Success") # End the session