trace-memory 1.0.4__tar.gz → 1.0.6__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: trace-memory
3
- Version: 1.0.4
3
+ Version: 1.0.6
4
4
  Summary: TRACE — Temporal Retrieval And Context Engine: self-healing hierarchical memory for LLM agents.
5
5
  License: Apache-2.0
6
6
  Project-URL: Homepage, https://github.com/husain34/TRACE
@@ -25,6 +25,7 @@ License-File: NOTICE
25
25
  Requires-Dist: openai>=1.0.0
26
26
  Requires-Dist: python-dotenv>=1.0.0
27
27
  Requires-Dist: sentence-transformers>=2.2.0
28
+ Requires-Dist: numpy>=1.21.0
28
29
  Dynamic: license-file
29
30
 
30
31
  <div align="center">
@@ -208,14 +209,13 @@ Result: The AI aggressively stopped the user.
208
209
  ```
209
210
  Phase 1 — Collect all frozen (inactive) TopicNodes
210
211
  Phase 2 — Generate missing summaries; embed all candidates
211
- Phase 3 — Compute pairwise cosine similarity
212
- Phase 4 — For each pair above threshold, apply 4 axioms:
212
+ Phase 3 — Compute pairwise cosine similarity; for each pair above threshold, apply 4 axioms:
213
213
  Axiom 1 — Chronological Guard
214
214
  Axiom 2 — Frozen State check
215
215
  Axiom 3 — Similarity threshold (default 0.55)
216
216
  Axiom 4 — LLM Veto
217
217
  If all 4 pass → merge (newer becomes child of older)
218
- Phase 5 — Optional: prune trivial leaf messages
218
+ Phase 4 — Optional: prune trivial leaf messages
219
219
  ```
220
220
 
221
221
  #### The Four Axioms in Detail
@@ -403,16 +403,13 @@ tree.save("sessions/chat_001.json", save_conversation=True)
403
403
 
404
404
  `save_conversation=True` embeds the raw message list so the session can be fully restored later.
405
405
 
406
- ---
407
-
408
- ##### `CTree.load(filepath: str, api_key: str = None, base_url: str = None, model: str = None) -> CTree`
406
+ ##### `CTree.load(filepath: str, api_key: str = None, base_url: str = None, model: str = None, embed_fn: Callable = None) -> CTree`
409
407
 
410
408
  Restore a tree from a JSON file.
411
409
 
412
410
  ```python
413
411
  tree = CTree.load("sessions/chat_001.json", api_key="sk-...", embed_fn=embed)
414
- tree.vdb = VectorDatabase("sessions/chat_001.db")
415
-
412
+ tree.vdb = VectorDatabase("sessions/chat_001.db")
416
413
  ```
417
414
 
418
415
  ---
@@ -473,7 +470,7 @@ ROOT (sub-nodes: 3)
473
470
 
474
471
  A local SQLite vector store with two active tables: conversation vectors and topic summaries.
475
472
 
476
- > **Note on Scaling:** TRACE uses pure Python (`sqlite3` and `struct`) to avoid heavy dependencies and run instantly on any machine. This is perfectly sufficient for local agents and moderate conversation histories. However, for massive scale (millions of vectors), you will need to swap this module for a dedicated vector database like FAISS or Chroma.
473
+ > **Note on Scaling:** TRACE uses SQLite (`sqlite3`) for storage, `numpy` for fast cosine similarity, and `struct` for compact binary packing all chosen to keep the footprint small while running on any machine. For massive scale (millions of vectors), swap this module for FAISS or Chroma.
477
474
 
478
475
  ```python
479
476
  from trace_memory import VectorDatabase
@@ -487,8 +484,6 @@ vdb = VectorDatabase("path/to/session.db") # creates the DB if it doesn't exist
487
484
 
488
485
  ---
489
486
 
490
- ---
491
-
492
487
  #### Conversation methods
493
488
 
494
489
  ##### `vdb.add_conversation_message(msg: ConversationVector)`
@@ -573,6 +568,7 @@ system_prompt = synth.synthesize_prompt(
573
568
  query_vector = embed("Is the cake safe for Sarah?"),
574
569
  active_node = tree.current_node,
575
570
  recent_messages = tree.conversation[-6:],
571
+ top_k_docs = 3, # max topic branch paths to surface
576
572
  top_k_history = 2, # max past messages to recall
577
573
  min_history_similarity = 0.50, # min cosine score for conversation recall
578
574
  )
@@ -641,15 +637,25 @@ tree = CTree(api_key="sk-...", model="gpt-4o-mini")
641
637
 
642
638
  ### 7.3 With NVIDIA NIM / any OpenAI-compatible endpoint
643
639
 
644
- ```python
645
- import os
646
- os.environ["OPENAI_BASE_URL"] = "https://integrate.api.nvidia.com/v1"
647
- os.environ["OPENAI_API_KEY"] = "nvapi-..."
640
+ Point the `base_url` to your endpoint.
648
641
 
