from LLM.llm_models import context_compression_model
from .get_context_compression_prompt import get_context_compression_prompt
from smolagents import PlanningStep
from typing import List

client = context_compression_model


def compress_context(context: List[PlanningStep]):
    completion = client.generate(
        messages=[
            {
                "role": "system",
                "content": get_context_compression_prompt()
            },
            {
                "role": "user",
                "content": "\n".join([str(step) for step in context])
            }
        ],
    )
    return completion.content