multi provider genai gateway test

gistfile1.txt

"""
GenAI Gateway — End-User Quickstart
=====================================

This script shows how to use the Multi-Provider GenAI Gateway from your Python code.
Built on AWS with LiteLLM — routes to Bedrock models, self-hosted models, and more.

SETUP
-----
1. Get your API key from the LiteLLM UI:
   https://<your-gateway-url>/ui
   → Log in with SSO → Keys → Create Key

2. Set your key and gateway URL:
   export GATEWAY_KEY="sk-your-key-here"
   export GATEWAY_URL="https://your-gateway-url.cloudfront.net/v1"

3. Install dependencies:
   pip install openai boto3

WHAT THIS SCRIPT DEMONSTRATES
------------------------------
- Calling any model through a single unified endpoint
- Chat completions (Claude, Nova, Llama)
- Streaming responses
- Embeddings
- Multi-turn conversation with history
- Switching models with zero code changes
- AWS Bedrock SDK compatibility — change one line, no rewrite needed
- Cross-SDK session persistence via RDS (OpenAI SDK + Bedrock SDK share context)

No AWS credentials needed — the gateway handles all provider auth.
"""

import os
import sys
from openai import OpenAI

# ── Configuration ─────────────────────────────────────────────────────────────

# Replace with your gateway CloudFront URL
GATEWAY_URL = os.getenv("GATEWAY_URL", "https://YOUR-GATEWAY-URL.cloudfront.net/v1")

# Replace with your key from the LiteLLM UI, or set GATEWAY_KEY env var
API_KEY = os.getenv("GATEWAY_KEY", os.getenv("HMS_GATEWAY_KEY", "YOUR_KEY_HERE"))

if API_KEY == "YOUR_KEY_HERE":
    print("ERROR: Please set your API key.")
    print("    Either edit this script or run:")
    print("    export GATEWAY_KEY='sk-your-key-here'")
    sys.exit(1)

if "YOUR-GATEWAY-URL" in GATEWAY_URL:
    print("ERROR: Please set your gateway URL.")
    print("    Either edit this script or run:")
    print("    export GATEWAY_URL='https://your-gateway-url.cloudfront.net/v1'")
    sys.exit(1)

# Single client — works with all models
client = OpenAI(
    api_key=API_KEY,
    base_url=GATEWAY_URL,
)

# ── Available Models ───────────────────────────────────────────────────────────
# Update these with the model names registered in your gateway

MODELS = {
    # On-Prem — Self-hosted via vLLM on EC2 (VPC Peering)
    # Replace with your self-hosted model name
    "onprem-llama":  "hms-onprem/llama3-8b",

    # Anthropic Claude 4.x via Amazon Bedrock
    "claude-opus":   "anthropic.claude-opus-4-7",
    "claude-sonnet": "anthropic.claude-sonnet-4-6",
    "claude-haiku":  "anthropic.claude-haiku-4-5-20251001-v1:0",

    # Amazon Nova via Bedrock
    "nova-pro":      "amazon.nova-pro-v1:0",
    "nova-lite":     "amazon.nova-lite-v1:0",
    "nova-micro":    "amazon.nova-micro-v1:0",

    # Meta Llama 3.x via Bedrock
    "llama-70b":     "meta.llama3-3-70b-instruct-v1:0",
    "llama-8b":      "meta.llama3-1-8b-instruct-v1:0",

    # Embeddings
    "titan-embed":   "amazon.titan-embed-text-v2:0",
    "cohere-embed":  "cohere.embed-english-v3",
}


# ── Example 1: Simple chat completion ─────────────────────────────────────────

def example_simple_chat():
    print("\n" + "="*60)
    print("Example 1: Simple Chat — On-Prem Self-Hosted Model")
    print("="*60)

    response = client.chat.completions.create(
        model=MODELS["onprem-llama"],
        messages=[
            {"role": "user", "content": "What are the top 3 benefits of cloud computing? Be concise."}
        ],
    )

    print(f"Model  : {response.model}")
    print(f"Reply  : {response.choices[0].message.content}")
    print(f"Tokens : {response.usage.prompt_tokens} in / {response.usage.completion_tokens} out")