649
- from trace_memory import CTree
650
- tree = CTree(model="meta/llama-3.1-70b-instruct")
642
+ ```python
643
+ tree = CTree(
644
+ api_key="your-nvidia-api-key",
645
+ base_url="https://integrate.api.nvidia.com/v1",
646
+ model="meta/llama-3.1-8b-instruct"
647
+ )
651
648
  ```
652
649
 
650
+ ### 7.4 Best Practices for LLM Selection
651
+
652
+ **Do not use "Reasoning" models (DeepSeek R1, Claude 3.7 Sonnet thinking mode, o1) for the internal `CTree` engine.**
653
+
654
+ When `CTree` runs in the background, it performs simple, deterministic classification tasks (e.g., naming a topic in 3 words, or extracting a JSON boolean). Reasoning models are built for complex logic and will waste massive amounts of compute "thinking" about simple tasks. Worse, reasoning models often exceed TRACE's strict internal token limits (`max_tokens=50` or `200`) designed to keep background operations fast, causing the LLM to get cut off mid-thought and breaking the internal JSON parsers.
655
+
656
+ * **For the `CTree` internal model:** Always use a fast, standard model (e.g., `gpt-4o-mini`, `meta-llama-3.1-8b-instruct`).
657
+ * **For your actual chat application:** You can (and should!) use whatever massive reasoning model you prefer to generate the final response to the user. TRACE is completely decoupled from your front-end chat completion.
658
+
653
659
  ---
654
660
 
655
661
  ## 8. Edge-Case Stress Tests & Validation
@@ -684,6 +690,8 @@ TRACE loads `.env` automatically if `python-dotenv` is installed.
684
690
  |---|---|---|
685
691
  | `openai` | ≥ 1.0.0 | ✅ Yes |
686
692
  | `python-dotenv` | ≥ 1.0.0 | ✅ Yes |
693
+ | `sentence-transformers` | ≥ 2.2.0 | ✅ Yes (local embed fallback) |
694
+ | `numpy` | ≥ 1.21.0 | ✅ Yes (cosine similarity) |
687
695
  | `sqlite3` | built-in | ✅ Yes (no install needed) |
688
696
  | `struct` | built-in | ✅ Yes (no install needed) |
689
697
 
@@ -179,14 +179,13 @@ Result: The AI aggressively stopped the user.
179
179
  ```
180
180
  Phase 1 — Collect all frozen (inactive) TopicNodes
181
181
  Phase 2 — Generate missing summaries; embed all candidates
182
- Phase 3 — Compute pairwise cosine similarity
183
- Phase 4 — For each pair above threshold, apply 4 axioms:
182
+ Phase 3 — Compute pairwise cosine similarity; for each pair above threshold, apply 4 axioms:
184
183
  Axiom 1 — Chronological Guard
185
184
  Axiom 2 — Frozen State check
186
185
  Axiom 3 — Similarity threshold (default 0.55)
187
186
  Axiom 4 — LLM Veto
188
187
  If all 4 pass → merge (newer becomes child of older)
189
- Phase 5 — Optional: prune trivial leaf messages
188
+ Phase 4 — Optional: prune trivial leaf messages
190
189
  ```
191
190
 
192
191
  #### The Four Axioms in Detail
@@ -374,16 +373,13 @@ tree.save("sessions/chat_001.json", save_conversation=True)
374
373
 
375
374
  `save_conversation=True` embeds the raw message list so the session can be fully restored later.
376
375
 
377
- ---
378
-
379
- ##### `CTree.load(filepath: str, api_key: str = None, base_url: str = None, model: str = None) -> CTree`
376
+ ##### `CTree.load(filepath: str, api_key: str = None, base_url: str = None, model: str = None, embed_fn: Callable = None) -> CTree`
380
377
 
381
378
  Restore a tree from a JSON file.
382
379
 
383
380
  ```python
384
381
  tree = CTree.load("sessions/chat_001.json", api_key="sk-...", embed_fn=embed)
385
- tree.vdb = VectorDatabase("sessions/chat_001.db")
386
-
382
+ tree.vdb = VectorDatabase("sessions/chat_001.db")
387
383
  ```
