Import ducklm runtime

2026-05-10 23:37:56 +08:00
parent fd1a045488
commit dc8267880a
90 changed files with 9171 additions and 0 deletions
@@ -0,0 +1,24 @@
+MEMORY_AVAILABLE = False
+VECTOR_AVAILABLE = False
+
+try:
+    from app.memory.store import MemoryStore
+    from app.memory.vector_index import VectorIndex
+    from app.memory.interface import MemoryInterface
+    from app.memory.write_policy import MemoryWritePolicy
+    MEMORY_AVAILABLE = True
+    VECTOR_AVAILABLE = True
+except ImportError:
+    MemoryStore = None
+    VectorIndex = None
+    MemoryInterface = None
+    MemoryWritePolicy = None
+
+__all__ = [
+    "MemoryStore",
+    "VectorIndex",
+    "MemoryInterface",
+    "MemoryWritePolicy",
+    "MEMORY_AVAILABLE",
+    "VECTOR_AVAILABLE",
+]
@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Literal
+
+import numpy as np
+
+from app.core.contracts import MemoryEntry
+from app.memory.store import MemoryStore
+from app.memory.vector_index import VectorIndex
+from app.models.embeddings import EmbeddingsAdapter
+
+
+class MemoryInterface:
+    def __init__(
+        self,
+        store: MemoryStore,
+        vector_index: VectorIndex,
+        embeddings: EmbeddingsAdapter,
+    ) -> None:
+        self._store = store
+        self._vector_index = vector_index
+        self._embeddings = embeddings
+
+    def insert(
+        self,
+        text: str,
+        kind: Literal["tool_result", "plan", "critique", "fact", "summary", "user_preference"],
+        source: Literal["tool", "critic", "user", "system"],
+        task_id: str | None = None,
+        session_id: str | None = None,
+        weight: float = 0.5,
+        metadata: dict[str, Any] | None = None,
+    ) -> MemoryEntry:
+        entry = MemoryEntry(
+            text=text,
+            kind=kind,
+            source=source,
+            weight=weight,
+            task_id=task_id,
+            session_id=session_id,
+            metadata=metadata or {},
+            embedding_model=self._embeddings.__class__.__name__,
+            embedding_dim=self._embeddings.embedding_dim,
+        )
+
+        embedding = self._embeddings.encode(text)
+        embedding_bytes = embedding.astype("float32").tobytes()
+
+        self._store.insert(entry, embedding_bytes)
+        self._vector_index.insert(entry.id, embedding)
+        self._vector_index.save()
+
+        self.cleanup()
+
+        return entry
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 5,
+        kind: str | None = None,
+        session_id: str | None = None,
+    ) -> list[tuple[MemoryEntry, float]]:
+        query_embedding = self._embeddings.encode(query)
+        memory_ids, scores = self._vector_index.search(query_embedding, k=top_k)
+
+        results: list[tuple[MemoryEntry, float]] = []
+        for memory_id, score in zip(memory_ids, scores):
+            entry = self._store.get(memory_id)
+            if entry:
+                if kind and entry.kind != kind:
+                    continue
+                if session_id and entry.session_id != session_id:
+                    continue
+                results.append((entry, score))
+
+        return results[:top_k]
+
+    def get(self, memory_id: str) -> MemoryEntry | None:
+        return self._store.get(memory_id)
+
+    def delete(self, memory_id: str) -> bool:
+        entry = self._store.get(memory_id)
+        if entry:
+            self._vector_index.delete(memory_id)
+            return self._store.delete(memory_id)
+        return False
+
+    def get_by_task(self, task_id: str) -> list[MemoryEntry]:
+        return self._store.get_by_task(task_id)
+
+    def get_by_session(self, session_id: str, limit: int = 100) -> list[MemoryEntry]:
+        return self._store.get_by_session(session_id, limit)
+
+    def get_recent(self, limit: int = 10) -> list[MemoryEntry]:
+        return self._store.get_all(limit)
+
+    def count(self) -> int:
+        return self._store.count()
+
+    def reindex(self) -> None:
+        entries = self._store.get_all(limit=10000)
+        self._vector_index.save()
+        for entry in entries:
+            text = entry.text
+            embedding = self._embeddings.encode(text)
+            self._vector_index.insert(entry.id, embedding)
+        self._vector_index.save()
+
+    def close(self) -> None:
+        self._store.close()
+
+    def cleanup(self, max_items: int = 750, decay_factor: float = 0.95) -> int:
+        """Remove low-weight entries when exceeding max_items limit.
+
+        Applies weight decay based on freshness before cleanup.
+        Returns number of removed entries.
+        """
+        current_count = self._store.count()
+        if current_count <= max_items:
+            return 0
+
+        removed = 0
+        entries_to_remove = current_count - max_items
+
+        all_entries = self._store.get_all(limit=current_count)
+
+        def effective_weight(entry: MemoryEntry) -> float:
+            entry_weight = entry.weight
+            if entry.created_at:
+                age_days = (datetime.now(timezone.utc) - entry.created_at).total_seconds() / 86400
+                freshness_factor = max(0.1, decay_factor ** age_days)
+                return entry_weight * freshness_factor
+            return entry_weight
+
+        sorted_entries = sorted(all_entries, key=effective_weight)
+
+        for entry in sorted_entries[:entries_to_remove]:
+            self._store.delete(entry.id)
+            removed += 1
+
+        return removed
@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import json
+import sqlite3
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Sequence
+from uuid import uuid4
+
+from app.core.contracts import MemoryEntry
+
+
+def utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class MemoryStore:
+    def __init__(self, db_path: str | Path) -> None:
+        self._db_path = Path(db_path)
+        self._db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._conn = sqlite3.connect(str(self._db_path), check_same_thread=False)
+        self._conn.row_factory = sqlite3.Row
+        self._init_tables()
+
+    def _init_tables(self) -> None:
+        self._conn.executescript("""
+            CREATE TABLE IF NOT EXISTS memory_items (
+                id TEXT PRIMARY KEY,
+                text TEXT NOT NULL,
+                kind TEXT NOT NULL,
+                source TEXT NOT NULL,
+                weight REAL NOT NULL DEFAULT 0.5,
+                task_id TEXT,
+                session_id TEXT,
+                metadata_json TEXT,
+                created_at TEXT NOT NULL,
+                updated_at TEXT NOT NULL
+            );
+
+            CREATE TABLE IF NOT EXISTS memory_embeddings (
+                memory_id TEXT PRIMARY KEY,
+                embedding BLOB NOT NULL,
+                embedding_model TEXT NOT NULL,
+                embedding_dim INTEGER NOT NULL,
+                created_at TEXT NOT NULL,
+                FOREIGN KEY (memory_id) REFERENCES memory_items(id) ON DELETE CASCADE
+            );
+
+            CREATE INDEX IF NOT EXISTS idx_memory_items_task ON memory_items(task_id);
+            CREATE INDEX IF NOT EXISTS idx_memory_items_session ON memory_items(session_id);
+            CREATE INDEX IF NOT EXISTS idx_memory_items_kind ON memory_items(kind);
+            CREATE INDEX IF NOT EXISTS idx_memory_embeddings_model ON memory_embeddings(embedding_model);
+        """)
+        self._conn.commit()
+
+    def insert(self, entry: MemoryEntry, embedding: bytes) -> None:
+        cursor = self._conn.cursor()
+        cursor.execute(
+            """
+            INSERT INTO memory_items (id, text, kind, source, weight, task_id, session_id, metadata_json, created_at, updated_at)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                entry.id,
+                entry.text,
+                entry.kind,
+                entry.source,
+                entry.weight,
+                entry.task_id,
+                entry.session_id,
+                json.dumps(entry.metadata) if entry.metadata else None,
+                entry.created_at.isoformat(),
+                utc_now().isoformat(),
+            ),
+        )
+        cursor.execute(
+            """
+            INSERT INTO memory_embeddings (memory_id, embedding, embedding_model, embedding_dim, created_at)
+            VALUES (?, ?, ?, ?, ?)
+            """,
+            (
+                entry.id,
+                embedding,
+                entry.embedding_model,
+                entry.embedding_dim,
+                utc_now().isoformat(),
+            ),
+        )
+        self._conn.commit()
+
+    def get(self, memory_id: str) -> MemoryEntry | None:
+        cursor = self._conn.cursor()
+        row = cursor.execute(
+            "SELECT * FROM memory_items WHERE id = ?", (memory_id,)
+        ).fetchone()
+        if not row:
+            return None
+        return self._row_to_entry(row)
+
+    def get_embedding(self, memory_id: str) -> bytes | None:
+        cursor = self._conn.cursor()
+        row = cursor.execute(
+            "SELECT embedding FROM memory_embeddings WHERE memory_id = ?", (memory_id,)
+        ).fetchone()
+        return bytes(row["embedding"]) if row else None
+
+    def get_all(self, limit: int = 1000) -> list[MemoryEntry]:
+        cursor = self._conn.cursor()
+        rows = cursor.execute(
+            "SELECT * FROM memory_items ORDER BY created_at DESC LIMIT ?", (limit,)
+        ).fetchall()
+        return [self._row_to_entry(row) for row in rows]
+
+    def get_by_task(self, task_id: str) -> list[MemoryEntry]:
+        cursor = self._conn.cursor()
+        rows = cursor.execute(
+            "SELECT * FROM memory_items WHERE task_id = ? ORDER BY created_at DESC", (task_id,)
+        ).fetchall()
+        return [self._row_to_entry(row) for row in rows]
+
+    def get_by_session(self, session_id: str, limit: int = 100) -> list[MemoryEntry]:
+        cursor = self._conn.cursor()
+        rows = cursor.execute(
+            "SELECT * FROM memory_items WHERE session_id = ? ORDER BY created_at DESC LIMIT ?",
+            (session_id, limit),
+        ).fetchall()
+        return [self._row_to_entry(row) for row in rows]
+
+    def get_by_kind(self, kind: str, limit: int = 100) -> list[MemoryEntry]:
+        cursor = self._conn.cursor()
+        rows = cursor.execute(
+            "SELECT * FROM memory_items WHERE kind = ? ORDER BY created_at DESC LIMIT ?", (kind, limit)
+        ).fetchall()
+        return [self._row_to_entry(row) for row in rows]
+
+    def delete(self, memory_id: str) -> bool:
+        cursor = self._conn.cursor()
+        cursor.execute("DELETE FROM memory_embeddings WHERE memory_id = ?", (memory_id,))
+        cursor.execute("DELETE FROM memory_items WHERE id = ?", (memory_id,))
+        self._conn.commit()
+        return cursor.rowcount > 0
+
+    def update_weight(self, memory_id: str, weight: float) -> bool:
+        cursor = self._conn.cursor()
+        cursor.execute(
+            "UPDATE memory_items SET weight = ?, updated_at = ? WHERE id = ?",
+            (weight, utc_now().isoformat(), memory_id),
+        )
+        self._conn.commit()
+        return cursor.rowcount > 0
+
+    def search_text(self, query: str, limit: int = 10) -> list[MemoryEntry]:
+        cursor = self._conn.cursor()
+        rows = cursor.execute(
+            "SELECT * FROM memory_items WHERE text LIKE ? ORDER BY created_at DESC LIMIT ?",
+            (f"%{query}%", limit),
+        ).fetchall()
+        return [self._row_to_entry(row) for row in rows]
+
+    def count(self) -> int:
+        cursor = self._conn.cursor()
+        row = cursor.execute("SELECT COUNT(*) FROM memory_items").fetchone()
+        return row[0] if row else 0
+
+    def close(self) -> None:
+        self._conn.close()
+
+    def _row_to_entry(self, row: sqlite3.Row) -> MemoryEntry:
+        metadata = {}
+        if row["metadata_json"]:
+            import json
+            metadata = json.loads(row["metadata_json"])
+        return MemoryEntry(
+            id=row["id"],
+            text=row["text"],
+            kind=row["kind"],
+            source=row["source"],
+            weight=row["weight"],
+            task_id=row["task_id"],
+            session_id=row["session_id"],
+            metadata=metadata,
+            created_at=datetime.fromisoformat(row["created_at"]),
+            embedding_model="",
+            embedding_dim=0,
+        )
@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import logging
+import numpy as np
+import hnswlib
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class VectorIndex:
+    def __init__(
+        self,
+        index_path: str | Path | None = None,
+        embedding_dim: int = 384,
+        max_elements: int = 10000,
+    ) -> None:
+        self._embedding_dim = embedding_dim
+        self._index_path = Path(index_path) if index_path else None
+        self._index: hnswlib.Index | None = None
+        self._max_elements = max_elements
+        self._loading = False  # Prevent recursion
+
+        self._init_index()
+
+    def _init_index(self) -> None:
+        if self._loading:
+            return
+        self._loading = True
+        try:
+            if self._index_path and self._index_path.exists():
+                self._load()
+            else:
+                self._index = hnswlib.Index(
+                    space="l2",
+                    dim=self._embedding_dim,
+                )
+                self._index.init_index(
+                    max_elements=self._max_elements,
+                    ef_construction=200,
+                    M=16,
+                )
+        except Exception as e:
+            logger.warning(f"VectorIndex init failed: {e}")
+            self._index = hnswlib.Index(
+                space="l2",
+                dim=self._embedding_dim,
+            )
+            self._index.init_index(
+                max_elements=self._max_elements,
+                ef_construction=100,
+                M=16,
+            )
+        finally:
+            self._loading = False
+
+    def insert(self, memory_id: str, embedding: np.ndarray) -> None:
+        if self._index is None:
+            self._init_index()
+            if self._index is None:
+                return
+
+        try:
+            vector = self._normalize(embedding)
+            internal_id = self._get_internal_id(memory_id)
+            self._index.add_items(vector, ids=np.array([internal_id]))
+        except Exception as e:
+            logger.warning(f"VectorIndex insert failed: {e}")
+
+    def search(
+        self,
+        query_embedding: np.ndarray,
+        k: int = 5,
+    ) -> tuple[list[str], list[float]]:
+        if self._index is None:
+            return [], []
+
+        try:
+            if self._index.get_current_count() == 0:
+                return [], []
+
+            # Set ef to at least k for proper search
+            self._index.set_ef(max(k * 2, 50))
+
+            vector = self._normalize(query_embedding)
+            labels, distances = self._index.knn_query(vector, k=k)
+
+            memory_ids = [self._get_memory_id(int(label)) for label in labels[0]]
+            scores = [1.0 - dist for dist in distances[0]]
+            return memory_ids, scores
+        except Exception as e:
+            logger.warning(f"VectorIndex search failed: {e}")
+            return [], []
+
+    def delete(self, memory_id: str) -> bool:
+        return False
+
+    def get_items(self, memory_ids: list[str]) -> np.ndarray:
+        if self._index is None:
+            raise RuntimeError("Index not initialized")
+        internal_ids = [self._get_internal_id(mid) for mid in memory_ids]
+        return self._index.get_items(np.array(internal_ids))
+
+    def save(self) -> None:
+        if self._index and self._index_path:
+            try:
+                self._index_path.parent.mkdir(parents=True, exist_ok=True)
+                self._index.save_index(str(self._index_path))
+            except Exception as e:
+                logger.warning(f"VectorIndex save failed: {e}")
+
+    def _load(self) -> None:
+        if self._loading:
+            return
+        self._loading = True
+        try:
+            if self._index_path and self._index_path.exists():
+                self._index = hnswlib.Index(space="l2", dim=self._embedding_dim)
+                self._index.load_index(
+                    str(self._index_path),
+                    max_elements=self._max_elements
+                )
+        except Exception as e:
+            logger.warning(f"VectorIndex load failed: {e}")
+            self._init_index()
+        finally:
+            self._loading = False
+
+    def _normalize(self, vector: np.ndarray) -> np.ndarray:
+        vec = vector.flatten()
+        norm = np.linalg.norm(vec)
+        if norm > 0:
+            vec = vec / norm
+        return vec.reshape(1, -1)
+
+    def _get_internal_id(self, memory_id: str) -> int:
+        return hash(memory_id) % (2**31)
+
+    def _get_memory_id(self, internal_id: int) -> str:
+        return str(internal_id)
+
+    @property
+    def embedding_dim(self) -> int:
+        return self._embedding_dim
+
+    @property
+    def element_count(self) -> int:
+        return self._index.get_current_count() if self._index else 0
@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from app.core.contracts import CriticScore, MemoryEntry
+
+
+class MemoryWritePolicy:
+    def __init__(
+        self,
+        store_threshold: float = 0.7,
+        min_usefulness: float = 0.3,
+        max_entries_per_session: int = 50,
+    ) -> None:
+        self._store_threshold = store_threshold
+        self._min_usefulness = min_usefulness
+        self._max_entries_per_session = max_entries_per_session
+
+    def decide(
+        self,
+        critic_score: CriticScore,
+        memory_type: MemoryEntry.Kind,
+        session_id: str | None = None,
+        has_duplicate: bool = False,
+        current_session_count: int = 0,
+    ) -> Literal["store", "store_with_weight", "skip", "merge"]:
+        if critic_score.safety < 0.5:
+            return "skip"
+
+        if has_duplicate:
+            return "merge"
+
+        if not critic_score.memory_store:
+            return "skip"
+
+        if critic_score.usefulness < self._min_usefulness:
+            return "skip"
+
+        if session_id and current_session_count >= self._max_entries_per_session:
+            return "skip"
+
+        base_decision = self._evaluate_scores(critic_score, memory_type)
+
+        if base_decision == "store" and critic_score.weight < self._store_threshold:
+            adjusted_weight = self._adjust_weight(critic_score, memory_type)
+            if adjusted_weight >= self._store_threshold:
+                return "store_with_weight"
+            return base_decision
+
+        return base_decision
+
+    def _evaluate_scores(
+        self,
+        critic_score: CriticScore,
+        memory_type: MemoryEntry.Kind,
+    ) -> Literal["store", "store_with_weight", "skip", "merge"]:
+        avg_score = (critic_score.correctness + critic_score.usefulness + critic_score.safety) / 3.0
+
+        if memory_type in ("fact", "plan", "summary"):
+            if avg_score >= 0.8:
+                return "store"
+            elif avg_score >= 0.6:
+                return "store_with_weight"
+
+        if memory_type in ("tool_result", "critique"):
+            if avg_score >= self._store_threshold:
+                return "store"
+            elif avg_score >= 0.5:
+                return "store_with_weight"
+
+        if memory_type == "user_preference":
+            if avg_score >= 0.5:
+                return "store"
+
+        return "skip"
+
+    def _adjust_weight(
+        self,
+        critic_score: CriticScore,
+        memory_type: MemoryEntry.Kind,
+    ) -> float:
+        base_weight = critic_score.weight
+
+        type_boost = {
+            "fact": 0.15,
+            "plan": 0.1,
+            "summary": 0.1,
+            "user_preference": 0.2,
+            "tool_result": 0.05,
+            "critique": 0.05,
+        }.get(memory_type, 0.0)
+
+        safety_boost = 0.0
+        if critic_score.safety >= 0.9:
+            safety_boost = 0.1
+
+        adjusted = base_weight + type_boost + safety_boost
+        return min(adjusted, 1.0)