{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Weave tracing for Koog agents\n",
    "\n",
    "This notebook demonstrates how to trace Koog agents to W&B Weave using OpenTelemetry (OTLP).\n",
    "You will create a simple Koog `AIAgent`, enable the Weave exporter, run a prompt, and view\n",
    "rich traces in the Weave UI.\n",
    "\n",
    "For background, see Weave OpenTelemetry docs: https://weave-docs.wandb.ai/guides/tracking/otel/\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Prerequisites\n",
    "\n",
    "Before running the example, make sure you have:\n",
    "\n",
    "- A Weave/W&B account: https://wandb.ai\n",
    "- Your API key from https://wandb.ai/authorize exposed as an environment variable: `WEAVE_API_KEY`\n",
    "- Your Weave entity (team or user) name exposed as `WEAVE_ENTITY`\n",
    "  - Find it on your W&B dashboard: https://wandb.ai/home (left sidebar \"Teams\")\n",
    "- A project name exposed as `WEAVE_PROJECT_NAME` (if not set, this example uses `koog-tracing`)\n",
    "- An OpenAI API key exposed as `OPENAI_API_KEY` to run the Koog agent\n",
    "\n",
    "Example (macOS/Linux):\n",
    "```bash\n",
    "export WEAVE_API_KEY=...  # required by Weave\n",
    "export WEAVE_ENTITY=your-team-or-username\n",
    "export WEAVE_PROJECT_NAME=koog-tracing\n",
    "export OPENAI_API_KEY=...\n",
    "```\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Notebook setup\n",
    "\n",
    "We use the latest Kotlin Jupyter descriptors. If you have Koog preconfigured as a `%use` plugin,\n",
    "you can uncomment the line below.\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "%useLatestDescriptors\n",
    "//%use koog\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Create an agent and enable Weave tracing\n",
    "\n",
    "We construct a minimal `AIAgent` and install the `OpenTelemetry` feature with the Weave exporter.\n",
    "The exporter sends OTLP spans to Weave using your environment configuration:\n",
    "- `WEAVE_API_KEY` — authentication to Weave\n",
    "- `WEAVE_ENTITY` — which team/user owns the traces\n",
    "- `WEAVE_PROJECT_NAME` — the Weave project to store traces in\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import ai.koog.agents.core.agent.AIAgent\n",
    "import ai.koog.agents.features.opentelemetry.feature.OpenTelemetry\n",
    "import ai.koog.agents.features.opentelemetry.integration.weave.addWeaveExporter\n",
    "import ai.koog.prompt.executor.clients.openai.OpenAIModels\n",
    "import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor\n",
    "\n",
    "val entity = System.getenv()[\"WEAVE_ENTITY\"] ?: throw IllegalArgumentException(\"WEAVE_ENTITY is not set\")\n",
    "val projectName = System.getenv()[\"WEAVE_PROJECT_NAME\"] ?: \"koog-tracing\"\n",
    "\n",
    "val agent = AIAgent(\n",
    "    executor = simpleOpenAIExecutor(System.getenv(\"OPENAI_API_KEY\")),\n",
    "    llmModel = OpenAIModels.Reasoning.GPT4oMini,\n",
    "    systemPrompt = \"You are a code assistant. Provide concise code examples.\"\n",
    ") {\n",
    "    install(OpenTelemetry) {\n",
    "        addWeaveExporter(\n",
    "            weaveEntity = entity,\n",
    "            weaveProjectName = projectName\n",
    "        )\n",
    "    }\n",
    "}\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Run the agent and view traces in Weave\n",
    "\n",
    "Execute a simple prompt. After completion, open the printed link to view the trace in Weave.\n",
    "You should see spans for the agent’s run, model calls, and other instrumented operations.\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import kotlinx.coroutines.runBlocking\n",
    "\n",
    "println(\"Running agent with Weave tracing\")\n",
    "\n",
    "runBlocking {\n",
    "    val result = agent.run(\"Tell me a joke about programming\")\n",
    "    \"Result: $result\\nSee traces on https://wandb.ai/$entity/$projectName/weave/traces\"\n",
    "}\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Troubleshooting\n",
    "\n",
    "- If you don't see traces, verify `WEAVE_API_KEY`, `WEAVE_ENTITY`, and `WEAVE_PROJECT_NAME` are set in your environment.\n",
    "- Ensure your network allows outbound HTTPS to Weave's OTLP endpoint.\n",
    "- Confirm your OpenAI key is valid and the selected model is accessible from your account."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Kotlin",
   "language": "kotlin",
   "name": "kotlin"
  },
  "language_info": {
   "name": "kotlin",
   "version": "2.2.20-Beta2",
   "mimetype": "text/x-kotlin",
   "file_extension": ".kt",
   "pygments_lexer": "kotlin",
   "codemirror_mode": "text/x-kotlin",
   "nbconvert_exporter": ""
  },
  "ktnbPluginMetadata": {
   "projectDependencies": [
    "koog-agents.examples.main"
   ],
   "projectLibraries": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}