| |
| |
| @@ -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 = ( |
| |
| |
| @@ -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. |
| |
| |
| |
| @@ -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}.") |
| |
| |
| |
| @@ -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]): |
| """ |
| |
| |
| @@ -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) |
|
|
