Merge pull request #1 from huggingface/compacting-context

agent/codex_agent_demo.py

@@ -201,7 +201,7 @@ class Session:
     """
 
     def __init__(self, event_queue: asyncio.Queue):
-        self.context_manager = ContextManager()
+        self.context_manager = ContextManager(tool_specs=[])
         self.event_queue = event_queue
         self.is_running = True
         self.current_task: Optional[asyncio.Task] = None

agent/config.py

@@ -20,7 +20,6 @@ class Config(BaseModel):
 
     model_name: str
     tools: list[Tool] = []
-    system_prompt_path: str = ""
     mcpServers: dict[str, MCPServerConfig] = {}
 
 

agent/context_manager/manager.py

@@ -2,43 +2,101 @@
 Context management for conversation history
 """
 
-from litellm import Message
+from pathlib import Path
+from typing import Any
+
+import yaml
+from jinja2 import Template
+from litellm import Message, acompletion
 
 
 class ContextManager:
     """Manages conversation context and message history for the agent"""
 
-    def __init__(self):
-        self.system_prompt = self._load_system_prompt()
+    def __init__(
+        self,
+        max_context: int = 180_000,
+        compact_size: float = 0.1,
+        untouched_messages: int = 5,
+        tool_specs: list[dict[str, Any]] | None = None,
+    ):
+        self.system_prompt = self._load_system_prompt(tool_specs or [])
+        self.max_context = max_context
+        self.compact_size = int(max_context * compact_size)
+        self.context_length = len(self.system_prompt) // 4
+        self.untouched_messages = untouched_messages
         self.items: list[Message] = [Message(role="system", content=self.system_prompt)]
 
-    def _load_system_prompt(self):
-        """Load the system prompt"""
+    def _load_system_prompt(self, tool_specs: list[dict[str, Any]]):
+        """Load and render the system prompt from YAML file with Jinja2"""
+        prompt_file = Path(__file__).parent.parent / "prompts" / "system_prompt.yaml"
+
+        with open(prompt_file, "r") as f:
+            prompt_data = yaml.safe_load(f)
+            template_str = prompt_data.get("system_prompt", "")
 
-        # TODO: get system prompt from jinja template
-        return "You are a helpful assistant."
+        template = Template(template_str)
+        return template.render(
+            tools=tool_specs,
+            num_tools=len(tool_specs),
+        )
 
-    def add_message(self, message: Message) -> None:
+    def add_message(self, message: Message, token_count: int = None) -> None:
         """Add a message to the history"""
+        if token_count:
+            self.context_length = token_count
+            print(f"DEBUG : token_count = {self.context_length}")
         self.items.append(message)
 
     def get_messages(self) -> list[Message]:
         """Get all messages for sending to LLM"""
         return self.items
 
-    def compact(self, target_size: int) -> None:
+    async def compact(self, model_name: str) -> None:
         """Remove old messages to keep history under target size"""
-        # Keep system prompt (first message) and remove oldest user/assistant messages
-        if len(self.items) <= target_size:
+        if (self.context_length <= self.max_context) or not self.items:
             return
 
-        # Always keep system prompt
         system_msg = (
             self.items[0] if self.items and self.items[0].role == "system" else None
         )
-        messages_to_keep = self.items[-(target_size - 1) :]
 
+        # Don't summarize a certain number of just-preceding messages
+        # Walk back to find a user message to make sure we keep an assistant -> user ->
+        # assistant general conversation structure
+        idx = len(self.items) - self.untouched_messages
+        while idx > 1 and self.items[idx].role != "user":
+            idx -= 1
+
+        recent_messages = self.items[idx:]
+        messages_to_summarize = self.items[1:idx]
+
+        # improbable, messages would have to very long
+        if not messages_to_summarize:
+            return
+
+        messages_to_summarize.append(
+            Message(
+                role="user",
+                content="Please provide a concise summary of the conversation above, focusing on key decisions, code changes, problems solved, and important context needed for future turns.",
+            )
+        )
+
+        response = await acompletion(
+            model=model_name,
+            messages=messages_to_summarize,
+            max_completion_tokens=self.compact_size,
+        )
+        summarized_message = Message(
+            role="assistant", content=response.choices[0].message.content
+        )
+
+        # Reconstruct: system + summary + recent messages (includes tools)
         if system_msg:
-            self.items = [system_msg] + messages_to_keep
+            self.items = [system_msg, summarized_message] + recent_messages
         else:
-            self.items = messages_to_keep
+            self.items = [summarized_message] + recent_messages
+
+        self.context_length = (
+            len(self.system_prompt) // 4 + response.usage.completion_tokens
+        )