# ── Example 2: Switch models with zero code changes ───────────────────────────

def example_model_comparison():
    print("\n" + "="*60)
    print("Example 2: Same Question, Different Models")
    print("         (Cloud + On-Prem side by side)")
    print("="*60)

    question = "In one sentence, what is machine learning?"

    for label, model_id in [
        ("On-Prem / Llama 3.1 8B (vLLM)",  MODELS["onprem-llama"]),
        ("Claude Haiku 4.5 (Bedrock)",       MODELS["claude-haiku"]),
        ("Amazon Nova Micro (Bedrock)",       MODELS["nova-micro"]),
    ]:
        response = client.chat.completions.create(
            model=model_id,
            messages=[{"role": "user", "content": question}],
        )
        answer = response.choices[0].message.content.strip()
        print(f"\n[{label}]\n{answer}")


# ── Example 3: Streaming response ─────────────────────────────────────────────

def example_streaming():
    print("\n" + "="*60)
    print("Example 3: Streaming Response")
    print("="*60)
    print("Model: Claude Sonnet 4.6")
    print("Response: ", end="", flush=True)

    stream = client.chat.completions.create(
        model=MODELS["claude-sonnet"],
        messages=[
            {"role": "user", "content": "Write a two-sentence summary of why AI is important in healthcare."}
        ],
        stream=True,
    )

    for chunk in stream:
        content = chunk.choices[0].delta.content
        if content:
            print(content, end="", flush=True)
    print()


# ── Example 4: Multi-turn conversation ────────────────────────────────────────

def example_multi_turn():
    print("\n" + "="*60)
    print("Example 4: Multi-Turn Conversation")
    print("="*60)

    conversation = [
        {"role": "system", "content": "You are a helpful medical research assistant. Be concise."},
        {"role": "user",   "content": "What is CRISPR?"},
    ]

    # Turn 1
    response = client.chat.completions.create(
        model=MODELS["nova-pro"],
        messages=conversation,
    )
    reply1 = response.choices[0].message.content
    print(f"User  : What is CRISPR?")
    print(f"Model : {reply1}\n")

    # Add assistant reply and follow-up to conversation history
    conversation.append({"role": "assistant", "content": reply1})
    conversation.append({"role": "user", "content": "What are its main medical applications?"})

    # Turn 2 — model has full context
    response = client.chat.completions.create(
        model=MODELS["nova-pro"],
        messages=conversation,
    )
    reply2 = response.choices[0].message.content
    print(f"User  : What are its main medical applications?")
    print(f"Model : {reply2}")


# ── Example 5: Embeddings ──────────────────────────────────────────────────────

def example_embeddings():
    print("\n" + "="*60)
    print("Example 5: Text Embeddings (for RAG / semantic search)")
    print("="*60)

    texts = [
        "The patient presented with acute chest pain.",
        "Machine learning models require large datasets.",
        "Cardiac arrest requires immediate intervention.",
    ]

    response = client.embeddings.create(
        model=MODELS["titan-embed"],
        input=texts,
    )

    print(f"Model      : {response.model}")
    print(f"Dimensions : {len(response.data[0].embedding)}")
    print(f"Texts embedded: {len(response.data)}")

    # Simple cosine similarity to show semantic relatedness
    import math

    def cosine_sim(a, b):
        dot = sum(x*y for x, y in zip(a, b))
        mag_a = math.sqrt(sum(x**2 for x in a))
        mag_b = math.sqrt(sum(x**2 for x in b))
        return dot / (mag_a * mag_b)

    e = [d.embedding for d in response.data]
    sim_01 = cosine_sim(e[0], e[1])
    sim_02 = cosine_sim(e[0], e[2])

    print(f"\nSimilarity: '{texts[0][:40]}...'")
    print(f"  vs '{texts[1][:40]}...' → {sim_01:.3f} (low — different topics)")
    print(f"  vs '{texts[2][:40]}...' → {sim_02:.3f} (high — both medical)")


