heaven_base.memory.heaven_history

  1import sys
  2import uuid
  3import json
  4import os
  5import re
  6from datetime import datetime
  7from typing import List, Optional, Dict, Any, Union, Literal
  8from langchain_core.messages import (
  9    BaseMessage, 
 10    SystemMessage, 
 11    HumanMessage, 
 12    AIMessage, 
 13    ToolMessage
 14)
 15from pydantic import BaseModel, Field
 16from .base_piece import BasePiece
 17from .history import History, AgentStatus
 18from .heaven_event import HeavenEvent, is_heaven_event_dict, convert_langchain_to_heaven, convert_heaven_to_langchain
 19from ..utils.name_utils import normalize_agent_name
 20
 21# Define SystemPrompt class for HEAVEN events
 22class SystemPrompt(BaseModel):
 23    """Represents a system prompt in the HEAVEN system.
 24    
 25    System prompts provide instructions, context, and guidelines to agents.
 26    They influence the agent's behavior and responses.
 27    """
 28    content: str = ""
 29    metadata: Dict[str, Any] = Field(default_factory=dict)
 30    
 31    def to_langchain(self) -> SystemMessage:
 32        """Convert to LangChain SystemMessage."""
 33        return SystemMessage(content=self.content, additional_kwargs=self.metadata)
 34    
 35    @classmethod
 36    def from_langchain(cls, message: SystemMessage) -> 'SystemPrompt':
 37        """Create from LangChain SystemMessage."""
 38        return cls(
 39            content=message.content,
 40            metadata=message.additional_kwargs
 41        )
 42    
 43    def to_heaven_event(self) -> Dict[str, Any]:
 44        """Convert to HEAVEN event format."""
 45        return HeavenEvent(
 46            event_type="SYSTEM_MESSAGE",
 47            data={
 48                "content": self.content,
 49                "metadata": self.metadata
 50            }
 51        ).to_dict()
 52    
 53    @classmethod
 54    def from_heaven_event(cls, event: Dict[str, Any]) -> Optional['SystemPrompt']:
 55        """Create from HEAVEN event."""
 56        if not is_heaven_event_dict(event) or event.get("event_type") != "SYSTEM_MESSAGE":
 57            return None
 58        
 59        data = event.get("data", {})
 60        return cls(
 61            content=data.get("content", ""),
 62            metadata=data.get("metadata", {})
 63        )
 64
 65class HeavenHistory(History):
 66    """Enhanced History class that supports both LangChain messages and HEAVEN events"""
 67    # We keep the original messages field for backward compatibility
 68    messages: List[BaseMessage]
 69    # Add a new field for HEAVEN events
 70    heaven_events: List[Dict[str, Any]] = Field(default_factory=list)
 71    # Add a field for system prompts
 72    system_prompts: List[SystemPrompt] = Field(default_factory=list)
 73    # Flag to track the primary storage format
 74    primary_format: Literal["langchain", "heaven"] = "langchain"
 75    
 76    def __init__(self, **data):
 77        super().__init__(**data)
 78        # Initialize heaven_events if not provided
 79        if not hasattr(self, "heaven_events") or not self.heaven_events:
 80            # Convert existing messages to HEAVEN events
 81            self.heaven_events = convert_langchain_to_heaven(self.messages)
 82        
 83        # Initialize system_prompts if not provided
 84        if not hasattr(self, "system_prompts") or not self.system_prompts:
 85            # Extract system messages and convert to SystemPrompt objects
 86            system_messages = [msg for msg in self.messages if isinstance(msg, SystemMessage)]
 87            self.system_prompts = [SystemPrompt.from_langchain(msg) for msg in system_messages]
 88    
 89    def add_message(self, message: BaseMessage) -> None:
 90        """Add a LangChain message to the history."""
 91        self.messages.append(message)
 92        
 93        # Also add as HEAVEN event
 94        heaven_event = HeavenEvent.from_langchain_message(message).to_dict()
 95        self.heaven_events.append(heaven_event)
 96        
 97        # If it's a system message, also add to system_prompts
 98        if isinstance(message, SystemMessage):
 99            system_prompt = SystemPrompt.from_langchain(message)
