Full Chain Example¶

import requests
import json
from typing import Dict, Any

# Configuration
API_URL = "http://localhost:8003/api"
ADMIN_KEY = "sk-admin-key-change-me"  # For org/team creation

class SaaSLiteLLMClient:
    """Complete client for SaaS LiteLLM platform"""

    def __init__(self, api_url: str = "http://localhost:8003/api"):
        self.api_url = api_url
        self.admin_headers = {
            "Authorization": f"Bearer {ADMIN_KEY}",
            "Content-Type": "application/json"
        }
        self.team_headers = None
        self.team_id = None
        self.virtual_key = None

    # ========================================================================
    # Step 1: Organization Management
    # ========================================================================

    def create_organization(self, org_id: str, name: str, metadata: Dict = None) -> Dict[str, Any]:
        """Create a new organization"""
        response = requests.post(
            f"{self.api_url}/organizations/create",
            headers=self.admin_headers,
            json={
                "organization_id": org_id,
                "name": name,
                "metadata": metadata or {}
            }
        )

        if response.status_code == 400 and "already exists" in response.text:
            print(f"Organization {org_id} already exists, fetching details...")
            return self.get_organization(org_id)

        response.raise_for_status()
        return response.json()

    def get_organization(self, org_id: str) -> Dict[str, Any]:
        """Get organization details"""
        response = requests.get(
            f"{self.api_url}/organizations/{org_id}",
            headers=self.admin_headers
        )
        response.raise_for_status()
        return response.json()

    # ========================================================================
    # Step 2: Model Group Management
    # ========================================================================

    def create_model_group(
        self,
        group_name: str,
        display_name: str,
        description: str,
        models: list
    ) -> Dict[str, Any]:
        """Create a model group with prioritized models"""
        response = requests.post(
            f"{self.api_url}/model-groups/create",
            headers=self.admin_headers,
            json={
                "group_name": group_name,
                "display_name": display_name,
                "description": description,
                "models": models
            }
        )

        if response.status_code == 400 and "already exists" in response.text:
            print(f"Model group {group_name} already exists")
            return self.get_model_group(group_name)

        response.raise_for_status()
        return response.json()

    def get_model_group(self, group_name: str) -> Dict[str, Any]:
        """Get model group details"""
        response = requests.get(
            f"{self.api_url}/model-groups/{group_name}",
            headers=self.admin_headers
        )
        response.raise_for_status()
        return response.json()

    # ========================================================================
    # Step 3: Team Management
    # ========================================================================

    def create_team(
        self,
        organization_id: str,
        team_id: str,
        team_alias: str,
        model_groups: list,
        credits_allocated: int
    ) -> Dict[str, Any]:
        """Create team with LiteLLM integration and credits"""
        response = requests.post(
            f"{self.api_url}/teams/create",
            headers=self.admin_headers,
            json={
                "organization_id": organization_id,
                "team_id": team_id,
                "team_alias": team_alias,
                "model_groups": model_groups,
                "credits_allocated": credits_allocated
            }
        )

        if response.status_code == 400 and "already exists" in response.text:
            print(f"Team {team_id} already exists, fetching details...")
            return self.get_team(team_id)

        response.raise_for_status()
        result = response.json()

        # Store virtual key for future requests
        self.virtual_key = result.get("virtual_key")
        self.team_id = team_id
        self.team_headers = {
            "Authorization": f"Bearer {self.virtual_key}",
            "Content-Type": "application/json"
        }

        return result

    def get_team(self, team_id: str) -> Dict[str, Any]:
        """Get team details"""
        response = requests.get(
            f"{self.api_url}/teams/{team_id}",
            headers=self.admin_headers
        )
        response.raise_for_status()
        return response.json()

    def set_virtual_key(self, virtual_key: str, team_id: str):
        """Set virtual key for team operations"""
        self.virtual_key = virtual_key
        self.team_id = team_id
        self.team_headers = {
            "Authorization": f"Bearer {virtual_key}",
            "Content-Type": "application/json"
        }

    # ========================================================================
    # Step 4: Credit Management
    # ========================================================================

    def get_credit_balance(self) -> Dict[str, Any]:
        """Get credit balance for authenticated team"""
        if not self.team_headers:
            raise ValueError("No team authenticated. Call create_team or set_virtual_key first.")

        response = requests.get(
            f"{self.api_url}/credits/teams/{self.team_id}/balance",
            headers=self.team_headers
        )
        response.raise_for_status()
        return response.json()

    def add_credits(self, credits: int, reason: str = "Manual allocation") -> Dict[str, Any]:
        """Add credits to team"""
        if not self.team_headers:
            raise ValueError("No team authenticated")

        response = requests.post(
            f"{self.api_url}/credits/teams/{self.team_id}/add",
            headers=self.team_headers,
            json={
                "credits": credits,
                "reason": reason
            }
        )
        response.raise_for_status()
        return response.json()

    # ========================================================================
    # Step 5: Job & LLM Call Workflow
    # ========================================================================

    def create_job(self, job_type: str, metadata: Dict = None) -> str:
        """Create a new job"""
        if not self.team_headers:
            raise ValueError("No team authenticated")

        response = requests.post(
            f"{self.api_url}/jobs/create",
            headers=self.team_headers,
            json={
                "team_id": self.team_id,
                "job_type": job_type,
                "metadata": metadata or {}
            }
        )
        response.raise_for_status()
        result = response.json()
        return result["job_id"]

    def make_llm_call(
        self,
        job_id: str,
        model_group: str,
        messages: list,
        purpose: str = None,
        temperature: float = 0.7,
        max_tokens: int = None
    ) -> Dict[str, Any]:
        """Make an LLM call within a job"""
        if not self.team_headers:
            raise ValueError("No team authenticated")

        payload = {
            "model_group": model_group,
            "messages": messages,
            "temperature": temperature
        }

        if purpose:
            payload["purpose"] = purpose
        if max_tokens:
            payload["max_tokens"] = max_tokens

        response = requests.post(
            f"{self.api_url}/jobs/{job_id}/llm-call",
            headers=self.team_headers,
            json=payload
        )
        response.raise_for_status()
        return response.json()

    def complete_job(
        self,
        job_id: str,
        status: str = "completed",
        error_message: str = None
    ) -> Dict[str, Any]:
        """Complete a job and get cost summary"""
        if not self.team_headers:
            raise ValueError("No team authenticated")

        payload = {"status": status}
        if error_message:
            payload["error_message"] = error_message

        response = requests.post(
            f"{self.api_url}/jobs/{job_id}/complete",
            headers=self.team_headers,
            json=payload
        )
        response.raise_for_status()
        return response.json()

    def get_job(self, job_id: str) -> Dict[str, Any]:
        """Get job details"""
        if not self.team_headers:
            raise ValueError("No team authenticated")

        response = requests.get(
            f"{self.api_url}/jobs/{job_id}",
            headers=self.team_headers
        )
        response.raise_for_status()
        return response.json()


