@ -1,5 +1,3 @@
import asyncio
from deepsearcher . agent . base import BaseAgent , describe_class from deepsearcher . agent . base import BaseAgent , describe_class
from deepsearcher . embedding . base import BaseEmbedding from deepsearcher . embedding . base import BaseEmbedding
from deepsearcher . llm . base import BaseLLM from deepsearcher . llm . base import BaseLLM
@ -7,20 +5,23 @@ from deepsearcher.utils import log
from deepsearcher . vector_db import RetrievalResult from deepsearcher . vector_db import RetrievalResult
from deepsearcher . vector_db . base import BaseVectorDB , deduplicate from deepsearcher . vector_db . base import BaseVectorDB , deduplicate
COLLECTION_ROUTE_PROMPT = """ COLLECTION_ROUTE_PROMPT = """
I provide you with collection_name ( s ) and corresponding collection_description ( s ) . I provide you with collection_name ( s ) and corresponding collection_description ( s ) .
Please select the collection names that may be related to the question and return a python list of str . Please select the collection names that may be related to the query and return a python list of str .
If there is no collection related to the question , you can return an empty list . If there is no collection related to the query , you can return an empty list .
" QUESTION " : { question } " Query " : { query }
" COLLECTION_INFO " : { collection_info } " COLLECTION_INFO " : { collection_info }
When you return , you can ONLY return a json convertable python list of str , WITHOUT any other additional content . When you return , you can ONLY return a json convertable python list of str , WITHOUT any other additional content .
Your selected collection name list is : Your selected collection name list is :
""" """
SUB_QUERY_PROMPT = """ SUB_QUERY_PROMPT = """
To answer this question more comprehensively , please break down the original question into few numbers of sub - questions ( more if necessary ) . To answer this question more comprehensively , please break down the original question into few numbers of sub - questions
( the less the better , but more if nesscessary to ensure answering the original question ) .
If this is a very simple question and no decomposition is necessary , then keep the only one original question . If this is a very simple question and no decomposition is necessary , then keep the only one original question .
Make sure each sub - question is clear , concise and atomic . Make sure each sub - question is clear , concise and atomic .
Return as list of str in python style and json convertable . Return as list of str in python style and json convertable .
@ -43,16 +44,18 @@ Example output:
Provide your response in a python code list of str format : Provide your response in a python code list of str format :
""" """
RERANK_PROMPT = """ RERANK_PROMPT = """
Based on the query questions and the retrieved chunks , determine whether each chunk is helpful in answering any of the query questions . Based on the query and the retrieved chunks , give a quick judge of whether each chunk is helpful in answering the query .
For each chunk , you must return " YES " or " NO " without any other information . For each chunk , you must return " YES " or " NO " python style list without any other information .
Query Questions : { query } Query : { query }
Retrieved Chunks : Retrieved Chunks :
{ retrieved_chunks } { retrieved_chunks }
Respond with a list of " YES " or " NO " values , one for each chunk , in the same order as the chunks are listed . For example a list of chunks of three : [ " YES " , " NO " , " YES " ] Respond with a list of " YES " or " NO " values , one for each chunk , in the same order as the chunks are listed .
For example , if there is a list of four chunks , the answer could be : [ " YES " , " NO " , " YES " , " YES " ]
""" """
@ -60,16 +63,16 @@ REFLECT_PROMPT = """
Determine whether additional search queries are needed based on the original query , previous sub queries , and all retrieved document chunks . Determine whether additional search queries are needed based on the original query , previous sub queries , and all retrieved document chunks .
If returned chunks does not cover all previous sub - queries , this means that there are no related documents can be retrieved . If returned chunks does not cover all previous sub - queries , this means that there are no related documents can be retrieved .
In this case , try generate simliar but slightly different queries to the previous sub - queries . In this case , try generate simliar but slightly different queries to the previous sub - queries .
And if further research is needed based on the new information , provide a Python list of more queries . And if further research is needed based on the new information which those chunks provided , give more queries on the basis of them .
( which is prefered , even if the previous sub - queries can be well answered by retrieved chunks , but ultimately according to your judge ) ( which is prefered , but ultimately according to your judge )
If no further research is needed , return an empty list . If no further research is needed , return an empty list .
Original Query : { question } Original Query : { original_query }
Previous Sub Queries : { mini_question s} Previous Sub Queries : { all_sub_querie s}
Related Chunks : Related Chunks :
{ mini_ chunk_ str } { chunks }
Respond exclusively in valid List of str format without any other text . """ Respond exclusively in valid List of str format without any other text . """
@ -79,14 +82,15 @@ You are a AI content analysis expert.
Please generate a long , specific and detailed answer or report based on the previous queries and the retrieved document chunks . Please generate a long , specific and detailed answer or report based on the previous queries and the retrieved document chunks .
If the chunks are not enough to answer the query or additional information is needed to enhance the content , you should answer with your own knowledge . If the chunks are not enough to answer the query or additional information is needed to enhance the content , you should answer with your own knowledge .
In this case , mark the part ( s ) that generated by your own with < unref > your knowledge here < / unref > In this case , mark the part ( s ) that generated by your own with < unref > your knowledge here < / unref >
( Don ' t place <unref></unref> part(s) individually into one paragraph, but insert it the proper place of the context) ( Don ' t place <unref></unref> part(s) individually into one paragraph, but insert it the proper place of the report)
Plus , you should give references in the report where you quote from the chunks using markdown links , and give a list of references at the end of the report .
Original Query : { question } Original Query : { original_query }
Previous Sub Queries : { mini_question s} Previous Sub Queries : { all_sub_querie s}
Related Chunks : Related Chunks :
{ mini_ chunk_ str } { chunks }
""" """
@ -108,7 +112,7 @@ class DeepSearch(BaseAgent):
embedding_model : BaseEmbedding , embedding_model : BaseEmbedding ,
vector_db : BaseVectorDB , vector_db : BaseVectorDB ,
max_iter : int = 3 , max_iter : int = 3 ,
route_collection : bool = Tru e, route_collection : bool = Fals e,
text_window_splitter : bool = True , text_window_splitter : bool = True ,
* * kwargs , * * kwargs ,
) : ) :
@ -162,7 +166,7 @@ class DeepSearch(BaseAgent):
) )
return [ the_only_collection ] return [ the_only_collection ]
vector_db_search_prompt = COLLECTION_ROUTE_PROMPT . format ( vector_db_search_prompt = COLLECTION_ROUTE_PROMPT . format (
question = query , query = query ,
collection_info = [ collection_info = [
{ {
" collection_name " : collection_info . collection_name , " collection_name " : collection_info . collection_name ,
@ -198,7 +202,7 @@ class DeepSearch(BaseAgent):
content = self . llm . remove_think ( content ) content = self . llm . remove_think ( content )
return self . llm . literal_eval ( content ) return self . llm . literal_eval ( content )
async def _search_chunks_from_vectordb ( self , query : str ) : def _search_chunks_from_vectordb ( self , query : str ) :
if self . route_collection : if self . route_collection :
selected_collections = self . invoke ( selected_collections = self . invoke (
query = query , dim = self . embedding_model . dimension query = query , dim = self . embedding_model . dimension
@ -222,7 +226,10 @@ class DeepSearch(BaseAgent):
# Format all chunks for batch processing # Format all chunks for batch processing
formatted_chunks = " " formatted_chunks = " "
for i , retrieved_result in enumerate ( retrieved_results ) : for i , retrieved_result in enumerate ( retrieved_results ) :
formatted_chunks + = f " <chunk_ { i } > \n { retrieved_result . text } \n </chunk_ { i } > \n " formatted_chunks + = f '''
< chunk_ { i + 1 } > \n { retrieved_result . text } \n < / chunk_ { i + 1 } > \n
< reference_ { i + 1 } > \n { retrieved_result . reference } \n < / reference_ { i + 1 } >
'''
# Batch process all chunks with a single LLM call # Batch process all chunks with a single LLM call
content = self . llm . chat ( content = self . llm . chat (
@ -278,21 +285,27 @@ class DeepSearch(BaseAgent):
) )
return all_retrieved_results return all_retrieved_results
def _generate_gap _queries ( def _generate_more_sub _queries (
self , original_query : str , all_sub_queries : list [ str ] , all_chunk s : list [ RetrievalResult ] self , original_query : str , all_sub_queries : list [ str ] , all_retrieved_result s : list [ RetrievalResult ]
) - > list [ str ] : ) - > list [ str ] :
chunks = [ ]
for i , chunk in enumerate ( all_retrieved_results ) :
if self . text_window_splitter and " wider_text " in chunk . metadata :
chunks . append ( chunk . metadata [ " wider_text " ] )
else :
chunks . append ( f ''' <chunk { i + 1 } > { chunk . text } </chunk { i + 1 } ><reference { i + 1 } > { chunk . reference } </reference { i + 1 } > ''' )
reflect_prompt = REFLECT_PROMPT . format ( reflect_prompt = REFLECT_PROMPT . format (
question = original_query , original_query = original_query ,
mini_questions = all_sub_queries , all_sub_querie s= all_sub_queries ,
mini_chunk_str = self . _format_chunk_texts ( [ chunk . text for chunk in all_chunks ] ) chunks = " \n " . join ( chunks )
if len ( all_chunks ) > 0 if len ( all_retrieved_result s ) > 0
else " NO RELATED CHUNKS FOUND. " , else " NO RELATED CHUNKS FOUND. " ,
) )
response = self . llm . chat ( [ { " role " : " user " , " content " : reflect_prompt } ] ) response = self . llm . chat ( [ { " role " : " user " , " content " : reflect_prompt } ] )
response = self . llm . remove_think ( response ) response = self . llm . remove_think ( response )
return self . llm . literal_eval ( response ) return self . llm . literal_eval ( response )
def retrieve ( self , original_query : str , * * kwargs ) - > tuple [ list [ RetrievalResult ] , dict ] : def retrieve ( self , original_query : str , * * kwargs ) - > tuple [ list [ RetrievalResult ] , list [ str ] ] :
""" """
Retrieve relevant documents from the knowledge base for the given query . Retrieve relevant documents from the knowledge base for the given query .
@ -308,15 +321,10 @@ class DeepSearch(BaseAgent):
- A list of retrieved document results - A list of retrieved document results
- Additional information about the retrieval process - Additional information about the retrieval process
""" """
return asyncio . run ( self . async_retrieve ( original_query , * * kwargs ) )
async def async_retrieve (
self , original_query : str , * * kwargs
) - > tuple [ list [ RetrievalResult ] , dict ] :
max_iter = kwargs . pop ( " max_iter " , self . max_iter ) max_iter = kwargs . pop ( " max_iter " , self . max_iter )
### SUB QUERIES ### ### SUB QUERIES ###
log . color_print ( f " <query> { original_query } </query> \n " ) log . color_print ( f " <query> { original_query } </query> \n " )
all_search_res = [ ] all_search_results = [ ]
all_sub_queries = [ ] all_sub_queries = [ ]
sub_queries = self . _generate_sub_queries ( original_query ) sub_queries = self . _generate_sub_queries ( original_query )
@ -324,54 +332,46 @@ class DeepSearch(BaseAgent):
log . color_print ( " No sub queries were generated by the LLM. Exiting. " ) log . color_print ( " No sub queries were generated by the LLM. Exiting. " )
return [ ] , { } return [ ] , { }
else : else :
log . color_print ( log . color_print ( f " </think> Break down the original query into new sub queries: { sub_queries } " )
f " <think> Break down the original query into new sub queries: { sub_queries } </think> \n "
)
all_sub_queries . extend ( sub_queries ) all_sub_queries . extend ( sub_queries )
sub_gap_queries = sub_queries
for iter in range ( max_iter ) :
log . color_print ( f " >> Iteration: { iter + 1 } \n " )
search_res_from_vectordb = [ ]
search_res_from_internet = [ ] # TODO
# Create all search tasks for it in range ( max_iter ) :
search_tasks = [ log . color_print ( f " >> Iteration: { it + 1 } \n " )
self . _search_chunks_from_vectordb ( query )
for query in sub_gap_queries
] # Execute all search tasks sequentially
# Execute all tasks in parallel and wait for results for query in sub_queries :
search_results = await asyncio . gather ( * search_tasks ) result = self . _search_chunks_from_vectordb ( query )
# Merge all results all_search_results . extend ( result )
for result in search_results : undeduped_len = len ( all_search_results )
search_res = result all_search_results = deduplicate ( all_search_results )
search_res_from_vectordb . extend ( search_res ) deduped_len = len ( all_search_results )
if undeduped_len - deduped_len != 0 :
search_res_from_vectordb = deduplicate ( search_res_from_vectordb ) log . color_print (
f " <search> Removed { undeduped_len - deduped_len } duplicates </search> "
)
# search_res_from_internet = deduplicate_results(search_res_from_internet) # search_res_from_internet = deduplicate_results(search_res_from_internet)
all_search_res . extend ( search_res_from_vectordb + search_res_from_internet ) # all_search_res.extend(search_res_from_vectordb + search_res_from_internet )
if iter == max_iter - 1 : if it == max_iter - 1 :
log . color_print ( " <think> Exceeded maximum iterations. Exiting. </think> \n " ) log . color_print ( " </ think> Exceeded maximum iterations. Exiting. " )
break break
### REFLECTION & GET GAP QUERIES ### ### REFLECTION & GET MORE SUB QUERIES ###
log . color_print ( " <think> Reflecting on the search results... </think> \n " ) log . color_print ( " </ think> Reflecting on the search results... " )
sub_gap_ queries = self . _generate_gap _queries ( sub_queries = self . _generate_more_sub _queries (
original_query , all_sub_queries , all_search_res original_query , all_sub_queries , all_search_results
) )
if not sub_gap_ queries or len ( sub_gap _queries ) == 0 : if not sub_queries or len ( sub_queries ) == 0 :
log . color_print ( " <think> No new search queries were generated. Exiting. </think> \n " ) log . color_print ( " </ think> No new search queries were generated. Exiting. " )
break break
else : else :
log . color_print ( log . color_print (
f " <think> New search queries for next iteration: { sub_gap_queries } </think> \n " f " </think> New search queries for next iteration: { sub_queries } " )
) all_sub_queries . extend ( sub_queries )
all_sub_queries . extend ( sub_gap_queries )
all_search_res = deduplicate ( all_search_res ) all_search_results = deduplicate ( all_search_results )
additional_info = { " all_sub_queries " : all_sub_queries } return all_search_results , all_sub_queries
return all_search_res , additional_info
def query ( self , query : str , * * kwargs ) - > tuple [ str , list [ RetrievalResult ] ] : def query ( self , original_ query: str , * * kwargs ) - > tuple [ str , list [ RetrievalResult ] ] :
""" """
Query the agent and generate an answer based on retrieved documents . Query the agent and generate an answer based on retrieved documents .
@ -387,31 +387,24 @@ class DeepSearch(BaseAgent):
- The generated answer - The generated answer
- A list of retrieved document results - A list of retrieved document results
""" """
all_retrieved_results , additional_info = self . retrieve ( query , * * kwargs ) all_retrieved_results , all_sub_queries = self . retrieve ( original_ query, * * kwargs )
if not all_retrieved_results or len ( all_retrieved_results ) == 0 : if not all_retrieved_results or len ( all_retrieved_results ) == 0 :
return f " No relevant information found for query ' { query } ' . " , [ ] return f " No relevant information found for query ' { original_query } ' . " , [ ]
all_sub_queries = additional_info [ " all_sub_queries " ] chunks = [ ] # type: list[str]
chunk_texts = [ ] for i , chunk in enumerate ( all_retrieved_results ) :
for chunk in all_retrieved_results :
if self . text_window_splitter and " wider_text " in chunk . metadata : if self . text_window_splitter and " wider_text " in chunk . metadata :
chunk_text s . append ( chunk . metadata [ " wider_text " ] ) chunks . append ( chunk . metadata [ " wider_text " ] )
else : else :
chunk_text s . append ( chunk . text ) chunks . append ( f ''' <chunk { i + 1 } > { chunk . text } </chunk { i + 1 } ><reference { i + 1 } > { chunk . reference } </reference { i + 1 } > ''' )
log . color_print ( log . color_print (
f " <think> Summarize answer from all { len ( all_retrieved_results ) } retrieved chunks... </think> \n " f " <think> Summarize answer from all { len ( all_retrieved_results ) } retrieved chunks... </think> \n "
) )
summary_prompt = SUMMARY_PROMPT . format ( summary_prompt = SUMMARY_PROMPT . format (
question = query , original_query = original_ query,
mini_question s= all_sub_queries , all_sub_querie s= all_sub_queries ,
mini_chunk_str = self . _format_chunk_texts ( chunk_texts ) , chunks = " \n " . join ( chunks )
) )
response = self . llm . chat ( [ { " role " : " user " , " content " : summary_prompt } ] ) response = self . llm . chat ( [ { " role " : " user " , " content " : summary_prompt } ] )
log . color_print ( " \n ==== FINAL ANSWER==== \n " ) log . color_print ( " \n ==== FINAL ANSWER==== \n " )
log . color_print ( self . llm . remove_think ( response ) ) log . color_print ( self . llm . remove_think ( response ) )
return self . llm . remove_think ( response ) , all_retrieved_results return self . llm . remove_think ( response ) , all_retrieved_results
def _format_chunk_texts ( self , chunk_texts : list [ str ] ) - > str :
chunk_str = " "
for i , chunk in enumerate ( chunk_texts ) :
chunk_str + = f """ <chunk_ { i } > \n { chunk } \n </chunk_ { i } > \n """
return chunk_str
tanxing