File size: 9,741 Bytes
750b88d | 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | --- a/REMem/src/remem/agent/answer_generator.py
+++ b/REMem/src/remem/agent/answer_generator.py
@@ -1,3 +1,4 @@
+# Modified.
from collections import defaultdict
from typing import Any, Callable, Dict, List, Optional
@@ -270,6 +271,7 @@
remem_instance: ReMem instance (for verbatim aggregation mode).
skip_empty_step: If True, omit steps that have no retrieved passages/content.
"""
+ max_items = remem_instance.global_config.qa_top_k if remem_instance else 10
qa_context = ""
# Add preliminary answer if available, before agent interaction history
@@ -303,17 +305,20 @@
# Fallback to original format if verbatim aggregation fails
if self.logger:
self.logger.warning("Verbatim aggregation failed, falling back to retrieved contexts display")
- qa_context += self._build_retrieved_contexts(previous_steps, skip_empty_step=skip_empty_step)
+ qa_context += self._build_retrieved_contexts(previous_steps, skip_empty_step=skip_empty_step,
+ max_items=max_items)
else:
# Use retrieved contexts display
- qa_context += self._build_retrieved_contexts(previous_steps, skip_empty_step=skip_empty_step)
+ qa_context += self._build_retrieved_contexts(previous_steps, skip_empty_step=skip_empty_step,
+ max_items=max_items)
if visited_nodes and not self.only_use_verbatim_contexts:
qa_context += f"Total Nodes Explored: {len(visited_nodes)}\n"
return qa_context
- def _build_retrieved_contexts(self, previous_steps: List[Dict[str, Any]], skip_empty_step: bool = False) -> str:
+ def _build_retrieved_contexts(self, previous_steps: List[Dict[str, Any]], skip_empty_step: bool = False,
+ max_items: int = 10) -> str:
"""Build retrieved contexts display for all content types (gists, facts, verbatim, etc.).
Args:
@@ -370,7 +375,7 @@
if num_passages > 0 and passages:
# Group contents by type using the same pattern as tool_selector
content_groups = self._group_contents_by_type(passages)
- context += self._format_grouped_contents(content_groups)
+ context += self._format_grouped_contents(content_groups, max_items)
context += "\n"
else:
@@ -405,7 +410,7 @@
return grouped_contents
- def _format_grouped_contents(self, content_groups: Dict[str, List[str]]) -> str:
+ def _format_grouped_contents(self, content_groups: Dict[str, List[str]], max_items: int = 10) -> str:
"""Format grouped contents for display.
Args:
@@ -420,7 +425,6 @@
if contents:
formatted_output += f" - Truncated top {content_type}:\n"
max_content_length = 1500
- max_items = 10
for i, content in enumerate(contents[:max_items], 1):
display_content = (
--- a/REMem/src/remem/rag_strategies/base_strategy.py
+++ b/REMem/src/remem/rag_strategies/base_strategy.py
@@ -1,3 +1,4 @@
+# Modified.
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple, Union
@@ -14,7 +15,7 @@
self.remem = remem_instance
@abstractmethod
- def index(self, docs: List[str]) -> None:
+ def index(self, docs: List[str], doc_visits: List = None) -> None:
"""
Index documents using the specific strategy.
--- a/REMem/src/remem/rag_strategies/episodic_gist_strategy.py
+++ b/REMem/src/remem/rag_strategies/episodic_gist_strategy.py
@@ -1,3 +1,4 @@
+# Modified.
import json
import logging
import os
@@ -34,7 +35,7 @@
datetime_cache_path = os.path.join(remem_instance.working_dir, "datetime_cache.pkl")
set_datetime_cache_path(datetime_cache_path)
- def index(self, docs: List) -> None:
+ def index(self, docs: List, doc_visits: List = None) -> None:
"""Internal implementation of index_with_gist."""
logger.info("Indexing Documents")
logger.info("Performing Extraction")
@@ -91,6 +92,7 @@
chunk_id_to_verbatim_id = {}
entries = ["verbatim", "facts", "gists"]
+ _gist_src = []
for chunk_idx, chunk in enumerate(episode_results_dict.values()):
chunk_id = chunk["chunk_id"]
@@ -101,7 +103,9 @@
for v in docs[chunk_idx]:
split_results.append([make_chunk_content("message", v), v])
elif isinstance(docs[chunk_idx], str):
- split_results.append([docs[chunk_idx], {"content": docs[chunk_idx]}])
+ split_results.append([docs[chunk_idx],
+ {"content": docs[chunk_idx],
+ "doc_index": chunk_idx}])
else:
raise ValueError(
f"Expected docs[chunk_idx] to be a list when splitting verbatim, but got type: {type(docs[chunk_idx])}"
@@ -119,6 +123,7 @@
chunk_id_to_verbatim_id[chunk_id] = verbatim_keys_for_chunk
elif entry in ["verbatim", "gists"]:
+ _n0 = len(element_to_encode[entry])
if isinstance(chunk[entry], str):
element_to_encode[entry].append(chunk[entry])
elif isinstance(chunk[entry], list):
@@ -129,6 +134,10 @@
else:
# Default behavior: each gist becomes a separate node
element_to_encode[entry].extend(chunk[entry])
+ if entry == "gists":
+ _lbl = (doc_visits[chunk_idx]
+ if doc_visits and chunk_idx < len(doc_visits) else None)
+ _gist_src.extend([_lbl] * (len(element_to_encode[entry]) - _n0))
if entry == "verbatim":
if isinstance(docs[chunk_idx], str):
if docs[chunk_idx].startswith("{") and docs[chunk_idx].endswith("}"):
@@ -140,7 +149,7 @@
metadata["content"] = content # session_metadata
else:
content = docs[chunk_idx]
- metadata = {"content": content}
+ metadata = {"content": content, "doc_index": chunk_idx}
else:
chunk_id = chunk.get("chunk_id") if isinstance(chunk, dict) else chunk.chunk_id
metadata = chunk_dict[chunk_id]["metadata"]
@@ -217,6 +226,11 @@
self.remem.episodic_embedding_stores[entry].insert_chunk_dicts(
verbatim_metadata, extract_method="message", embed=True
)
+ elif entry == "gists":
+ self.remem.episodic_embedding_stores[entry].insert_chunk_dicts(
+ [{"content": g, "source_visit": (_gist_src[i] if i < len(_gist_src) else None)}
+ for i, g in enumerate(element_to_encode[entry])],
+ extract_method="message", embed=True)
else:
embed = False if entry == "entity" else True
self.remem.episodic_embedding_stores[entry].insert_strings(element_to_encode[entry], embed=embed)
@@ -913,7 +927,8 @@
# Extract content from the batch result
top_k_chunks_content = [chunk_rows[hash_id]["content"] for hash_id in hash_ids_to_fetch]
- top_k_chunks_metadata = [] # Empty list instead of None to fix type error
+ top_k_chunks_metadata = [chunk_rows[hash_id].get("metadata", {}) or {}
+ for hash_id in hash_ids_to_fetch]
else:
raise ValueError(f"Invalid return_chunk type: {return_chunk}.")
--- a/REMem/src/remem/remem.py
+++ b/REMem/src/remem/remem.py
@@ -1,3 +1,4 @@
+# Modified.
import json
import logging
import os
@@ -327,7 +328,7 @@
)
return loaded_graph
- def index(self, docs: List[str]) -> None:
+ def index(self, docs: List[str], doc_visits: List = None) -> None:
"""
Indexes the given documents using the configured RAG strategy.
@@ -335,7 +336,7 @@
docs : List[str]
A list of documents to be indexed.
"""
- self.rag_strategy.index(docs)
+ self.rag_strategy.index(docs, doc_visits=doc_visits)
def index_original(self, docs: List[str]):
"""
--- a/REMem/src/remem/rerank.py
+++ b/REMem/src/remem/rerank.py
@@ -1,3 +1,4 @@
+# Modified.
import ast
import difflib
import json
@@ -95,7 +96,7 @@
)
# call openai
- self.default_gen_kwargs["max_completion_tokens"] = 512
+ self.default_gen_kwargs["max_completion_tokens"] = 10000
try:
response = self.llm_infer_fn(messages=messages, model=self.model_name, **self.default_gen_kwargs)
@@ -241,7 +242,7 @@
}
)
- self.default_gen_kwargs["max_completion_tokens"] = 512
+ self.default_gen_kwargs["max_completion_tokens"] = 10000
try:
response = self.llm_infer_fn(messages=messages, model=self.model_name, **self.default_gen_kwargs)
|