# ============================================================================
# Complete End-to-End Example
# ============================================================================

def full_workflow_example():
    """
    Complete workflow:
    1. Create organization
    2. Create model groups
    3. Create team
    4. Create job
    5. Make LLM calls
    6. Complete job
    7. Check credits
    """

    client = SaaSLiteLLMClient()

    print("="*70)
    print("SAAS LITELLM - FULL WORKFLOW EXAMPLE")
    print("="*70)

    # Step 1: Create Organization
    print("\n[1/7] Creating Organization...")
    org = client.create_organization(
        org_id="acme-corp",
        name="Acme Corporation",
        metadata={"industry": "technology", "tier": "enterprise"}
    )
    print(f"✓ Organization created: {org['organization_id']}")

    # Step 2: Create Model Groups
    print("\n[2/7] Creating Model Groups...")

    # Fast model group for quick tasks
    fast_group = client.create_model_group(
        group_name="ChatFast",
        display_name="Fast Chat Model",
        description="Quick responses for simple queries",
        models=[
            {"model_name": "gpt-3.5-turbo", "priority": 0},
            {"model_name": "gpt-4o-mini", "priority": 1}
        ]
    )
    print(f"✓ Model group created: {fast_group['group_name']}")

    # Advanced model group for complex tasks
    advanced_group = client.create_model_group(
        group_name="ChatAdvanced",
        display_name="Advanced Chat Model",
        description="Powerful model for complex reasoning",
        models=[
            {"model_name": "gpt-4", "priority": 0},
            {"model_name": "gpt-4-turbo", "priority": 1}
        ]
    )
    print(f"✓ Model group created: {advanced_group['group_name']}")

    # Step 3: Create Team
    print("\n[3/7] Creating Team with Credits...")
    team = client.create_team(
        organization_id="acme-corp",
        team_id="engineering-team",
        team_alias="Engineering Team",
        model_groups=["ChatFast", "ChatAdvanced"],
        credits_allocated=100
    )
    print(f"✓ Team created: {team['team_id']}")
    print(f"  Virtual Key: {team['virtual_key'][:50]}...")
    print(f"  Credits: {team['credits_allocated']}")
    print(f"  Model Groups: {', '.join(team['model_groups'])}")

    # Step 4: Check Initial Credits
    print("\n[4/7] Checking Initial Credit Balance...")
    credits = client.get_credit_balance()
    print(f"✓ Credits available: {credits['credits_remaining']}/{credits['credits_allocated']}")

    # Step 5: Create Job
    print("\n[5/7] Creating Job...")
    job_id = client.create_job(
        job_type="customer_support",
        metadata={
            "customer_id": "customer_12345",
            "priority": "high"
        }
    )
    print(f"✓ Job created: {job_id}")

    # Step 6: Make LLM Calls
    print("\n[6/7] Making LLM Calls...")

    # First call - initial response
    print("  → Call 1: Initial customer query...")
    response1 = client.make_llm_call(
        job_id=job_id,
        model_group="ChatFast",
        messages=[
            {
                "role": "system",
                "content": "You are a helpful customer support assistant."
            },
            {
                "role": "user",
                "content": "How do I reset my password?"
            }
        ],
        purpose="initial_response"
    )
    print(f"    Response: {response1['response']['content'][:100]}...")
    print(f"    Tokens: {response1['metadata']['tokens_used']}")
    print(f"    Latency: {response1['metadata']['latency_ms']}ms")

    # Second call - follow-up
    print("\n  → Call 2: Follow-up question...")
    response2 = client.make_llm_call(
        job_id=job_id,
        model_group="ChatFast",
        messages=[
            {
                "role": "system",
                "content": "You are a helpful customer support assistant."
            },
            {
                "role": "user",
                "content": "How do I reset my password?"
            },
            {
                "role": "assistant",
                "content": response1['response']['content']
            },
            {
                "role": "user",
                "content": "I didn't receive the reset email."
            }
        ],
        purpose="follow_up"
    )
    print(f"    Response: {response2['response']['content'][:100]}...")
    print(f"    Tokens: {response2['metadata']['tokens_used']}")

    # Third call - complex analysis using advanced model
    print("\n  → Call 3: Complex analysis...")
    response3 = client.make_llm_call(
        job_id=job_id,
        model_group="ChatAdvanced",
        messages=[
            {
                "role": "user",
                "content": "Analyze this customer interaction and suggest improvements to our password reset process."
            }
        ],
        purpose="analysis",
        temperature=0.3
    )
    print(f"    Response: {response3['response']['content'][:100]}...")
    print(f"    Tokens: {response3['metadata']['tokens_used']}")

    # Step 7: Complete Job
    print("\n[7/7] Completing Job...")
    completion = client.complete_job(
        job_id=job_id,
        status="completed"
    )

    print(f"✓ Job completed: {completion['job_id']}")
    print(f"\n  Cost Summary:")
    print(f"    Total calls: {completion['costs']['total_calls']}")
    print(f"    Successful: {completion['costs']['successful_calls']}")
    print(f"    Failed: {completion['costs']['failed_calls']}")
    print(f"    Total tokens: {completion['costs']['total_tokens']}")
    print(f"    Total cost: ${completion['costs']['total_cost_usd']:.6f}")
    print(f"    Avg latency: {completion['costs']['avg_latency_ms']}ms")
    print(f"    Credit applied: {completion['costs']['credit_applied']}")
    print(f"    Credits remaining: {completion['costs']['credits_remaining']}")

    # Final credit check
    print("\n[FINAL] Credit Balance After Job...")
    final_credits = client.get_credit_balance()
    print(f"✓ Credits: {final_credits['credits_remaining']}/{final_credits['credits_allocated']}")
    print(f"  Credits used: {final_credits['credits_used']}")

    print("\n" + "="*70)
    print("WORKFLOW COMPLETED SUCCESSFULLY!")
    print("="*70)

    return {
        "organization": org,
        "team": team,
        "job_id": job_id,
        "completion": completion
    }


if __name__ == "__main__":
    try:
        result = full_workflow_example()
        print("\n✅ All steps completed successfully!")
    except requests.exceptions.ConnectionError as e:
        print(f"\n❌ Connection Error: Could not connect to API")
        print("Make sure SaaS API is running: http://localhost:8003")
    except Exception as e:
        print(f"\n❌ Error: {str(e)}")
        import traceback
        traceback.print_exc()