388
384
 
389
385
  ---
@@ -444,7 +440,7 @@ ROOT (sub-nodes: 3)
444
440
 
445
441
  A local SQLite vector store with two active tables: conversation vectors and topic summaries.
446
442
 
447
- > **Note on Scaling:** TRACE uses pure Python (`sqlite3` and `struct`) to avoid heavy dependencies and run instantly on any machine. This is perfectly sufficient for local agents and moderate conversation histories. However, for massive scale (millions of vectors), you will need to swap this module for a dedicated vector database like FAISS or Chroma.
443
+ > **Note on Scaling:** TRACE uses SQLite (`sqlite3`) for storage, `numpy` for fast cosine similarity, and `struct` for compact binary packing all chosen to keep the footprint small while running on any machine. For massive scale (millions of vectors), swap this module for FAISS or Chroma.
448
444
 
449
445
  ```python
450
446
  from trace_memory import VectorDatabase
@@ -458,8 +454,6 @@ vdb = VectorDatabase("path/to/session.db") # creates the DB if it doesn't exist
458
454
 
459
455
  ---
460
456
 
461
- ---
462
-
463
457
  #### Conversation methods
464
458
 
465
459
  ##### `vdb.add_conversation_message(msg: ConversationVector)`
@@ -544,6 +538,7 @@ system_prompt = synth.synthesize_prompt(
544
538
  query_vector = embed("Is the cake safe for Sarah?"),
545
539
  active_node = tree.current_node,
546
540
  recent_messages = tree.conversation[-6:],
541
+ top_k_docs = 3, # max topic branch paths to surface
547
542
  top_k_history = 2, # max past messages to recall
548
543
  min_history_similarity = 0.50, # min cosine score for conversation recall
549
544
  )
@@ -612,15 +607,25 @@ tree = CTree(api_key="sk-...", model="gpt-4o-mini")
612
607
 
613
608
  ### 7.3 With NVIDIA NIM / any OpenAI-compatible endpoint
614
609
 
615
- ```python
616
- import os
617
- os.environ["OPENAI_BASE_URL"] = "https://integrate.api.nvidia.com/v1"
618
- os.environ["OPENAI_API_KEY"] = "nvapi-..."
610
+ Point the `base_url` to your endpoint.
619
611
 
620
- from trace_memory import CTree
621
- tree = CTree(model="meta/llama-3.1-70b-instruct")
612
+ ```python
613
+ tree = CTree(
614
+ api_key="your-nvidia-api-key",
615
+ base_url="https://integrate.api.nvidia.com/v1",
616
+ model="meta/llama-3.1-8b-instruct"
617
+ )
622
618
  ```
623
619
 
620
+ ### 7.4 Best Practices for LLM Selection
621
+
622
+ **Do not use "Reasoning" models (DeepSeek R1, Claude 3.7 Sonnet thinking mode, o1) for the internal `CTree` engine.**
623
+
624
+ When `CTree` runs in the background, it performs simple, deterministic classification tasks (e.g., naming a topic in 3 words, or extracting a JSON boolean). Reasoning models are built for complex logic and will waste massive amounts of compute "thinking" about simple tasks. Worse, reasoning models often exceed TRACE's strict internal token limits (`max_tokens=50` or `200`) designed to keep background operations fast, causing the LLM to get cut off mid-thought and breaking the internal JSON parsers.
625
+
626
+ * **For the `CTree` internal model:** Always use a fast, standard model (e.g., `gpt-4o-mini`, `meta-llama-3.1-8b-instruct`).
627
+ * **For your actual chat application:** You can (and should!) use whatever massive reasoning model you prefer to generate the final response to the user. TRACE is completely decoupled from your front-end chat completion.
628
+
624
629
  ---
625
630
 
626
631
  ## 8. Edge-Case Stress Tests & Validation
@@ -655,6 +660,8 @@ TRACE loads `.env` automatically if `python-dotenv` is installed.
655
660
  |---|---|---|
656
661
  | `openai` | ≥ 1.0.0 | ✅ Yes |
657
662
  | `python-dotenv` | ≥ 1.0.0 | ✅ Yes |
663
+ | `sentence-transformers` | ≥ 2.2.0 | ✅ Yes (local embed fallback) |
664
+ | `numpy` | ≥ 1.21.0 | ✅ Yes (cosine similarity) |
658
665
  | `sqlite3` | built-in | ✅ Yes (no install needed) |
659
666
  | `struct` | built-in | ✅ Yes (no install needed) |
