{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "580c85ac",
   "metadata": {},
   "source": [
    "# Google Generative AI Example with AgentOps\n",
    "\n",
    "This notebook demonstrates how to use AgentOps with Google's Generative AI package for observing both synchronous and streaming text generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d52c7ce5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instal necessary packages\n",
    "%pip install agentops\n",
    "%pip install google-genai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d731924a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google import genai\n",
    "import agentops\n",
    "from dotenv import load_dotenv\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a94545c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "\n",
    "os.environ[\"AGENTOPS_API_KEY\"] = os.getenv(\"AGENTOPS_API_KEY\", \"your_api_key_here\")\n",
    "os.environ[\"GEMINI_API_KEY\"] = os.getenv(\"GEMINI_API_KEY\", \"your_gemini_api_key_here\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d632fe48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize AgentOps and Gemini client\n",
    "agentops.init(tags=[\"gemini-example\", \"agentops-example\"])\n",
    "client = genai.Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3923b6b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test synchronous generation\n",
    "print(\"Testing synchronous generation:\")\n",
    "response = client.models.generate_content(model=\"gemini-1.5-flash\", contents=\"What are the three laws of robotics?\")\n",
    "print(response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da54e521",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test streaming generation\n",
    "print(\"\\nTesting streaming generation:\")\n",
    "response_stream = client.models.generate_content_stream(\n",
    "    model=\"gemini-1.5-flash\", contents=\"Explain the concept of machine learning in simple terms.\"\n",
    ")\n",
    "\n",
    "for chunk in response_stream:\n",
    "    print(chunk.text, end=\"\")\n",
    "print()  # Add newline after streaming output\n",
    "\n",
    "# Test another synchronous generation\n",
    "print(\"\\nTesting another synchronous generation:\")\n",
    "response = client.models.generate_content(\n",
    "    model=\"gemini-1.5-flash\", contents=\"What is the difference between supervised and unsupervised learning?\"\n",
    ")\n",
    "print(response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbb2a59c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example of token counting\n",
    "print(\"\\nTesting token counting:\")\n",
    "token_response = client.models.count_tokens(\n",
    "    model=\"gemini-1.5-flash\", contents=\"This is a test sentence to count tokens.\"\n",
    ")\n",
    "print(f\"Token count: {token_response.total_tokens}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "agentops (3.11.11)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}