File size: 3,952 Bytes
942216e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
from datetime import datetime
from typing import Dict, List, Optional, Any, Union
from uuid import uuid4
import json


class Context:
    def __init__(self, context_id: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, max_history_length: int = 100):
        self.context_id = context_id or str(uuid4())
        self.metadata = metadata or {}
        self.max_history_length = max_history_length
        self.history: List[Dict[str, Any]] = []
        self.state: Dict[str, Any] = {}
        self.created_at = datetime.utcnow()
        self.updated_at = self.created_at
        
    def add_interaction(self, input_data: Dict[str, Any], output_data: Dict[str, Any]) -> None:
        interaction = {
            "timestamp": datetime.utcnow().isoformat(),
            "input": input_data,
            "output": output_data
        }
        
        self.history.append(interaction)
        
        if len(self.history) > self.max_history_length:
            self.history = self.history[-self.max_history_length:]
            
        self.updated_at = datetime.utcnow()
        
    def update_state(self, new_state: Dict[str, Any], merge: bool = True) -> None:
        if merge:
            self.state.update(new_state)
        else:
            self.state = new_state
            
        self.updated_at = datetime.utcnow()
        
    def get_history(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
        if limit is None or limit >= len(self.history):
            return self.history
        return self.history[-limit:]
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "context_id": self.context_id,
            "metadata": self.metadata,
            "state": self.state,
            "history": self.history,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
        }
    
    def to_json(self) -> str:
        return json.dumps(self.to_dict())
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Context':
        context = cls(context_id=data.get("context_id"))
        context.metadata = data.get("metadata", {})
        context.state = data.get("state", {})
        context.history = data.get("history", [])
        
        if "created_at" in data:
            context.created_at = datetime.fromisoformat(data["created_at"])
        if "updated_at" in data:
            context.updated_at = datetime.fromisoformat(data["updated_at"])
            
        return context
    
    @classmethod
    def from_json(cls, json_str: str) -> 'Context':
        data = json.loads(json_str)
        return cls.from_dict(data)


class ContextManager:
    def __init__(self):
        self.contexts: Dict[str, Context] = {}
        
    def create_context(self, metadata: Optional[Dict[str, Any]] = None, context_id: Optional[str] = None) -> Context:
        context = Context(context_id=context_id, metadata=metadata)
        self.contexts[context.context_id] = context
        return context
    
    def get_context(self, context_id: str) -> Optional[Context]:
        return self.contexts.get(context_id)
    
    def update_context(self, context_id: str, updates: Dict[str, Any]) -> Optional[Context]:
        context = self.get_context(context_id)
        if not context:
            return None
        
        if "metadata" in updates:
            context.metadata.update(updates["metadata"])
            
        if "state" in updates:
            context.update_state(updates["state"])
            
        return context
    
    def delete_context(self, context_id: str) -> bool:
        if context_id in self.contexts:
            del self.contexts[context_id]
            return True
        return False
    
    def list_contexts(self) -> List[str]:
        return list(self.contexts.keys())