Caching & Cost Control

import time

class SemanticCache:
    def __init__(self, vector_db, similarity_threshold=0.95, ttl_hours=24):
        self.vector_db = vector_db
        self.threshold = similarity_threshold
        self.ttl = ttl_hours * 3600

    def get(self, query):
        """Find a cached response for a similar query."""
        query_emb = get_embedding(query)
        results = self.vector_db.search(query_emb, top_k=1)

        if results and results[0]["similarity"] >= self.threshold:
            entry = results[0]
            if time.time() - entry["timestamp"] < self.ttl:
                log_cache_hit(query, entry["similarity"])
                return entry["response"]

        log_cache_miss(query)
        return None

    def set(self, query, response):
        """Cache a query-response pair."""
        embedding = get_embedding(query)
        self.vector_db.insert(
            text=query,
            embedding=embedding,
            metadata={
                "response": response,
                "timestamp": time.time(),
            }
        )

def batch_classify(alerts, batch_size=10):
    """Classify multiple alerts in a single API call."""
    results = []
    for i in range(0, len(alerts), batch_size):
        batch = alerts[i:i + batch_size]

        # Combine into a single prompt
        alert_list = "\n".join([
            f"{j+1}. {a['name']}: {a['description']}"
            for j, a in enumerate(batch)
        ])

        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            temperature=0.0,
            messages=[
                {
                    "role": "system",
                    "content": "Classify each alert as P1-P4. Return JSON array."
                },
                {"role": "user", "content": alert_list}
            ]
        )

        batch_results = json.loads(response["choices"][0]["message"]["content"])
        results.extend(batch_results)

    return results

def compress_context(documents, max_tokens=2000):
    """Compress retrieved documents to fit within token budget."""
    # Strategy 1: Extract key sentences
    compressed = []
    total_tokens = 0

    for doc in documents:
        sentences = doc.split(". ")
        for sentence in sentences:
            sentence_tokens = len(sentence) // 4
            if total_tokens + sentence_tokens > max_tokens:
                break
            compressed.append(sentence)
            total_tokens += sentence_tokens

    return ". ".join(compressed)

# Strategy 2: Use a smaller model to summarize
def summarize_context(documents, max_tokens=2000):
    """Use GPT-3.5 to summarize context before sending to GPT-4."""
    combined = "\n---\n".join(documents)
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "user", "content": f"Summarize these documents in under {max_tokens} tokens:\n{combined}"}
        ]
    )
    return response["choices"][0]["message"]["content"]

class TokenBudget:
    def __init__(self):
        self.budgets = {
            "alert_triage": {"daily": 100000, "per_request": 2000},
            "runbook_lookup": {"daily": 50000, "per_request": 4000},
            "incident_summary": {"daily": 200000, "per_request": 8000},
        }
        self.usage = defaultdict(int)

    def check_budget(self, feature, estimated_tokens):
        """Check if we are within budget."""
        budget = self.budgets.get(feature, {"daily": 100000})
        if self.usage[feature] + estimated_tokens > budget["daily"]:
            return {"allowed": False, "reason": f"Daily budget exceeded for {feature}"}
        return {"allowed": True}

    def record_usage(self, feature, tokens_used):
        """Record token usage."""
        self.usage[feature] += tokens_used

    def get_remaining(self, feature):
        """Get remaining budget for today."""
        budget = self.budgets.get(feature, {"daily": 100000})
        return budget["daily"] - self.usage[feature]