100            self.system_prompts.append(system_prompt)
101    
102    def add_heaven_event(self, event: Union[Dict[str, Any], HeavenEvent]) -> None:
103        """Add a HEAVEN event to the history."""
104        # Convert to dict if it's a HeavenEvent object
105        if isinstance(event, HeavenEvent):
106            event_dict = event.to_dict()
107        else:
108            event_dict = event
109        
110        # Add to heaven_events
111        self.heaven_events.append(event_dict)
112        
113        # Check if it's a system message event
114        if is_heaven_event_dict(event_dict) and event_dict.get("event_type") == "SYSTEM_MESSAGE":
115            # Create a SystemPrompt from the event
116            system_prompt = SystemPrompt.from_heaven_event(event_dict)
117            if system_prompt:
118                self.system_prompts.append(system_prompt)
119        
120        # Convert to LangChain message if possible and add to messages
121        if is_heaven_event_dict(event_dict):
122            heaven_event = HeavenEvent.from_dict(event_dict)
123            message = heaven_event.to_langchain_message()
124            if message:
125                self.messages.append(message)
126    
127    def get_heaven_events(self) -> List[Dict[str, Any]]:
128        """Get all HEAVEN events in the history."""
129        return self.heaven_events
130    
131    def get_langchain_messages(self) -> List[BaseMessage]:
132        """Get all LangChain messages in the history."""
133        return self.messages
134    
135    def get_system_prompts(self) -> List[SystemPrompt]:
136        """Get all system prompts in the history."""
137        return self.system_prompts
138    
139    def get_combined_system_prompt(self) -> str:
140        """Get all system prompts combined into a single string."""
141        return "\n\n".join([sp.content for sp in self.system_prompts if sp.content])
142    
143    def add_system_prompt(self, content: str, metadata: Optional[Dict[str, Any]] = None) -> None:
144        """Add a system prompt to the history."""
145        # Create SystemPrompt object
146        system_prompt = SystemPrompt(content=content, metadata=metadata or {})
147        self.system_prompts.append(system_prompt)
148        
149        # Add as LangChain message
150        self.messages.append(system_prompt.to_langchain())
151        
152        # Add as HEAVEN event
153        self.heaven_events.append(system_prompt.to_heaven_event())
154    
155    def to_json(self) -> dict:
156        """Convert to JSON format with added HEAVEN events and system prompts."""
157        base_json = super().to_json()
158        base_json["heaven_events"] = self.heaven_events
159        base_json["primary_format"] = self.primary_format
160        base_json["system_prompts"] = [sp.dict() for sp in self.system_prompts]
161        return base_json
162    
163    @classmethod
164    def from_json(cls, data: dict, project: str = None) -> "HeavenHistory":
165        """Create HeavenHistory from JSON."""
166        # First create a standard History object
167        history = super().from_json(data, project)
168        
169        # Extract heaven_events if present
170        heaven_events = data.get("heaven_events", [])
171        primary_format = data.get("primary_format", "langchain")
172        
173        # Extract system_prompts if present
174        system_prompts_data = data.get("system_prompts", [])
175        system_prompts = [SystemPrompt(**sp_data) for sp_data in system_prompts_data]
176        
177        # Create and return a HeavenHistory
178        return cls(
179            messages=history.messages,
180            created_datetime=history.created_datetime,
181            metadata=history.metadata,
182            project=history.project,
183            agent_status=history.agent_status,
184            history_id=history.history_id,
185            json_md_path=history.json_md_path,
186            heaven_events=heaven_events,
187            primary_format=primary_format,
188            system_prompts=system_prompts
189        )
190    
191    @classmethod
192    def from_history(cls, history: History) -> "HeavenHistory":
193        """Convert a standard History to HeavenHistory."""
194        # Convert messages to HEAVEN events
195        heaven_events = convert_langchain_to_heaven(history.messages)
196        
197        # Extract system messages
198        system_messages = [msg for msg in history.messages if isinstance(msg, SystemMessage)]
199        system_prompts = [SystemPrompt.from_langchain(msg) for msg in system_messages]
200        
201        # Create and return a HeavenHistory
202        return cls(
203            messages=history.messages,
204            created_datetime=history.created_datetime,
205            metadata=history.metadata,
206            project=history.project,
207            agent_status=history.agent_status,
208            history_id=history.history_id,
209            json_md_path=history.json_md_path,
210            heaven_events=heaven_events,
211            system_prompts=system_prompts,
212            primary_format="langchain"
213        )
214    
215    @classmethod
216    def from_heaven_events(cls, events: List[Dict[str, Any]], metadata: Dict[str, Any] = None) -> "HeavenHistory":
217        """Create a HeavenHistory from a list of HEAVEN events."""
218        # Convert HEAVEN events to LangChain messages
219        messages = convert_heaven_to_langchain(events)
220        
221        # Extract system message events
222        system_events = [event for event in events 
223                        if is_heaven_event_dict(event) and event.get("event_type") == "SYSTEM_MESSAGE"]
224        system_prompts = [SystemPrompt.from_heaven_event(event) for event in system_events 
225                         if SystemPrompt.from_heaven_event(event) is not None]
226        
227        # Create and return a HeavenHistory
228        return cls(
229            messages=messages,
230            created_datetime=datetime.now(),
231            metadata=metadata or {},
232            heaven_events=events,
233            system_prompts=system_prompts,
234            primary_format="heaven"
235        )
236    
237    def _compute_heaven_iterations(self) -> Dict[str, List[Dict[str, Any]]]:
238        """Split heaven_events into iterations."""
239        iterations: Dict[str, List[Dict[str, Any]]] = {}
240        current: List[Dict[str, Any]] = []
241        idx = 0
242        
243        for event in self.heaven_events:
244            if is_heaven_event_dict(event) and event.get("event_type") == "USER_MESSAGE":
245                if current:
246                    iterations[f"iteration_{idx}"] = current
247                    idx += 1
248                    current = []
249            else:
250                current.append(event)
251        
252        if current:
253            iterations[f"iteration_{idx}"] = current
254        
255        return iterations
256    
257    @property
258    def heaven_iterations(self) -> Dict[str, List[Dict[str, Any]]]:
259        """Returns a dict of all iterations in HEAVEN event format."""
260        return self._compute_heaven_iterations()
261
262
263def get_heaven_iteration_view(
264    history: HeavenHistory,
265    start: int,
266    end: int
267) -> Dict[str, Any]:
268    """Get a view of iterations in HEAVEN event format."""
269    if start > end:
270        raise ValueError(f"start ({start}) must be <= end ({end})")
271    
272    all_iters = history.heaven_iterations
273    total = len(all_iters)
274    
275    # Build the view
276    view: Dict[str, List[Dict[str, Any]]] = {}
277    for i in range(start, end + 1):
278        key = f"iteration_{i}"
279        if key not in all_iters:
280            raise KeyError(f"Missing iteration: {key}")
281        view[key] = all_iters[key]
282    
283    return {
284        "history_id": history.history_id,
285        "total_iterations": total,
286        "view_range": {"start": start, "end": end},
287        "view": view
288    }