InMemoryQueryEngine - AG2

InMemoryQueryEngine

`autogen.agents.experimental.InMemoryQueryEngine`

InMemoryQueryEngine(llm_config)

This engine stores ingested documents in memory and then injects them into an internal agent's system message for answering queries.

This implements the autogen.agentchat.contrib.rag.RAGQueryEngine protocol.

Source code in autogen/agents/experimental/document_agent/inmemory_query_engine.py

 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64  def __init__( self, llm_config: Union[LLMConfig, dict[str, Any]], ) -> None: # Deep copy the llm config to avoid changing the original structured_config = copy.deepcopy(llm_config) # The query agent will answer with a structured output structured_config["response_format"] = QueryAnswer # Our agents for querying self._query_agent = ConversableAgent( name="inmemory_query_agent", llm_config=structured_config, ) # In-memory storage for ingested documents self._ingested_documents: list[DocumentStore] = []


`<br>47<br>48<br>49<br>50<br>51<br>52<br>53<br>54<br>55<br>56<br>57<br>58<br>59<br>60<br>61<br>62<br>63<br>64<br>`	<br>def __init__(<br> self,<br> llm_config: Union[LLMConfig, dict[str, Any]],<br>) -> None:<br> # Deep copy the llm config to avoid changing the original<br> structured_config = copy.deepcopy(llm_config)<br> # The query agent will answer with a structured output<br> structured_config["response_format"] = QueryAnswer<br> # Our agents for querying<br> self._query_agent = ConversableAgent(<br> name="inmemory_query_agent",<br> llm_config=structured_config,<br> )<br> # In-memory storage for ingested documents<br> self._ingested_documents: list[DocumentStore] = []<br>

`query`

query(question, *args, **kwargs)

Run a query against the ingested documents and return the answer.

Source code in autogen/agents/experimental/document_agent/inmemory_query_engine.py

 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115  def query(self, question: str, *args: Any, **kwargs: Any) -> str: """Run a query against the ingested documents and return the answer.""" # If no documents have been ingested, return an empty response if not self._ingested_documents: return QUERY_NO_INGESTIONS_REPLY # Put the context into the system message context_parts = [] for i, doc in enumerate(self._ingested_documents, 1): context_parts.append(f"Ingested File/URL {i} - '{doc.ingestation_name}':\n{doc.content}\n") context = "\n".join(context_parts) system_message = ( "You are a query agent tasked with answering questions based on ingested documents.\n\n" "AVAILABLE DOCUMENTS:\n" + "\n".join([f"- {doc.ingestation_name}" for doc in self._ingested_documents]) + "\n\n" "When answering questions about these documents, use ONLY the information in the following context:\n\n" f"{context}\n\n" "IMPORTANT: The user will ask about these documents by name. When they do, provide helpful, detailed answers based on the document content above." ) self._query_agent.update_system_message(system_message) message = f"Using ONLY the document content in your system message, answer this question: {question}" response = self._query_agent.run( message=message, max_turns=1, ) response.process() try: # Get the structured output and return the answer answer_object = QueryAnswer.model_validate(json.loads(response.summary)) # type: ignore[arg-type] if answer_object.could_answer: return answer_object.answer else: if answer_object.answer: return COULD_NOT_ANSWER_REPLY + ": " + answer_object.answer else: return COULD_NOT_ANSWER_REPLY except Exception as e: # Error converting the response to the structured output return ERROR_RESPONSE_REPLY + str(e)


`<br> 66<br> 67<br> 68<br> 69<br> 70<br> 71<br> 72<br> 73<br> 74<br> 75<br> 76<br> 77<br> 78<br> 79<br> 80<br> 81<br> 82<br> 83<br> 84<br> 85<br> 86<br> 87<br> 88<br> 89<br> 90<br> 91<br> 92<br> 93<br> 94<br> 95<br> 96<br> 97<br> 98<br> 99<br>100<br>101<br>102<br>103<br>104<br>105<br>106<br>107<br>108<br>109<br>110<br>111<br>112<br>113<br>114<br>115<br>`	<br>def query(self, question: str, args: Any, *kwargs: Any) -> str:<br> """Run a query against the ingested documents and return the answer."""<br> # If no documents have been ingested, return an empty response<br> if not self._ingested_documents:<br> return QUERY_NO_INGESTIONS_REPLY<br> # Put the context into the system message<br> context_parts = []<br> for i, doc in enumerate(self._ingested_documents, 1):<br> context_parts.append(f"Ingested File/URL {i} - '{doc.ingestation_name}':\n{doc.content}\n")<br> context = "\n".join(context_parts)<br> system_message = (<br> "You are a query agent tasked with answering questions based on ingested documents.\n\n"<br> "AVAILABLE DOCUMENTS:\n"<br> + "\n".join([f"- {doc.ingestation_name}" for doc in self._ingested_documents])<br> + "\n\n"<br> "When answering questions about these documents, use ONLY the information in the following context:\n\n"<br> f"{context}\n\n"<br> "IMPORTANT: The user will ask about these documents by name. When they do, provide helpful, detailed answers based on the document content above."<br> )<br> self._query_agent.update_system_message(system_message)<br> message = f"Using ONLY the document content in your system message, answer this question: {question}"<br> response = self._query_agent.run(<br> message=message,<br> max_turns=1,<br> )<br> response.process()<br> try:<br> # Get the structured output and return the answer<br> answer_object = QueryAnswer.model_validate(json.loads(response.summary)) # type: ignore[arg-type]<br> if answer_object.could_answer:<br> return answer_object.answer<br> else:<br> if answer_object.answer:<br> return COULD_NOT_ANSWER_REPLY + ": " + answer_object.answer<br> else:<br> return COULD_NOT_ANSWER_REPLY<br> except Exception as e:<br> # Error converting the response to the structured output<br> return ERROR_RESPONSE_REPLY + str(e)<br>