# ── Example 6: List available models ──────────────────────────────────────────

def example_list_models():
    print("\n" + "="*60)
    print("Example 6: List All Available Models")
    print("="*60)

    models = client.models.list()
    for m in models.data:
        print(f"  - {m.id}")


# ── Example 7: Bedrock SDK — zero code change migration ───────────────────────

def example_bedrock_sdk():
    """
    Shows how existing code using the AWS Bedrock SDK works unchanged
    by just pointing the endpoint_url at the gateway.

    Migration story: change ONE line, get full gateway benefits:
    unified auth, budget tracking, Redis caching, on-prem model routing.
    """
    print("\n" + "="*60)
    print("Example 7: AWS Bedrock SDK — Zero Code Change Migration")
    print("="*60)

    try:
        import boto3
        from botocore.config import Config
        from botocore import UNSIGNED
    except ImportError:
        print("  boto3 not installed. Run: pip install boto3")
        print("  Skipping this example.")
        return

    # The ONLY change from real Bedrock: endpoint_url points at the gateway
    # Everything else is identical to existing Bedrock SDK code
    bedrock = boto3.client(
        service_name="bedrock-runtime",
        region_name="us-east-1",
        endpoint_url=f"{GATEWAY_URL.replace('/v1', '')}/bedrock",
        config=Config(signature_version=UNSIGNED),
    )

    # Add auth header — same sk- key, same as OpenAI examples above
    def add_auth(request, **kwargs):
        request.headers["Authorization"] = f"Bearer {API_KEY}"

    bedrock.meta.events.register("request-created.*", add_auth)

    print("\n--- Non-streaming (Bedrock Converse API) ---")
    response = bedrock.converse(
        modelId="anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{
            "role": "user",
            "content": [{"text": "What is the role of AI in modern medicine? One sentence."}]
        }]
    )
    content = response["output"]["message"]["content"][0]["text"]
    usage   = response["usage"]
    print(f"Model   : anthropic.claude-haiku-4-5-20251001-v1:0")
    print(f"Reply   : {content}")
    print(f"Tokens  : {usage['inputTokens']} in / {usage['outputTokens']} out")

    print("\n--- Streaming (Bedrock ConverseStream API) ---")
    stream_resp = bedrock.converse_stream(
        modelId="amazon.nova-micro-v1:0",
        messages=[{
            "role": "user",
            "content": [{"text": "Name three pioneering medical discoveries. Be brief."}]
        }]
    )
    print("Reply   : ", end="", flush=True)
    for event in stream_resp["stream"]:
        if "contentBlockDelta" in event:
            text = event["contentBlockDelta"]["delta"].get("text", "")
            print(text, end="", flush=True)
    print()

    print("\n--- On-Prem model via Bedrock SDK ---")
    onprem_resp = bedrock.converse(
        modelId="hms-onprem/llama3-8b",
        messages=[{
            "role": "user",
            "content": [{"text": "What is vLLM? One sentence."}]
        }]
    )
    onprem_content = onprem_resp["output"]["message"]["content"][0]["text"]
    print(f"Model   : hms-onprem/llama3-8b (self-hosted via VPC peering)")
    print(f"Reply   : {onprem_content}")


# ── Example 8: Cross-SDK Session — RDS as shared conversation store ───────────

