--- 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)