llm.py

# Day 7: LLM abstraction — supports OpenAI AND Bedrock.
# main.py calls THIS instead of OpenAI directly.

from openai import OpenAI
import boto3
import json
import os

# ─── OpenAI ───
def call_openai(messages: list) -> str:
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=messages,
    )
    return response.choices[0].message.content

# ─── AWS Bedrock ───
def call_bedrock(messages: list) -> str:
    bedrock = boto3.client("bedrock-runtime", region_name="us-east-1")

    # Bedrock needs system prompt separate from messages
    system = ""
    chat_messages = []
    for m in messages:
        if m["role"] == "system":
            system = m["content"]
        else:
            chat_messages.append({"role": m["role"], "content": m["content"]})

    response = bedrock.invoke_model(
        modelId="anthropic.claude-3-haiku-20240307-v1:0",
        body=json.dumps({
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 1000,
            "system": system,
            "messages": chat_messages,
        })
    )
    result = json.loads(response["body"].read())
    return result["content"][0]["text"]

# ─── Universal function — picks the right provider ───
def generate_response(messages: list, provider: str = "openai") -> str:
    """Call the LLM. Provider can be 'openai' or 'bedrock'."""
    if provider == "bedrock":
        return call_bedrock(messages)
    else:
        return call_openai(messages)