220 lines
7.0 KiB
Plaintext
220 lines
7.0 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# IBM Watsonx AI Streaming with AgentOps\n",
|
|
"\n",
|
|
"This notebook demonstrates how to use IBM Watsonx AI for streaming text generation and streaming chat completion with AgentOps instrumentation."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup\n",
|
|
"\n",
|
|
"First, let's import the necessary libraries and initialize AgentOps:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": "import agentops\nfrom ibm_watsonx_ai import Credentials\nfrom ibm_watsonx_ai.foundation_models import ModelInference\nfrom dotenv import load_dotenv\nimport os\n\n# Load environment variables\nload_dotenv()\nos.environ[\"AGENTOPS_API_KEY\"] = os.getenv(\"AGENTOPS_API_KEY\", \"your_api_key_here\")\n\n# Initialize AgentOps\nagentops.init(tags=[\"watsonx-streaming\", \"agentops-example\"])"
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Initialize IBM Watsonx AI Credentials\n",
|
|
"\n",
|
|
"To use IBM Watsonx AI, you need to set up your credentials and project ID."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": "# Initialize credentials - replace with your own API key\n# Best practice: Store API keys in environment variables\n# Ensure WATSONX_API_KEY is set in your .env file or environment\nos.environ[\"WATSONX_API_KEY\"] = os.getenv(\"WATSONX_API_KEY\", \"your_watsonx_api_key_here\")\n\ncredentials = Credentials(\n url=os.getenv(\"WATSONX_URL\", \"https://eu-de.ml.cloud.ibm.com\"),\n api_key=os.environ[\"WATSONX_API_KEY\"],\n)\n\n# Project ID for your IBM Watsonx project\nproject_id = os.getenv(\"WATSONX_PROJECT_ID\", \"your-project-id-here\")"
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Initialize Models\n",
|
|
"\n",
|
|
"Let's initialize models for our streaming examples:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Initialize text generation model\n",
|
|
"gen_model = ModelInference(model_id=\"google/flan-ul2\", credentials=credentials, project_id=project_id)\n",
|
|
"\n",
|
|
"# Initialize chat model\n",
|
|
"chat_model = ModelInference(\n",
|
|
" model_id=\"meta-llama/llama-3-3-70b-instruct\", credentials=credentials, project_id=project_id\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Streaming Text Generation\n",
|
|
"\n",
|
|
"Let's use IBM Watsonx AI to generate streaming text:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Streaming text generation\n",
|
|
"prompt = \"List 3 benefits of machine learning:\"\n",
|
|
"stream_response = gen_model.generate_text_stream(prompt)\n",
|
|
"\n",
|
|
"print(\"Streaming Response:\")\n",
|
|
"full_stream_response = \"\"\n",
|
|
"for chunk in stream_response:\n",
|
|
" if isinstance(chunk, str):\n",
|
|
" print(chunk, end=\"\", flush=True)\n",
|
|
" full_stream_response += chunk\n",
|
|
"print(\"\\n\\nComplete Response:\")\n",
|
|
"print(full_stream_response)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Streaming Chat Completion\n",
|
|
"\n",
|
|
"Now, let's try streaming chat completion:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Format messages for chat\n",
|
|
"chat_stream_messages = [\n",
|
|
" {\"role\": \"system\", \"content\": \"You are a concise assistant.\"},\n",
|
|
" {\"role\": \"user\", \"content\": \"Explain the concept of photosynthesis in one sentence.\"},\n",
|
|
"]\n",
|
|
"\n",
|
|
"# Get streaming chat response\n",
|
|
"chat_stream_response_gen = chat_model.chat_stream(messages=chat_stream_messages)\n",
|
|
"\n",
|
|
"print(\"Chat Stream Response:\")\n",
|
|
"full_chat_stream_response = \"\"\n",
|
|
"for chunk in chat_stream_response_gen:\n",
|
|
" try:\n",
|
|
" # Check structure based on SDK docstring example\n",
|
|
" if chunk and \"choices\" in chunk and chunk[\"choices\"]:\n",
|
|
" delta = chunk[\"choices\"][0].get(\"delta\", {})\n",
|
|
" content_chunk = delta.get(\"content\")\n",
|
|
" if content_chunk:\n",
|
|
" print(content_chunk, end=\"\", flush=True)\n",
|
|
" full_chat_stream_response += content_chunk\n",
|
|
" except Exception as e:\n",
|
|
" print(f\"Error processing chat stream chunk: {e}, Chunk: {chunk}\")\n",
|
|
"\n",
|
|
"print(\"\\n\\nComplete Chat Response:\")\n",
|
|
"print(full_chat_stream_response)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Another Streaming Chat Example\n",
|
|
"\n",
|
|
"Let's try another example with a more complex query:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# New chat messages for streaming\n",
|
|
"chat_stream_messages = [\n",
|
|
" {\"role\": \"system\", \"content\": \"You are a helpful assistant that provides step-by-step explanations.\"},\n",
|
|
" {\"role\": \"user\", \"content\": \"Explain how to make a simple chocolate cake.\"},\n",
|
|
"]\n",
|
|
"\n",
|
|
"# Get streaming chat response\n",
|
|
"chat_stream_response_gen = chat_model.chat_stream(messages=chat_stream_messages)\n",
|
|
"\n",
|
|
"print(\"Chat Stream Response:\")\n",
|
|
"full_chat_stream_response = \"\"\n",
|
|
"for chunk in chat_stream_response_gen:\n",
|
|
" try:\n",
|
|
" if chunk and \"choices\" in chunk and chunk[\"choices\"]:\n",
|
|
" delta = chunk[\"choices\"][0].get(\"delta\", {})\n",
|
|
" content_chunk = delta.get(\"content\")\n",
|
|
" if content_chunk:\n",
|
|
" print(content_chunk, end=\"\", flush=True)\n",
|
|
" full_chat_stream_response += content_chunk\n",
|
|
" except Exception as e:\n",
|
|
" print(f\"Error processing chat stream chunk: {e}, Chunk: {chunk}\")\n",
|
|
"\n",
|
|
"print(\"\\n\\nComplete Chat Response:\")\n",
|
|
"print(full_chat_stream_response)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Clean Up\n",
|
|
"\n",
|
|
"Finally, let's close the persistent connection with the models:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Close connections\n",
|
|
"gen_model.close_persistent_connection()\n",
|
|
"chat_model.close_persistent_connection()"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": ".venv",
|
|
"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.12"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 4
|
|
} |