`add_docs`

add_docs(new_doc_dir=None, new_doc_paths_or_urls=None)

Add additional documents to the in-memory store

Loads new Docling-parsed Markdown files from a specified directory or a list of file paths and inserts them into the in-memory store.

PARAMETER	DESCRIPTION
`new_doc_dir`	The directory path from which to load additional documents. If provided, all eligible files in this directory are loaded. TYPE:`Optional[Union[Path, str]]`DEFAULT:`None`
`new_doc_paths_or_urls`	A list of file paths specifying additional documents to load. Each file should be a Docling-parsed Markdown file. TYPE:`Optional[Sequence[Union[Path, str]]]`DEFAULT:`None`

Source code in autogen/agents/experimental/document_agent/inmemory_query_engine.py

 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136  def add_docs( self, new_doc_dir: Optional[Union[Path, str]] = None, new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None, ) -> None: """ Add additional documents to the in-memory store Loads new Docling-parsed Markdown files from a specified directory or a list of file paths and inserts them into the in-memory store. Args: new_doc_dir: The directory path from which to load additional documents. If provided, all eligible files in this directory are loaded. new_doc_paths_or_urls: A list of file paths specifying additional documents to load. Each file should be a Docling-parsed Markdown file. """ new_doc_dir = new_doc_dir or "" new_doc_paths = new_doc_paths_or_urls or [] self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths)


`<br>117<br>118<br>119<br>120<br>121<br>122<br>123<br>124<br>125<br>126<br>127<br>128<br>129<br>130<br>131<br>132<br>133<br>134<br>135<br>136<br>`	<br>def add_docs(<br> self,<br> new_doc_dir: Optional[Union[Path, str]] = None,<br> new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,<br>) -> None:<br> """<br> Add additional documents to the in-memory store<br> Loads new Docling-parsed Markdown files from a specified directory or a list of file paths<br> and inserts them into the in-memory store.<br> Args:<br> new_doc_dir: The directory path from which to load additional documents.<br> If provided, all eligible files in this directory are loaded.<br> new_doc_paths_or_urls: A list of file paths specifying additional documents to load.<br> Each file should be a Docling-parsed Markdown file.<br> """<br> new_doc_dir = new_doc_dir or ""<br> new_doc_paths = new_doc_paths_or_urls or []<br> self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths)<br>

`init_db`

init_db(new_doc_dir=None, new_doc_paths_or_urls=None, *args, **kwargs)

Not required nor implemented for InMemoryQueryEngine

Source code in autogen/agents/experimental/document_agent/inmemory_query_engine.py


`<br>200<br>201<br>202<br>203<br>204<br>205<br>206<br>207<br>208<br>`	`<br>def init_db(<br> self,<br> new_doc_dir: Optional[Union[Path, str]] = None,<br> new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,<br> args: Any,<br> *kwargs: Any,<br>) -> bool:<br> """Not required nor implemented for InMemoryQueryEngine"""<br> raise NotImplementedError("Method, init_db, not required nor implemented for InMemoryQueryEngine")<br>`

`connect_db`

connect_db(*args, **kwargs)

Not required nor implemented for InMemoryQueryEngine

Source code in autogen/agents/experimental/document_agent/inmemory_query_engine.py


`<br>210<br>211<br>212<br>`	`<br>def connect_db(self, args: Any, *kwargs: Any) -> bool:<br> """Not required nor implemented for InMemoryQueryEngine"""<br> raise NotImplementedError("Method, connect_db, not required nor implemented for InMemoryQueryEngine")<br>`