660
667
 
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "trace-memory"
7
- version = "1.0.4"
7
+ version = "1.0.6"
8
8
  description = "TRACE — Temporal Retrieval And Context Engine: self-healing hierarchical memory for LLM agents."
9
9
  readme = "README.md"
10
10
  license = { text = "Apache-2.0" }
@@ -30,9 +30,10 @@ classifiers = [
30
30
  ]
31
31
 
32
32
  dependencies = [
33
- "openai>=1.0.0", # LLM API client (supports any OpenAI-compatible endpoint)
34
- "python-dotenv>=1.0.0", # .env file support
35
- "sentence-transformers>=2.2.0", # local fallback for embeddings
33
+ "openai>=1.0.0", # LLM API client (supports any OpenAI-compatible endpoint)
34
+ "python-dotenv>=1.0.0", # .env file support
35
+ "sentence-transformers>=2.2.0", # local fallback embedder (BAAI/bge-base-en-v1.5)
36
+ "numpy>=1.21.0", # cosine similarity in vector_db
36
37
  ]
37
38
 
38
39
 
@@ -25,7 +25,7 @@ from .ctree import CTree, Node, TopicNode, MessageNode
25
25
  from .vector_db import VectorDatabase, ConversationVector
26
26
  from .prompt_synthesizer import PromptSynthesizer
27
27
 
28
- __version__ = "1.0.0"
28
+ __version__ = "1.0.6"
29
29
  __author__ = "Husain"
30
30
  __license__ = "Apache-2.0"
31
31
 
@@ -76,7 +76,6 @@ def _normalize_for_json(text: str) -> str:
76
76
  """Normalise common LLM JSON formatting quirks before parsing."""
77
77
  text = text.replace(": N/A", ": null")
78
78
  text = text.replace(": True", ": true").replace(": False", ": false").replace(": None", ": null")
79
- text = text.replace(": True,", ": true,").replace(": False,", ": false,").replace(": None,", ": null,")
80
79
  text = text.replace(",}", "}").replace(",]", "]")
81
80
  text = text.replace("\r\n", " ").replace("\n", " ").replace("\r", " ")
82
81
  text = " ".join(text.split())
@@ -13,7 +13,7 @@ exchange into a tree of named topic branches, with:
13
13
 
14
14
  Typical usage
15
15
  -------------
16
- from trace import CTree, VectorDatabase, PromptSynthesizer
16
+ from trace_memory import CTree, VectorDatabase, PromptSynthesizer
17
17
 
18
18
  # 1. Boot the engine
19
19
  tree = CTree(api_key="sk-...", model="gpt-4o-mini")
@@ -54,6 +54,7 @@ from typing import List, Dict, Optional, Tuple, Any, Callable
54
54
  from dataclasses import dataclass, field
55
55
 
56
56
  from ._llm_utils import ChatGPT_API, extract_json
57
+ from .vector_db import cosine_similarity
57
58
 
58
59
  # ── Helpers ───────────────────────────────────────────────────────────────────
59
60
 
@@ -94,16 +95,20 @@ class Node:
94
95
  ----------
95
96
  children : Direct child nodes.
96
97
  parent : Parent node (None for root).
97
- sub_node_count : Cached count of direct children.
98
+ sub_node_count : Live count of direct children (dynamic property).
98
99
  created_at : Unix timestamp of node birth (set once, never changed).
