11from enum import Enum
2- from typing import Dict , List , Optional
2+ from typing import Dict , List , Optional , Any , Coroutine
33
4- from .protocol import Document
5- from .providers import KnowledgeProvider , SearchableKnowledgeProvider , SimpleKnowledgeProvider
6- from . search import SearchHit , SearchResults
4+ from .protocol import Document , ProviderSearchResult
5+ from .providers import KnowledgeProvider , SearchableKnowledgeProvider , SimpleKnowledgeProvider , \
6+ ProviderSearchResultAggregator
77from ..models import Model
88from ..logging .logging import get_logging_config
99import logging .config
@@ -24,93 +24,64 @@ class Retrieval(Enum):
2424 Search within the document and retrieve only relevant sections matching the query.
2525 Use this for large or structured documents where only specific parts are needed.
2626 """
27- ALWAYS_WHOLE = "always" ,
28- SEARCHABLE_PIECES = "searchable" ,
29-
27+ WHOLE = "whole" ,
28+ SEARCHABLE = "searchable" ,
3029
3130class Memory :
32- namespace : str
33- providers : List [KnowledgeProvider ]
31+ scope : str
32+ knowledge_providers : List [KnowledgeProvider ]
3433 whole_provider : SimpleKnowledgeProvider
3534 searchable_provider : SearchableKnowledgeProvider
36- search_results : SearchResults
3735 max_knowledge_size : int
3836
3937 def __init__ (
4038 self ,
41- namespace : str ,
39+ scope : str ,
4240 model : Model = Model .default_embedding_model (),
4341 max_knowledge_size : int = 4096 ,
44- providers : Optional [List [KnowledgeProvider ]] = None ,
42+ knowledge_providers : Optional [List [KnowledgeProvider ]] = None ,
4543 ):
4644 """
4745 Initializes a Knowledge instance.
4846
49- :param namespace : A string representing the unique namespace for the knowledge.
47+ :param scope : A string representing the unique scope of the knowledge.
5048 :param max_knowledge_size: Maximum size of knowledge in characters.
51- :param providers : A list of KnowledgeProvider instances to be used for searching.
49+ :param knowledge_providers : A list of KnowledgeProvider instances to be used for searching.
5250 """
53- self .namespace = namespace
54- self .whole_provider = SimpleKnowledgeProvider (namespace )
55- self .searchable_provider = SearchableKnowledgeProvider (namespace , model )
56- self .providers = [self .whole_provider , self .searchable_provider ]
57- if providers :
58- self .providers .extend (providers )
59- self .search_results = SearchResults ()
51+ self .scope = scope
52+ self .whole_provider = SimpleKnowledgeProvider (scope )
53+ self .searchable_provider = SearchableKnowledgeProvider (scope , model )
54+ self .knowledge_providers = [self .whole_provider , self .searchable_provider ]
55+ if knowledge_providers :
56+ self .knowledge_providers .extend (knowledge_providers )
6057 self .max_knowledge_size = max_knowledge_size
6158
62- async def add_document (self , document : Document , retrieval : Retrieval = Retrieval .ALWAYS_WHOLE ) -> None :
59+ async def add (self , document : Document , retrieval : Retrieval = Retrieval .WHOLE ) -> None :
6360 """
6461 Add a document to the knowledge base.
6562
6663 :param document: The document to add.
6764 :param retrieval: Retrieval strategy to use.
6865 :raises ValueError: If the retrieval type is unknown.
6966 """
70- if retrieval == Retrieval .ALWAYS_WHOLE :
71- await self .whole_provider .add_document (document )
72- elif retrieval == Retrieval .SEARCHABLE_PIECES :
73- await self .searchable_provider .add_document (document )
67+ if retrieval == Retrieval .WHOLE :
68+ await self .whole_provider .add (document )
69+ elif retrieval == Retrieval .SEARCHABLE :
70+ await self .searchable_provider .add (document )
7471 else :
7572 raise ValueError (f"Unknown retrieval type: { retrieval } )
7673
77- async def add_query (
74+ async def search (
7875 self , query : str ,
79- ) -> None :
76+ ) -> ProviderSearchResultAggregator :
8077 """
8178 This method will search across all providers and store the results.
8279
8380 :param query: Search query string.
8481 """
85- hits : List [SearchHit ] = []
86- for provider in self .providers :
87- provider_hits = await provider .search (query )
88- hits .extend (provider_hits )
89- self .search_results .add (hits )
90-
91- def render_text (self ) -> Optional [str ]:
92- """
93- Renders the text from the search results into a string.
94-
95- :return: Formatted text from search results, or None when there are no results.
96- """
97- initial_knowledge_size = None
98- while True :
99- view : Dict [str , List [str ]] = self .search_results .view ()
100- if len (view ) > 0 :
101- contextual_knowledge : str = ""
102- for document_name , text_pieces in view .items ():
103- contextual_knowledge += f"Document name: { document_name } \n "
104- for text_piece in text_pieces :
105- contextual_knowledge += f"- { text_piece } \n "
106- contextual_knowledge += "\n "
107-
108- if len (contextual_knowledge ) <= self .max_knowledge_size :
109- if initial_knowledge_size is not None :
110- logger .info (f"Knowledge size reduced from { initial_knowledge_size } { len (contextual_knowledge )} )
111- return contextual_knowledge
112- else :
113- initial_knowledge_size = len (contextual_knowledge )
114- self .search_results .reduce_size ()
115- else :
116- return None
82+ results : List [ProviderSearchResult ] = []
83+ for provider in self .knowledge_providers :
84+ results .append (
85+ await provider .search (query )
86+ )
87+ return ProviderSearchResultAggregator (results )
0 commit comments