99 ANSWER_LABEL ,
1010 DEFAULT_MAX_QUERY_KEY_WORDS ,
1111 QUERY_LABEL ,
12+ WORKING_MEMORY_TYPE ,
1213 MemCubeID ,
1314 UserID ,
1415)
@@ -35,11 +36,12 @@ def __init__(self, config: GeneralSchedulerConfig):
3536
3637 # for evaluation
3738 def search_for_eval (
38- self ,
39- query : str ,
40- user_id : UserID | str ,
41- top_k : int ,
42- ) -> list [str ]:
39+ self , query : str , user_id : UserID | str , top_k : int , scheduler_flag : bool = True
40+ ) -> (list [str ], bool ):
41+ self .monitor .register_query_monitor_if_not_exists (
42+ user_id = user_id , mem_cube_id = self .current_mem_cube_id
43+ )
44+
4345 query_keywords = self .monitor .extract_query_keywords (query = query )
4446 logger .info (f'Extract keywords "{ query_keywords } { query } )
4547
@@ -48,35 +50,61 @@ def search_for_eval(
4850 keywords = query_keywords ,
4951 max_keywords = DEFAULT_MAX_QUERY_KEY_WORDS ,
5052 )
51- self .monitor .query_monitors .put (item = item )
52- logger .debug (
53- f"Queries in monitor are { self .monitor .query_monitors .get_queries_with_timesort ()}
54- )
53+ query_monitor = self .monitor .query_monitors [user_id ][self .current_mem_cube_id ]
54+ query_monitor .put (item = item )
55+ logger .debug (f"Queries in monitor are { query_monitor .get_queries_with_timesort ()} )
5556
5657 queries = [query ]
5758
5859 # recall
59- cur_working_memory , new_candidates = self .process_session_turn (
60- queries = queries ,
61- user_id = user_id ,
62- mem_cube_id = self .current_mem_cube_id ,
63- mem_cube = self .current_mem_cube ,
64- top_k = self .top_k ,
65- )
66- logger .info (f"Processed { queries } { len (new_candidates )} )
67-
68- # rerank
69- new_order_working_memory = self .replace_working_memory (
70- user_id = user_id ,
71- mem_cube_id = self .current_mem_cube_id ,
72- mem_cube = self .current_mem_cube ,
73- original_memory = cur_working_memory ,
74- new_memory = new_candidates ,
60+ mem_cube = self .current_mem_cube
61+ text_mem_base = mem_cube .text_mem
62+
63+ cur_working_memory : list [TextualMemoryItem ] = text_mem_base .get_working_memory ()
64+ text_working_memory : list [str ] = [w_m .memory for w_m in cur_working_memory ]
65+ intent_result = self .monitor .detect_intent (
66+ q_list = queries , text_working_memory = text_working_memory
7567 )
76- new_order_working_memory = new_order_working_memory [:top_k ]
77- logger .info (f"size of new_order_working_memory: { len (new_order_working_memory )} )
7868
79- return [m .memory for m in new_order_working_memory ]
69+ if not scheduler_flag :
70+ return text_working_memory , intent_result ["trigger_retrieval" ]
71+ else :
72+ if intent_result ["trigger_retrieval" ]:
73+ missing_evidences = intent_result ["missing_evidences" ]
74+ num_evidence = len (missing_evidences )
75+ k_per_evidence = max (1 , top_k // max (1 , num_evidence ))
76+ new_candidates = []
77+ for item in missing_evidences :
78+ logger .info (f"missing_evidences: { item } )
79+ results : list [TextualMemoryItem ] = self .retriever .search (
80+ query = item ,
81+ mem_cube = mem_cube ,
82+ top_k = k_per_evidence ,
83+ method = self .search_method ,
84+ )
85+ logger .info (
86+ f"search results for { missing_evidences } { [one .memory for one in results ]}
87+ )
88+ new_candidates .extend (results )
89+ print (
90+ f"missing_evidences: { missing_evidences } { len (new_candidates )}
91+ )
92+ else :
93+ new_candidates = []
94+ print (f"intent_result: { intent_result } )
95+
96+ # rerank
97+ new_order_working_memory = self .replace_working_memory (
98+ user_id = user_id ,
99+ mem_cube_id = self .current_mem_cube_id ,
100+ mem_cube = self .current_mem_cube ,
101+ original_memory = cur_working_memory ,
102+ new_memory = new_candidates ,
103+ )
104+ new_order_working_memory = new_order_working_memory [:top_k ]
105+ logger .info (f"size of new_order_working_memory: { len (new_order_working_memory )} )
106+
107+ return [m .memory for m in new_order_working_memory ], intent_result ["trigger_retrieval" ]
80108
81109 def _query_message_consumer (self , messages : list [ScheduleMessageItem ]) -> None :
82110 """
@@ -105,6 +133,10 @@ def _query_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
105133
106134 # update query monitors
107135 for msg in messages :
136+ self .monitor .register_query_monitor_if_not_exists (
137+ user_id = user_id , mem_cube_id = mem_cube_id
138+ )
139+
108140 query = msg .content
109141 query_keywords = self .monitor .extract_query_keywords (query = query )
110142 logger .info (f'Extract keywords "{ query_keywords } { query } )
@@ -114,9 +146,11 @@ def _query_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
114146 keywords = query_keywords ,
115147 max_keywords = DEFAULT_MAX_QUERY_KEY_WORDS ,
116148 )
117- self .monitor .query_monitors .put (item = item )
149+
150+ self .monitor .query_monitors [user_id ][mem_cube_id ].put (item = item )
118151 logger .debug (
119- f"Queries in monitor are { self .monitor .query_monitors .get_queries_with_timesort ()}
152+ f"Queries in monitor are "
153+ f"{ self .monitor .query_monitors [user_id ][mem_cube_id ].get_queries_with_timesort ()}
120154 )
121155
122156 queries = [msg .content for msg in messages ]
@@ -215,6 +249,9 @@ def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
215249 mem_type = mem_item .metadata .memory_type
216250 mem_content = mem_item .memory
217251
252+ if mem_type == WORKING_MEMORY_TYPE :
253+ continue
254+
218255 self .log_adding_memory (
219256 memory = mem_content ,
220257 memory_type = mem_type ,
@@ -289,18 +326,20 @@ def process_session_turn(
289326 new_candidates = []
290327 for item in missing_evidences :
291328 logger .info (f"missing_evidences: { item } )
329+ info = {
330+ "user_id" : user_id ,
331+ "session_id" : "" ,
332+ }
333+
292334 results : list [TextualMemoryItem ] = self .retriever .search (
293- query = item , mem_cube = mem_cube , top_k = k_per_evidence , method = self .search_method
335+ query = item ,
336+ mem_cube = mem_cube ,
337+ top_k = k_per_evidence ,
338+ method = self .search_method ,
339+ info = info ,
294340 )
295341 logger .info (
296342 f"search results for { missing_evidences } { [one .memory for one in results ]}
297343 )
298344 new_candidates .extend (results )
299-
300- if len (new_candidates ) == 0 :
301- logger .warning (
302- f"As new_candidates is empty, new_candidates is set same to working_memory.\n "
303- f"time_trigger_flag: { time_trigger_flag } { intent_result }
304- )
305- new_candidates = cur_working_memory
306345 return cur_working_memory , new_candidates
0 commit comments