99
100
  """
100
101
 
101
- def __init__(self, children=None, parent=None, sub_node_count=0):
102
+ def __init__(self, children=None, parent=None):
102
103
  self.children: List["Node"] = children or []
103
104
  self.parent: Optional["Node"] = parent
104
- self.sub_node_count: int = sub_node_count
105
105
  self.created_at: float = time.time()
106
106
 
107
+ @property
108
+ def sub_node_count(self) -> int:
109
+ """Dynamically computed — always reflects the real number of direct children."""
110
+ return len(self.children)
111
+
107
112
  def is_leaf(self) -> bool:
108
113
  return len(self.children) == 0
109
114
 
@@ -365,8 +370,9 @@ class CTree:
365
370
  end_index = node.end_index,
366
371
  depth = depth,
367
372
  )
368
- except Exception:
369
- pass # embedding failures must never crash the tree
373
+ except Exception as e:
374
+ import sys
375
+ print(f"[TRACE] \u26a0 VDB upsert failed for node '{node.topic_name}': {e}", file=sys.stderr)
370
376
 
371
377
  # ── Public: add exchange ──────────────────────────────────────────────────
372
378
 
@@ -541,34 +547,8 @@ class CTree:
541
547
  parent.children.append(new_topic)
542
548
  return new_topic
543
549
 
544
- def _create_new_node(self, message, start_index, parent, topic_name=""):
545
- if not topic_name or not topic_name.strip():
546
- topic_name = self._llm_generate_topic(message, parent)
547
- new_topic = TopicNode(
548
- topic_name=topic_name,
549
- start_index=start_index,
550
- end_index=start_index,
551
- parent=parent,
552
- )
553
- parent.children.append(new_topic)
554
- return new_topic
555
-
556
550
  # ── LLM helpers ───────────────────────────────────────────────────────────
557
551
 
558
- def _llm_generate_topic(self, message, parent=None) -> str:
559
- content = message.get("content", "")
560
- context = f"\nParent topic: {parent.topic_name}" if (parent and parent.topic_name != "ROOT") else ""
561
- prompt = (
562
- f"Given the following message, generate a concise topic name "
563
- f"(2-5 words) that captures its main subject.{context}\n\n"
564
- f"Message: {content}\n\nRespond with ONLY the topic name, nothing else."
565
- )
566
- try:
567
- return ChatGPT_API(self.model, prompt, api_key=self.api_key, base_url=self.base_url, temperature=0.3, max_tokens=50).strip()
568
- except Exception as e:
569
- print(f"LLM error in topic generation: {e}")
570
- return f"Topic at index {len(self.conversation)}"
571
-
572
552
  def _llm_generate_topic_from_message(
573
553
  self, user_msg, assistant_msg, parent=None, system_msg=None
574
554
  ) -> str:
@@ -829,10 +809,11 @@ class CTree:
829
809
  def _in_subtree(n, target):
830
810
  if n is target: return True
831
811
  return any(_in_subtree(c, target) for c in n.children)
832
- for sn in root_node.children:
833
- if isinstance(sn, TopicNode) and _in_subtree(sn, self.current_node):
834
- break
835
- else:
812
+ current_found = any(
813
+ isinstance(child, TopicNode) and _in_subtree(child, self.current_node)
814
+ for child in root_node.children
815
+ )
816
+ if not current_found:
836
817
  self.current_node = last_subtopic_node
837
818
 
838
819
  # ── Public utilities ──────────────────────────────────────────────────────
@@ -951,38 +932,42 @@ class CTree:
951
932
  def _reconstruct_node(self, node_data: dict, parent):
952
933
  node_type = node_data.get("type", "topic")
953
934
  if node_type == "message":
954
- msg_start_idx = node_data.get("message_index", node_data.get("pair_index", 0) * 2)
955
- user_msg = assistant_msg = {}
935
+ msg_idx = node_data.get("message_index", 0)
936
+ user_msg = assistant_msg = {}
956
937
  system_msg = None
957
- for i in range(max(0, msg_start_idx - 1), min(len(self.conversation), msg_start_idx + 4)):
958
- if i < len(self.conversation) and self.conversation[i].get("role") == "user":
959
- if i > 0 and self.conversation[i - 1].get("role") == "system":
960
- system_msg = self.conversation[i - 1]
961
- user_msg = self.conversation[i]
962
- if i + 1 < len(self.conversation) and self.conversation[i + 1].get("role") == "assistant":
963
- assistant_msg = self.conversation[i + 1]
938
+ if self.conversation:
939
+ for i in range(max(0, msg_idx - 1), min(len(self.conversation), msg_idx + 4)):
940
+ if self.conversation[i].get("role") == "user":
941
+ user_msg = self.conversation[i]
942
+ if i > 0 and self.conversation[i - 1].get("role") == "system":
943
+ system_msg = self.conversation[i - 1]
944
+ if i + 1 < len(self.conversation) and self.conversation[i + 1].get("role") == "assistant":
945
+ assistant_msg = self.conversation[i + 1]
964
946
  break
965
- node = MessageNode(
966
- user_message = user_msg,
967
- assistant_message = assistant_msg,
968
- system_message = system_msg,
969
- message_index = msg_start_idx,
970
- parent = parent,
971
- )
972
- else:
973
- node = TopicNode(
974
- topic_name = node_data.get("topic_name", ""),
975
- summary = node_data.get("summary", ""),
976
- start_index = node_data.get("start_index", 0),
977
- end_index = node_data.get("end_index", 0),
978
- parent = parent,
947
+ else:
948
+ import warnings
949
+ warnings.warn(
950
+ "[TRACE] Tree loaded without conversation history (save_conversation=False). "
951
+ "MessageNode content will be empty. Re-save with save_conversation=True.",
952
+ stacklevel=2,
953
+ )
954
+ return MessageNode(
955
+ user_message=user_msg, assistant_message=assistant_msg,
956
+ system_message=system_msg, message_index=msg_idx, parent=parent,
979
957
  )
980
- if "node_id" in node_data:
981
- node.node_id = node_data["node_id"]
982
- node.created_at = node_data.get("created_at", node.created_at)
983
- for child_data in node_data.get("children", []):
984
- node.children.append(self._reconstruct_node(child_data, parent=node))
985
- node.sub_node_count = node_data.get("sub_node_count", 0)
958
+ # TopicNode
959
+ node = TopicNode(
960
+ topic_name = node_data.get("topic_name", ""),
961
+ summary = node_data.get("summary", ""),
962
+ start_index = node_data.get("start_index", 0),
963
+ end_index = node_data.get("end_index", 0),
964
+ parent = parent,
965
+ )
966
+ if "node_id" in node_data:
967
+ node.node_id = node_data["node_id"]
968
+ node.created_at = node_data.get("created_at", node.created_at)
969
+ for child_data in node_data.get("children", []):
970
+ node.children.append(self._reconstruct_node(child_data, parent=node))
986
971
  return node
987
972
 
988
973
  def _find_current_node(self, node):
@@ -1116,10 +1101,7 @@ class CTree:
1116
1101
  )
1117
1102
  print(f"Merged {stats['merged']} branches in {stats['duration_secs']:.1f}s")
1118
1103
  """
