Update utils.py

utils.py

@@ -2,6 +2,7 @@ import os
 import json
 import faiss
 import numpy as np
+from uuid import uuid4
 from datetime import datetime
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -13,7 +14,8 @@ MEMORY_INDEX_PATH = "memory.index"
 MEMORY_TEXTS_PATH = "memory_texts.json"
 CHAT_LOG_PATH = "chatlog.jsonl"
 FEEDBACK_PATH = "feedback.jsonl"
-SUMMARY_TRIGGER = 100
+SUMMARY_TRIGGER = int(os.getenv("SUMMARY_TRIGGER", 100))
+CHUNK_SIZE = int(os.getenv("SUMMARY_CHUNK", 10))
 
 # === Load models ===
 embedder = SentenceTransformer(EMBED_MODEL)
@@ -23,25 +25,31 @@ summary_model = AutoModelForCausalLM.from_pretrained(SUMMARIZER_MODEL).eval()
 embedding_dim = embedder.get_sentence_embedding_dimension()
 
 # === Memory state ===
-if os.path.exists(MEMORY_INDEX_PATH) and os.path.exists(MEMORY_TEXTS_PATH):
-    memory_index = faiss.read_index(MEMORY_INDEX_PATH)
-    with open(MEMORY_TEXTS_PATH, "r") as f:
-        memory_texts = json.load(f)
-else:
+try:
+    if os.path.exists(MEMORY_INDEX_PATH) and os.path.exists(MEMORY_TEXTS_PATH):
+        memory_index = faiss.read_index(MEMORY_INDEX_PATH)
+        with open(MEMORY_TEXTS_PATH, "r") as f:
+            memory_texts = json.load(f)
+    else:
+        raise FileNotFoundError
+except:
     memory_index = faiss.IndexFlatL2(embedding_dim)
     memory_texts = []
 
-def embed(text):
-    """Embed text into vector"""
-    return embedder.encode([text])
+def embed(texts):
+    """Embed a list of texts into vectors"""
+    return embedder.encode(texts)
 
 def add_to_memory(text):
     """Add a memory item"""
-    vec = embed(text)
+    vec = embed([text])
     memory_index.add(np.array(vec))
-    memory_texts.append(text)
+    memory_texts.append({
+        "id": str(uuid4()),
+        "text": text,
+        "timestamp": datetime.now().isoformat()
+    })
     save_memory()
-
     if len(memory_texts) > SUMMARY_TRIGGER:
         summarize_old_memories()
 
@@ -49,36 +57,40 @@ def retrieve_memories(query, k=3):
     """Retrieve top relevant memories"""
     if memory_index.ntotal == 0:
         return []
-    vec = embed(query)
+    vec = embed([query])
     D, I = memory_index.search(np.array(vec), k)
-    return [memory_texts[i] for i in I[0] if i < len(memory_texts)]
+    return [memory_texts[i]["text"] for i in I[0] if i < len(memory_texts)]
 
 def save_memory():
     """Save FAISS and text memory to disk"""
     faiss.write_index(memory_index, MEMORY_INDEX_PATH)
     with open(MEMORY_TEXTS_PATH, "w") as f:
         json.dump(memory_texts, f)
+    print(f"[INFO] Memory saved: {len(memory_texts)} items")
 
 def reset_memory():
     """Reset memory entirely"""
     memory_index.reset()
     memory_texts.clear()
-    if os.path.exists(MEMORY_INDEX_PATH):
-        os.remove(MEMORY_INDEX_PATH)
-    if os.path.exists(MEMORY_TEXTS_PATH):
-        os.remove(MEMORY_TEXTS_PATH)
+    for path in [MEMORY_INDEX_PATH, MEMORY_TEXTS_PATH]:
+        if os.path.exists(path):
+            os.remove(path)
 
 def summarize_old_memories():
     """Replace older entries with a summary to save space"""
-    old = "\n".join(memory_texts[:10])
+    old = "\n".join(m["text"] for m in memory_texts[:CHUNK_SIZE])
     inputs = summary_tokenizer(f"Summarize: {old}", return_tensors="pt")
     output = summary_model.generate(**inputs, max_new_tokens=100)
     summary = summary_tokenizer.decode(output[0][inputs['input_ids'].shape[-1]:], skip_special_tokens=True)
 
-    memory_texts[:10] = [summary]
+    memory_texts[:CHUNK_SIZE] = [{
+        "id": str(uuid4()),
+        "text": summary,
+        "timestamp": datetime.now().isoformat()
+    }]
     memory_index.reset()
-    for text in memory_texts:
-        vec = embed(text)
+    for mem in memory_texts:
+        vec = embed([mem["text"]])
         memory_index.add(np.array(vec))
     save_memory()
 
@@ -109,4 +121,4 @@ def generate_suggestions(feedback_text):
     outputs = summary_model.generate(**inputs, max_new_tokens=100)
     suggestions = summary_tokenizer.decode(outputs[0][inputs['input_ids'].shape[-1]:], skip_special_tokens=True)
     log_feedback(feedback_text, suggestions)
-    return suggestions
+    return suggestions