def example_cross_sdk_session():
    """
    Demonstrates that RDS persists conversation history across SDK boundaries.

    A conversation started with the OpenAI SDK (Turn 1) can be continued
    with the Bedrock SDK (Turn 2) using the same session_id.

    The middleware stores chat history in RDS keyed by session_id.
    Neither SDK knows about the other — RDS is the shared memory.

    Real-world scenario:
      - A researcher's Python script (OpenAI SDK) starts a conversation
      - A separate clinical app (Bedrock SDK) continues the same conversation
    """
    print("\n" + "="*60)
    print("Example 8: Cross-SDK Multi-Turn Session")
    print("         (RDS as shared conversation store)")
    print("="*60)

    try:
        import boto3
        from botocore.config import Config
        from botocore import UNSIGNED
    except ImportError:
        print("  boto3 not installed. Skipping this example.")
        return

    bedrock = boto3.client(
        service_name="bedrock-runtime",
        region_name="us-east-1",
        endpoint_url=f"{GATEWAY_URL.replace('/v1', '')}/bedrock",
        config=Config(signature_version=UNSIGNED),
    )

    def add_auth(request, **kwargs):
        request.headers["Authorization"] = f"Bearer {API_KEY}"

    bedrock.meta.events.register("request-created.*", add_auth)

    print("\n  Turn 1 — OpenAI SDK starts the conversation")
    print("  " + "-"*50)

    # Turn 1: OpenAI SDK — start conversation, get session_id from response
    resp1 = client.chat.completions.create(
        model=MODELS["nova-lite"],
        messages=[{
            "role": "user",
            "content": "I am a researcher studying Alzheimer's disease. What is the main protein involved?"
        }],
        extra_body={"enable_history": True},   # tells middleware to store history in RDS
    )

    # session_id is returned in the response — this is the RDS key
    session_id = resp1.model_extra.get("session_id")
    reply1 = resp1.choices[0].message.content

    print(f"  User    : I am a researcher studying Alzheimer's disease. What is the main protein involved?")
    print(f"  Model   : {reply1.strip()}")
    print(f"  Session : {session_id}  ← stored in RDS")

    print("\n  Turn 2 — Bedrock SDK continues the SAME conversation")
    print("  " + "-"*50)

    # Turn 2: Bedrock SDK — continues using session_id from RDS
    resp2 = bedrock.converse(
        modelId="amazon.nova-lite-v1:0",
        messages=[{
            "role": "user",
            "content": [{"text": "What treatments are currently being researched for that?"}]
        }],
        additionalModelRequestFields={"session_id": session_id}  # same session
    )

    reply2 = resp2["output"]["message"]["content"][0]["text"]
    print(f"  User    : What treatments are currently being researched for that?")
    print(f"  Model   : {reply2.strip()}")
    print(f"  Session : {session_id}  ← same RDS record, context preserved")

    print("\n  Turn 3 — OpenAI SDK rejoins the same conversation")
    print("  " + "-"*50)

    # Turn 3: Back to OpenAI SDK — still uses the same session
    resp3 = client.chat.completions.create(
        model=MODELS["nova-lite"],
        messages=[{
            "role": "user",
            "content": "Summarise what we've discussed in one sentence."
        }],
        extra_body={"session_id": session_id},  # same session
    )

    reply3 = resp3.choices[0].message.content
    print(f"  User    : Summarise what we've discussed in one sentence.")
    print(f"  Model   : {reply3.strip()}")
    print(f"\n  Full conversation persisted in RDS across both SDKs")
    print(f"     Session ID: {session_id}")


# ── Main ───────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=" * 60)
    print("  Multi-Provider GenAI Gateway -- End-User Quickstart")
    print("=" * 60)

    try:
        example_list_models()
        example_simple_chat()
        example_model_comparison()
        example_streaming()
        example_multi_turn()
        example_embeddings()
        example_bedrock_sdk()
        example_cross_sdk_session()

        print("\n" + "="*60)
        print("  All examples completed successfully.")
        print("  You're ready to integrate the gateway into your code.")
        print("="*60 + "\n")

    except Exception as e:
        print(f"\nERROR: {e}")
        print("    Check your API key and network connectivity.")
        sys.exit(1)