1119
- import time as _time
1120
- from .vector_db import cosine_similarity as _cos
1121
-
1122
- t0 = _time.time()
1104
+ t0 = time.time()
1123
1105
  merged = pruned = skipped = 0
1124
1106
 
1125
1107
  # Phase 1 — collect frozen candidates
@@ -1165,7 +1147,7 @@ class CTree:
1165
1147
  if na.node_id not in embeddings or nb.node_id not in embeddings:
1166
1148
  skipped += 1
1167
1149
  continue
1168
- sim = _cos(embeddings[na.node_id], embeddings[nb.node_id])
1150
+ sim = cosine_similarity(embeddings[na.node_id], embeddings[nb.node_id])
1169
1151
  if sim < similarity_threshold:
1170
1152
  continue
1171
1153
  older, newer = (na, nb) if na.created_at <= nb.created_at else (nb, na)
@@ -1215,5 +1197,5 @@ class CTree:
1215
1197
  "merged": merged,
1216
1198
  "pruned": pruned,
1217
1199
  "skipped": skipped,
1218
- "duration_secs": _time.time() - t0,
1200
+ "duration_secs": time.time() - t0,
1219
1201
  }
@@ -92,7 +92,7 @@ class PromptSynthesizer:
92
92
  _walk(self.tree.root)
93
93
  return result
94
94
 
95
- def _get_surgical_context(self, query_vector) -> str:
95
+ def _get_surgical_context(self, query_vector, top_k_docs: int = 3) -> str:
96
96
  """
97
97
  Multi-path surgical retrieval (the heart of TRACE):
98
98
 
@@ -138,7 +138,7 @@ class PromptSynthesizer:
138
138
  unique_nodes_ordered.append((node, similarity))
139
139
 
140
140
  # Now rank paths by similarity and take the top-3
141
- all_ancestor_sets_ranked = sorted(all_ancestor_sets, key=lambda x: x[0], reverse=True)[:3]
141
+ all_ancestor_sets_ranked = sorted(all_ancestor_sets, key=lambda x: x[0], reverse=True)[:top_k_docs]
142
142
 
143
143
  blocks = ["── SURGICAL MEMORY CONTEXT (top matched topics + ancestry) ──"]
144
144
  for i, (similarity, ancestors) in enumerate(all_ancestor_sets_ranked, 1):
@@ -159,8 +159,9 @@ class PromptSynthesizer:
159
159
  query_vector: list,
160
160
  active_node,
161
161
  recent_messages: list,
162
- top_k_history: int = 2,
163
- min_history_similarity: float = 0.50,
162
+ top_k_docs: int = 3, # max topic branches to surface in surgical context
163
+ top_k_history: int = 5,
164
+ min_history_similarity: float = 0.25,
164
165
  ) -> str:
165
166
  """
166
167
  Assemble the full enriched system prompt for the next LLM turn.
@@ -171,6 +172,7 @@ class PromptSynthesizer:
171
172
  query_vector : Embedding of *user_query*.
172
173
  active_node : ``tree.current_node``.