agent/core/agent_loop.py

@@ -1,4 +1,4 @@
-"""
+"""loop
 Main agent implementation with integrated tool system and MCP support
 """
 
@@ -58,17 +58,17 @@ class Handlers:
                     tool_choice="auto",
                 )
 
+                # Extract text response, token usage, and tool calls
                 message = response.choices[0].message
-
-                # Extract content and tool calls
                 content = message.content
+                token_count = response.usage.total_tokens
                 tool_calls: list[ToolCall] = message.get("tool_calls", [])
 
                 # If no tool calls, add assistant message and we're done
                 if not tool_calls:
                     if content:
                         assistant_msg = Message(role="assistant", content=content)
-                        session.context_manager.add_message(assistant_msg)
+                        session.context_manager.add_message(assistant_msg, token_count)
                         await session.send_event(
                             Event(
                                 event_type="assistant_message",
@@ -81,9 +81,11 @@ class Handlers:
                 # Add assistant message with tool calls to history
                 # LiteLLM will format this correctly for the provider
                 assistant_msg = Message(
-                    role="assistant", content=content, tool_calls=tool_calls
+                    role="assistant",
+                    content=content,
+                    tool_calls=tool_calls,
                 )
-                session.context_manager.add_message(assistant_msg)
+                session.context_manager.add_message(assistant_msg, token_count)
 
                 if content:
                     await session.send_event(
@@ -139,6 +141,18 @@ class Handlers:
                 )
                 break
 
+        old_length = session.context_manager.context_length
+        await session.context_manager.compact(model_name=session.config.model_name)
+        new_length = session.context_manager.context_length
+
+        if new_length != old_length:
+            await session.send_event(
+                Event(
+                    event_type="compacted",
+                    data={"old_tokens": old_length, "new_tokens": new_length},
+                )
+            )
+
         await session.send_event(
             Event(
                 event_type="turn_complete",
@@ -156,14 +170,14 @@ class Handlers:
     @staticmethod
     async def compact(session: Session) -> None:
         """Handle compact (like compact in codex.rs:1317)"""
-        old_size = len(session.context_manager.items)
-        session.context_manager.compact(target_size=10)
-        new_size = len(session.context_manager.items)
+        old_length = session.context_manager.context_length
+        await session.context_manager.compact(model_name=session.config.model_name)
+        new_length = session.context_manager.context_length
 
         await session.send_event(
             Event(
                 event_type="compacted",
-                data={"removed": old_size - new_size, "remaining": new_size},
+                data={"removed": old_length, "remaining": new_length},
             )
         )
 
@@ -231,9 +245,8 @@ async def submission_loop(
     This is the core of the agent (like submission_loop in codex.rs:1259-1340)
     """
 
-    # Create session and assign tool router
-    session = Session(event_queue, config=config)
-    session.tool_router = tool_router
+    # Create session with tool router
+    session = Session(event_queue, config=config, tool_router=tool_router)
     print("🤖 Agent loop started")
 
     # Main processing loop

agent/core/session.py

@@ -4,6 +4,8 @@ from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Optional
 
+from litellm import get_max_tokens
+
 from agent.config import Config
 from agent.context_manager.manager import ContextManager
 
@@ -33,18 +35,24 @@ class Session:
         self,
         event_queue: asyncio.Queue,
         config: Config | None = None,
+        tool_router=None,
     ):
-        self.context_manager = ContextManager()
+        self.tool_router = tool_router
+        tool_specs = tool_router.get_tool_specs_for_llm() if tool_router else []
+        self.context_manager = ContextManager(
+            max_context=get_max_tokens(config.model_name),
+            compact_size=0.1,
+            untouched_messages=5,
+            tool_specs=tool_specs,
+        )
         self.event_queue = event_queue
         self.session_id = str(uuid.uuid4())
         self.config = config or Config(
             model_name="anthropic/claude-sonnet-4-5-20250929",
             tools=[],
-            system_prompt_path="",
         )
         self.is_running = True
         self.current_task: asyncio.Task | None = None
-        self.tool_router = None  # Set by submission_loop
 
     async def send_event(self, event: Event) -> None:
         """Send event back to client"""

agent/main.py

