Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-memory/src/lexigram/ai/memory/working/manager.py: 30%
66 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Working memory manager — assembles token-budgeted context windows."""
3from __future__ import annotations
5from lexigram.ai.memory.config import WorkingMemoryConfig
6from lexigram.ai.memory.working.token_budget import TokenBudgetAllocator
7from lexigram.contracts.ai.memory import (
8 EpisodicMemoryProtocol,
9 MemoryEntry,
10 MemoryQuery,
11 SemanticMemoryProtocol,
12)
13from lexigram.contracts.core.health import HealthCheckResult, HealthStatus
14from lexigram.logging import (
15 get_logger,
16)
18logger = get_logger(__name__)
21class WorkingMemoryManager:
22 """Assembles the optimal context window from all memory tiers.
24 Pulls recent turns from episodic memory and relevant facts from
25 semantic memory, fitting everything within a configurable token budget.
26 """
28 def __init__(
29 self,
30 episodic: EpisodicMemoryProtocol | None = None,
31 semantic: SemanticMemoryProtocol | None = None,
32 config: WorkingMemoryConfig | None = None,
33 ) -> None:
34 """Initialise the working memory manager.
36 Args:
37 episodic: Episodic memory source for past interactions.
38 semantic: Semantic memory source for structured facts.
39 config: Token budget configuration.
40 """
41 self._episodic = episodic
42 self._semantic = semantic
43 self._config = config or WorkingMemoryConfig()
44 self._allocator = TokenBudgetAllocator(self._config)
45 self._current_entries: list[MemoryEntry] = []
47 async def assemble(
48 self,
49 query: str,
50 token_budget: int,
51 *,
52 owner_id: str,
53 session_id: str | None = None,
54 ) -> list[MemoryEntry]:
55 """Assemble context window from all available memory tiers.
57 Args:
58 query: Current user query used to retrieve relevant memories.
59 token_budget: Total token budget for the assembled context.
60 owner_id: Owner scope restricting memory retrieval.
61 session_id: Optional session scope for retrieval filtering.
63 Returns:
64 Ordered list of memory entries ready for LLM context injection.
65 """
66 budget = self._allocator.allocate(token_budget)
67 entries: list[MemoryEntry] = []
69 if self._episodic:
70 filters: dict = {}
71 if session_id:
72 filters["session_id"] = session_id
74 # Recent conversation turns
75 recent = await self._episodic.recall(
76 MemoryQuery(
77 owner_id=owner_id,
78 query=query,
79 top_k=self._config.max_recent_turns,
80 recency_weight=0.8,
81 importance_weight=0.1,
82 relevance_weight=0.1,
83 filters={**filters, "type": "turn"}
84 if filters
85 else {"type": "turn"},
86 )
87 )
88 # Episodic semantic recall
89 episodic_results = await self._episodic.recall(
90 MemoryQuery(
91 owner_id=owner_id,
92 query=query,
93 top_k=10,
94 recency_weight=0.2,
95 importance_weight=0.3,
96 relevance_weight=0.5,
97 filters=filters,
98 )
99 )
101 # Merge: recent turns first, then semantic episodic hits
102 seen_ids: set[str] = set()
103 for result in recent:
104 if result.entry.id not in seen_ids:
105 entries.append(result.entry)
106 seen_ids.add(result.entry.id)
107 for result in episodic_results:
108 if result.entry.id not in seen_ids:
109 entries.append(result.entry)
110 seen_ids.add(result.entry.id)
112 logger.debug(
113 "episodic_assembled",
114 count=len(entries),
115 budget=budget["episodic"] + budget["recent_turns"],
116 )
118 if self._semantic:
119 # Inject top semantic facts as synthetic memory entries
120 from datetime import UTC, datetime
121 from uuid import uuid4
123 facts = await self._semantic.query_facts(subject=query[:100])
124 for fact in facts[: budget["semantic"] // 20]: # rough token estimate
125 fact_entry = MemoryEntry(
126 id=str(uuid4()),
127 owner_id=owner_id,
128 content=f"{fact.get('subject')} {fact.get('predicate')} {fact.get('object_')}",
129 role="system",
130 timestamp=datetime.now(UTC),
131 importance=fact.get("confidence", 0.5),
132 metadata={"source": "semantic", **fact},
133 )
134 entries.append(fact_entry)
136 logger.debug("semantic_injected", count=len(facts))
138 self._current_entries = entries
139 return entries
141 async def add(self, entry: MemoryEntry) -> None:
142 """Add an entry to the current working memory stream.
144 Args:
145 entry: Memory entry to add.
146 """
147 self._current_entries.append(entry)
148 if self._episodic:
149 await self._episodic.record(entry)
151 async def get_context_entries(self) -> list[MemoryEntry]:
152 """Return the entries currently assembled in working context.
154 Returns:
155 Current context entry list.
156 """
157 return list(self._current_entries)
159 async def flush(self) -> None:
160 """Clear the current context assembly."""
161 self._current_entries.clear()
163 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult:
164 """Check health of underlying memory tiers.
166 Args:
167 timeout: Maximum seconds for the health check.
169 Returns:
170 HealthCheckResult aggregating episodic and semantic tier health.
171 """
172 try:
173 ep_healthy = True
174 sem_healthy = True
175 if self._episodic:
176 ep_result = await self._episodic.health_check(timeout)
177 ep_healthy = ep_result.status == HealthStatus.HEALTHY
178 if self._semantic:
179 sem_result = await self._semantic.health_check(timeout)
180 sem_healthy = sem_result.status == HealthStatus.HEALTHY
181 status = (
182 HealthStatus.HEALTHY
183 if ep_healthy and sem_healthy
184 else HealthStatus.DEGRADED
185 )
186 return HealthCheckResult(component="working_memory", status=status)
187 except Exception as exc: # noqa: BLE001
188 return HealthCheckResult(
189 component="working_memory",
190 status=HealthStatus.UNHEALTHY,
191 message=str(exc),
192 )
195__all__ = ["WorkingMemoryManager"]