173
174
  recent_messages : Tail of ``tree.conversation`` (e.g. last 6 msgs).
175
+ top_k_docs : Max number of topic-branch paths to surface in surgical context.
174
176
  top_k_history : Max past conversation messages to recall.
175
177
  min_history_similarity : Minimum cosine score for conversation recall.
176
178
 
@@ -196,7 +198,7 @@ class PromptSynthesizer:
196
198
  )
197
199
  """
198
200
  # 1. Surgical multi-path memory context
199
- narrative_block = self._get_surgical_context(query_vector)
201
+ narrative_block = self._get_surgical_context(query_vector, top_k_docs)
200
202
 
201
203
  # 2. Cross-thread conversation recall
202
204
  history_matches = self.vdb.search_conversation(
@@ -29,11 +29,7 @@ def unpack_float_vector(blob):
29
29
  return []
30
30
  num_floats = len(blob) // 4
31
31
  format_string = '<' + str(num_floats) + 'f'
32
- unpacked = struct.unpack(format_string, blob)
33
- result = []
34
- for val in unpacked:
35
- result.append(val)
36
- return result
32
+ return list(struct.unpack(format_string, blob))
37
33
 
38
34
  def cosine_similarity(v1, v2):
39
35
  if not v1 or not v2 or len(v1) != len(v2):
@@ -94,7 +90,7 @@ class VectorDatabase:
94
90
 
95
91
 
96
92
 
97
- def search_conversation(self, query_vector, top_k=2, min_similarity=0.5):
93
+ def search_conversation(self, query_vector, top_k=5, min_similarity=0.25):
98
94
  if not query_vector:
99
95
  return []
100
96
  conn = sqlite3.connect(self.db_path)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: trace-memory
3
- Version: 1.0.4
3
+ Version: 1.0.6
4
4
  Summary: TRACE — Temporal Retrieval And Context Engine: self-healing hierarchical memory for LLM agents.
5
5
  License: Apache-2.0
6
6
  Project-URL: Homepage, https://github.com/husain34/TRACE
@@ -25,6 +25,7 @@ License-File: NOTICE
25
25
  Requires-Dist: openai>=1.0.0
26
26
  Requires-Dist: python-dotenv>=1.0.0
27
27
  Requires-Dist: sentence-transformers>=2.2.0
28
+ Requires-Dist: numpy>=1.21.0
28
29
  Dynamic: license-file
29
30
 
30
31
  <div align="center">
@@ -208,14 +209,13 @@ Result: The AI aggressively stopped the user.
208
209
  ```
209
210
  Phase 1 — Collect all frozen (inactive) TopicNodes
210
211
  Phase 2 — Generate missing summaries; embed all candidates
211
- Phase 3 — Compute pairwise cosine similarity
212
- Phase 4 — For each pair above threshold, apply 4 axioms:
212
+ Phase 3 — Compute pairwise cosine similarity; for each pair above threshold, apply 4 axioms:
213
213
  Axiom 1 — Chronological Guard
214
214
  Axiom 2 — Frozen State check
215
215
  Axiom 3 — Similarity threshold (default 0.55)
216
216
  Axiom 4 — LLM Veto
217
217
  If all 4 pass → merge (newer becomes child of older)
218
- Phase 5 — Optional: prune trivial leaf messages
218
+ Phase 4 — Optional: prune trivial leaf messages
219
219
  ```
220
220
 
221
221
  #### The Four Axioms in Detail
@@ -403,16 +403,13 @@ tree.save("sessions/chat_001.json", save_conversation=True)
403
403
 
404
404
  `save_conversation=True` embeds the raw message list so the session can be fully restored later.
405
405
 
406
- ---
407
-
408
- ##### `CTree.load(filepath: str, api_key: str = None, base_url: str = None, model: str = None) -> CTree`
406
+ ##### `CTree.load(filepath: str, api_key: str = None, base_url: str = None, model: str = None, embed_fn: Callable = None) -> CTree`
409
407
 
410
408
  Restore a tree from a JSON file.
411
409
 
412
410
  ```python
413
411
  tree = CTree.load("sessions/chat_001.json", api_key="sk-...", embed_fn=embed)
414
- tree.vdb = VectorDatabase("sessions/chat_001.db")
415
-
412
+ tree.vdb = VectorDatabase("sessions/chat_001.db")
416
413
  ```
417
414
 
418
415
  ---
@@ -473,7 +470,7 @@ ROOT (sub-nodes: 3)
473
470
 
474
471
  A local SQLite vector store with two active tables: conversation vectors and topic summaries.
475
472
 
476
- > **Note on Scaling:** TRACE uses pure Python (`sqlite3` and `struct`) to avoid heavy dependencies and run instantly on any machine. This is perfectly sufficient for local agents and moderate conversation histories. However, for massive scale (millions of vectors), you will need to swap this module for a dedicated vector database like FAISS or Chroma.
473
+ > **Note on Scaling:** TRACE uses SQLite (`sqlite3`) for storage, `numpy` for fast cosine similarity, and `struct` for compact binary packing all chosen to keep the footprint small while running on any machine. For massive scale (millions of vectors), swap this module for FAISS or Chroma.
477
474
 
478
475
  ```python
479
476
  from trace_memory import VectorDatabase
@@ -487,8 +484,6 @@ vdb = VectorDatabase("path/to/session.db") # creates the DB if it doesn't exist
487
484
 
488
485
  ---
489
486
 
490
- ---
491
-
492
487
  #### Conversation methods
493
488
 
494
489
  ##### `vdb.add_conversation_message(msg: ConversationVector)`
@@ -573,6 +568,7 @@ system_prompt = synth.synthesize_prompt(
573
568
  query_vector = embed("Is the cake safe for Sarah?"),
574
569
  active_node = tree.current_node,
575
570
  recent_messages = tree.conversation[-6:],
571
+ top_k_docs = 3, # max topic branch paths to surface
576
572
  top_k_history = 2, # max past messages to recall
577
573
  min_history_similarity = 0.50, # min cosine score for conversation recall
578
574
  )
@@ -641,15 +637,25 @@ tree = CTree(api_key="sk-...", model="gpt-4o-mini")
641
637
 
642
638
  ### 7.3 With NVIDIA NIM / any OpenAI-compatible endpoint
643
639
 
644
- ```python
645
- import os
646
- os.environ["OPENAI_BASE_URL"] = "https://integrate.api.nvidia.com/v1"
647
- os.environ["OPENAI_API_KEY"] = "nvapi-..."
640
+ Point the `base_url` to your endpoint.
648
641
 
649
- from trace_memory import CTree
650
- tree = CTree(model="meta/llama-3.1-70b-instruct")
642
+ ```python
643
+ tree = CTree(
644
+ api_key="your-nvidia-api-key",
645
+ base_url="https://integrate.api.nvidia.com/v1",
646
+ model="meta/llama-3.1-8b-instruct"
647
+ )
651
648
  ```
652
649
 
650
+ ### 7.4 Best Practices for LLM Selection
651
+
652
+ **Do not use "Reasoning" models (DeepSeek R1, Claude 3.7 Sonnet thinking mode, o1) for the internal `CTree` engine.**
653
+
654
+ When `CTree` runs in the background, it performs simple, deterministic classification tasks (e.g., naming a topic in 3 words, or extracting a JSON boolean). Reasoning models are built for complex logic and will waste massive amounts of compute "thinking" about simple tasks. Worse, reasoning models often exceed TRACE's strict internal token limits (`max_tokens=50` or `200`) designed to keep background operations fast, causing the LLM to get cut off mid-thought and breaking the internal JSON parsers.
655
+
656
+ * **For the `CTree` internal model:** Always use a fast, standard model (e.g., `gpt-4o-mini`, `meta-llama-3.1-8b-instruct`).
657
+ * **For your actual chat application:** You can (and should!) use whatever massive reasoning model you prefer to generate the final response to the user. TRACE is completely decoupled from your front-end chat completion.
658
+
653
659
  ---
654
660
 
655
661
  ## 8. Edge-Case Stress Tests & Validation
@@ -684,6 +690,8 @@ TRACE loads `.env` automatically if `python-dotenv` is installed.
684
690
  |---|---|---|
685
691
  | `openai` | ≥ 1.0.0 | ✅ Yes |
686
692
  | `python-dotenv` | ≥ 1.0.0 | ✅ Yes |
693
+ | `sentence-transformers` | ≥ 2.2.0 | ✅ Yes (local embed fallback) |
694
+ | `numpy` | ≥ 1.21.0 | ✅ Yes (cosine similarity) |
687
695
  | `sqlite3` | built-in | ✅ Yes (no install needed) |
688
696
  | `struct` | built-in | ✅ Yes (no install needed) |
689
697
 
@@ -1,3 +1,4 @@
1
1
  openai>=1.0.0
2
2
  python-dotenv>=1.0.0
3
3
  sentence-transformers>=2.2.0
4
+ numpy>=1.21.0
File without changes
File without changes
File without changes