@@ -88,6 +88,10 @@ async def event_listener(
                 break
             elif event.event_type == "processing":
                 print("⏳ Processing...", flush=True)
+            elif event.event_type == "compacted":
+                old_tokens = event.data.get("old_tokens", 0) if event.data else 0
+                new_tokens = event.data.get("new_tokens", 0) if event.data else 0
+                print(f"📦 Compacted context: {old_tokens} → {new_tokens} tokens")
             # Silently ignore other events
 
         except asyncio.CancelledError:

agent/prompts/system_prompt.yaml

@@ -0,0 +1,112 @@
+system_prompt: |
+  You are HF Agent, a powerful AI assistant for Machine Learning Engineering, particularly training Large Language Models. You have access to {{ num_tools }} tools for interacting with Hugging Face Hub and performing ML tasks.
+
+  # Available Tools
+
+  You have access to the following categories of tools:
+
+  - Hugging Face Hub: Search and interact with models, datasets, papers, and documentation
+  - Spaces: Use and discover ML applications
+  - Jobs: Manage compute jobs for training and inference
+  - Image Generation: Generate and transform images
+
+  # Agency
+
+  You take initiative when the user asks you to do something, maintaining an appropriate balance between:
+
+  1. Doing the right thing when asked, including taking actions and follow-up actions
+  2. Not surprising the user with actions you take without asking
+  3. Not adding unnecessary explanations after completing tasks
+
+  # Task Approach
+
+  For ML engineering tasks:
+  1. Use all available tools to complete the task
+  2. Search for relevant models, datasets, and documentation on Hugging Face Hub
+  3. Leverage existing resources before creating new ones
+  4. Invoke multiple independent tools simultaneously for efficiency
+
+  # Examples
+
+  <example>
+  <user>Find the best text generation models</user>
+  <response>[uses mcp__hf-mcp-server__model_search with task="text-generation" and sort="trendingScore"]
+
+  Top trending text generation models:
+  - meta-llama/Llama-3.1-405B-Instruct
+  - mistralai/Mistral-Large-2
+  </response>
+  </example>
+
+  <example>
+  <user>Search for papers about reinforcement learning from human feedback</user>
+  <response>[uses mcp__hf-mcp-server__paper_search with query="reinforcement learning from human feedback"]
+
+  Found 5 relevant papers on RLHF including "Training language models to follow instructions with human feedback" (Ouyang et al.)
+  </response>
+  </example>
+
+  <example>
+  <user>Find datasets for sentiment analysis</user>
+  <response>[uses mcp__hf-mcp-server__dataset_search with query="sentiment analysis" and tags for task_categories]
+
+  Top sentiment analysis datasets:
+  - stanfordnlp/imdb (25k reviews)
+  - tweet_eval (sentiment task)
+  </response>
+  </example>
+
+  <example>
+  <user>How do I use the transformers library for text generation?</user>
+  <response>[uses mcp__hf-mcp-server__hf_doc_search with query="text generation transformers"]
+
+  [provides concise answer based on documentation]
+  </response>
+  </example>
+
+  <example>
+  <user>Generate an image of a sunset over mountains</user>
+  <response>[uses mcp__hf-mcp-server__gr1_flux1_schnell_infer with prompt="sunset over mountains"]
+
+  [returns generated image]
+  </response>
+  </example>
+
+  <example>
+  <user>Get details about the bert-base-uncased model</user>
+  <response>[uses mcp__hf-mcp-server__hub_repo_details with repo_ids=["google-bert/bert-base-uncased"]]
+
+  BERT base uncased: 110M parameters, trained on English Wikipedia and BookCorpus, commonly used for text classification and NER.
+  </response>
+  </example>
+
+  # Conventions
+
+  - Always search Hugging Face Hub for existing resources before suggesting custom implementations
+  - When referencing models, datasets, or papers, include direct links from search results
+  - Never assume a library is available - check documentation first
+  - Follow ML best practices: proper train/val/test splits, reproducibility, evaluation metrics
+  - For training tasks, consider compute requirements and suggest appropriate hardware
+  - Never expose or log API keys, tokens, or secrets
+
+  # Communication Style
+
+  - Be concise and direct
+  - Skip flattery and unnecessary preamble
+  - Respond in 1-3 sentences when possible
+  - No emojis, minimal exclamation points
+  - Don't apologize for limitations - offer alternatives or keep responses short
+  - Don't thank the user for results
+  - Explain what you're doing for non-trivial operations
+
+  Answer the user's question directly without elaboration unless they ask for detail. One word answers are best when appropriate.
+
+  <example>
+  <user>What's the state-of-the-art model for image classification?</user>
+  <response>EVA-CLIP-18B or ConvNeXt-XXLarge depending on your constraints</response>
+  </example>
+
+  <example>
+  <user>How many parameters does GPT-3 have?</user>
+  <response>175 billion</response>
+  </example>