Tutorial: Running inference with Kimi K2 using W&B Inference

pip install openai wandb weave

# Set environment variables in your notebook
%env WANDB_API_KEY=your-wandb-api-key-here
%env WANDB_PROJECT=your-team/your-project

export WANDB_API_KEY="your-wandb-api-key-here"
export WANDB_PROJECT="your-team/your-project"

import os

# Set environment variables programmatically
os.environ["WANDB_API_KEY] = "your-wandb-api-key-here"
os.environ["WANDB_PROJECT"] = "your-team/your-project"

import openai
import weave
import os

# Initialize Weave for automatic logging and monitoring
# This will use the WANDB_PROJECT environment variable
weave.init(os.getenv("WANDB_PROJECT"))

# Create the OpenAI client configured for W&B Inference
client = openai.OpenAI(
    # W&B Inference endpoint
    base_url='https://api.inference.wandb.ai/v1',
    
    # Use environment variable for API key (your W&B API key)
    api_key=os.getenv("WANDB_API_KEY"),
    
    # Specify your team and project for usage tracking
    project=os.getenv("WANDB_PROJECT"),
)

# Make a request to Kimi K2
response = client.chat.completions.create(
    model="moonshotai/Kimi-K2-Instruct",
    messages=[
        {"role": "system", "content": "You are a helpful assistant specialized in providing detailed explanations."},
        {"role": "user", "content": "Explain quantum computing in simple terms."}
    ],
    temperature=0.7,
    max_tokens=1000
)

# Print the response
print(response.choices[0].message.content)

response = client.chat.completions.create(
    model="moonshotai/Kimi-K2-Instruct",
    messages=[
        {"role": "system", "content": "You are a creative writing assistant."},
        {"role": "user", "content": "Write a short story about time travel."}
    ],
    temperature=0.8,        # Controls creativity (0.0=focused, 1.0=creative)
    max_tokens=2000,        # Maximum response length
    top_p=0.9,             # Nucleus sampling (0.1=focused, 1.0=diverse)
    frequency_penalty=0.1,  # Reduces repetitive phrases (0.0-2.0)
    presence_penalty=0.1,   # Encourages topic diversity (0.0-2.0)
)

response = client.chat.completions.create(
    model="moonshotai/Kimi-K2-Instruct",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell me a comprehensive story about space exploration."}
    ],
    stream=True,
    temperature=0.7
)

# Print the response as it streams
print("Response: ", end="")
for chunk in response:
    if chunk.choices[0].delta.content is not None:
        print(chunk.choices[0].delta.content, end="", flush=True)
print("\n")

# Real example: Analyzing a technical specification
technical_doc = """
API Documentation: RESTful Web Services v2.1

OVERVIEW
This document outlines the complete specification for our RESTful web services API, 
including authentication, endpoints, request/response formats, error handling, 
and rate limiting policies.

AUTHENTICATION
All API requests must include a valid Bearer token in the Authorization header:
Authorization: Bearer <your-access-token>

Tokens expire after 24 hours and must be renewed using the /auth/refresh endpoint.

ENDPOINTS

1. User Management
GET /api/v2/users - Retrieve user list
POST /api/v2/users - Create new user
PUT /api/v2/users/{id} - Update existing user
DELETE /api/v2/users/{id} - Delete user

2. Data Operations  
GET /api/v2/data - Query data with filtering
POST /api/v2/data - Submit new data
PUT /api/v2/data/{id} - Update existing data

RATE LIMITING
- Standard tier: 1000 requests/hour
- Premium tier: 10000 requests/hour
- Enterprise tier: Unlimited

ERROR HANDLING
All errors return JSON with 'error' and 'message' fields:
{
  "error": "INVALID_TOKEN",
  "message": "The provided token is expired or invalid"
}

[Document continues with detailed specifications...]
"""

response = client.chat.completions.create(
    model="moonshotai/Kimi-K2-Instruct",
    messages=[
        {
            "role": "system", 
            "content": "You are an expert API analyst. Provide detailed technical analysis focusing on potential issues, missing information, and recommendations for improvement."
        },
        {
            "role": "user", 
            "content": f"Analyze this API documentation and provide a comprehensive review:\n\n{technical_doc}"
        }
    ],
    temperature=0.3,  # Lower temperature for focused technical analysis
    max_tokens=2000
)

print(response.choices[0].message.content)

mixed_language_code = """
// English comments with Chinese variable names
function 计算总价(商品列表, 折扣率) {
    let 总价 = 0;
    
    // Calculate base total
    for (const 商品 of 商品列表) {
        总价 += 商品.价格 * 商品.数量;
    }
    
    // Apply discount
    const 折扣金额 = 总价 * 折扣率;
    return 总价 - 折扣金额;
}

# Python function with mixed comments
def validate_用户输入(user_data):
    '''
    验证用户输入数据的完整性和有效性
    Validates user input data for completeness and validity
    '''
    required_fields = ['name', 'email', '年龄']
    
    for field in required_fields:
        if field not in user_data:
            raise ValueError(f"Missing required field: {field}")
    
    return True
"""

response = client.chat.completions.create(
    model="moonshotai/Kimi-K2-Instruct",
    messages=[
        {
            "role": "system",
            "content": "You are a senior software engineer fluent in multiple languages. Analyze code for best practices, potential issues, and provide refactoring suggestions."
        },
        {
            "role": "user",
            "content": f"Please review this multilingual code and suggest improvements for maintainability:\n\n{mixed_language_code}"
        }
    ],
    temperature=0.2
)

print(response.choices[0].message.content)

response = client.chat.completions.create(
    model="moonshotai/Kimi-K2-Instruct",
    messages=[
        {
            "role": "system",
            "content": "You are a strategic business analyst. Break down complex problems into logical steps and provide actionable recommendations."
        },
        {
            "role": "user",
            "content": """
            Our SaaS company is experiencing a 15% monthly churn rate. Our customer acquisition cost (CAC) is $150, 
            average customer lifetime value (CLV) is $800, and monthly recurring revenue per customer is $50. 
            We have 2,000 active customers and are spending $60,000/month on marketing.
            
            Please analyze this situation and provide a comprehensive strategy to improve our metrics, 
            including specific actions, expected timelines, and success metrics.
            """
        }
    ],
    temperature=0.4,  # Balanced creativity and focus for strategic thinking
    max_tokens=2500
)

print(response.choices[0].message.content)

import weave

@weave.op()
def analyze_customer_feedback(feedback_text, sentiment_threshold=0.5):
    """
    Analyzes customer feedback and categorizes sentiment.
    This function will be tracked with custom metadata.
    """
    response = client.chat.completions.create(
        model="moonshotai/Kimi-K2-Instruct",
        messages=[
            {
                "role": "system", 
                "content": f"You are a customer feedback analyst. Analyze sentiment and provide a score from -1 (very negative) to 1 (very positive). Also categorize the main topics mentioned."
            },
            {
                "role": "user", 
                "content": f"Analyze this customer feedback: {feedback_text}"
            }
        ],
        temperature=0.1  # Consistent analysis
    )
    return response.choices[0].message.content

# This call will be logged with function name, parameters, and execution context
result = analyze_customer_feedback(
    "The new update is confusing and slow. I can't find the features I used daily.",
    sentiment_threshold=0.3
)
print(result)