Skip to content

Commit ff77a43

Browse files
Eric-ShangmarklyszedavorrunjerjambrecicAgentGenie
authored
Create LlamaIndex query engine (#1311)
* Create a LlamaIndex query engine * Create unit tests * Fix file name * Fix file name * Change name * Updated notebooks, more comments, tidy. * Update autogen/agentchat/contrib/rag/llamaindex_query_engine.py Co-authored-by: Mark Sze <[email protected]> * Update autogen/agentchat/contrib/rag/llamaindex_query_engine.py Co-authored-by: Mark Sze <[email protected]> * Update autogen/agentchat/contrib/rag/llamaindex_query_engine.py Co-authored-by: Mark Sze <[email protected]> * Update autogen/agentchat/contrib/rag/llamaindex_query_engine.py Co-authored-by: Mark Sze <[email protected]> * Update autogen/agentchat/contrib/rag/llamaindex_query_engine.py Co-authored-by: Mark Sze <[email protected]> * Notebook metadata, pre-commit fix * Termination messages (#1336) * Addition of termination messages * Tests and tidy * Tidy, remove full stop * Correct code execution text (#1335) * Add silent_override for suppressing overriding function messages, updated dependency injection parameter to avoid pydantic warning message (#1334) * [Docs]: Swarm Patterns documentation (#1255) * Hierarchical pattern * Hierarchal updates, Star, Escalation patterns * Feedback Loop, Pipeline, Section overview * Context-Aware Routing and Redundant, tidy * Organic pattern * Triage with Tasks pattern, move to Pattern Cookbook * Remove blank. * Diagrams * Add package version to import utils (#1331) * wip * wip * wipp * versioning added to import utils * fix for missing version * Implement GoogleSearchTool (#1315) * wip * wip * wip * Google search tool WIP * Google search tool WIP * Initial GoogleSearchTool implementation * Add gemini tests * Update dependencies and initial tests for the google search tool added * Update tests * Cleanup * Cleanup * Fix tests * Fix tests * Fix tests * Google search notebook WIP * Documentation polishing * Update user guides * Refactor and fix tests * Update docs --------- Co-authored-by: Davor Runje <[email protected]> * Release 0.8.2rc0 (#1346) * [DocAgent] Add Text Citation (#1338) * Add table output option * Fix bug * Add VectorChromaCitationQueryEngine to support citations * Update DocAgent to support citations * Update DocAgent Notebook with Citation Example * polishing * update Summary Agent prompt to handle citations and non-citations * type hint with string literal --------- Co-authored-by: Davor Runje <[email protected]> Co-authored-by: Mark Sze <[email protected]> * Ollama response_format fix (#1347) * Ollama response_format fix * Ollama tests with response format passed in and conversion check * pre-commit tidy * Fix get load param if needed function (#1344) * test added * test added * DocAgent fix * DocAgent fix * DocAgent fix * DocAgent fix * Main and tweaks to function field annotations --------- Co-authored-by: Mark Sze <[email protected]> * Make Document Reader customizable. * Move SimpleDirectoryReader into init body --------- Co-authored-by: Mark Sze <[email protected]> Co-authored-by: Mark Sze <[email protected]> Co-authored-by: Davor Runje <[email protected]> Co-authored-by: Robert Jambrečić <[email protected]> Co-authored-by: Davor Runje <[email protected]> Co-authored-by: AgentGenie <[email protected]>
1 parent 88f7df2 commit ff77a43

File tree

7 files changed

+687
-15
lines changed

7 files changed

+687
-15
lines changed

‎autogen/agentchat/contrib/rag/__init__.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,9 @@
22
#
33
# SPDX-License-Identifier: Apache-2.0
44

5-
from .chroma_db_query_engine import ChromaDBQueryEngine
5+
from .chromadb_query_engine import ChromaDBQueryEngine
6+
from .llamaindex_query_engine import LlamaIndexQueryEngine
67
from .mongodb_query_engine import MongoDBQueryEngine
78
from .query_engine import RAGQueryEngine
89

9-
__all__ = ["ChromaDBQueryEngine", "MongoDBQueryEngine", "RAGQueryEngine"]
10+
__all__ = ["ChromaDBQueryEngine", "LlamaIndexQueryEngine", "MongoDBQueryEngine", "RAGQueryEngine"]

autogen/agentchat/contrib/rag/chroma_db_query_engine.py renamed to autogen/agentchat/contrib/rag/chromadb_query_engine.py

File renamed without changes.

‎autogen/agentchat/contrib/rag/llamaindex_query_engine.py‎

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
import logging
6+
import os
7+
from pathlib import Path
8+
from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
9+
10+
from ....doc_utils import export_module
11+
from ....import_utils import optional_import_block, require_optional_import
12+
13+
with optional_import_block():
14+
from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex
15+
from llama_index.core.llms import LLM
16+
from llama_index.core.readers.base import BaseReader
17+
from llama_index.core.schema import Document as LlamaDocument
18+
from llama_index.core.vector_stores.types import BasePydanticVectorStore
19+
from llama_index.llms.openai import OpenAI
20+
21+
__all__ = ["LlamaIndexQueryEngine"]
22+
23+
24+
EMPTY_RESPONSE_TEXT = "Empty Response" # Indicates that the query did not return any results
25+
EMPTY_RESPONSE_REPLY = "Sorry, I couldn't find any information on that. If you haven't ingested any documents, please try that." # Default response for queries without results
26+
27+
28+
# Set up logging
29+
logging.basicConfig(level=logging.INFO)
30+
logging.getLogger("httpx").setLevel(logging.WARNING)
31+
logger = logging.getLogger(__name__)
32+
33+
34+
@require_optional_import("llama_index", "rag")
35+
@export_module("autogen.agentchat.contrib.rag")
36+
class LlamaIndexQueryEngine:
37+
"""
38+
This engine leverages LlamaIndex's VectorStoreIndex to efficiently index and retrieve documents, and generate an answer in response
39+
to natural language queries. It use any LlamaIndex [vector store](https://docs.llamaindex.ai/en/stable/module_guides/storing/vector_stores/).
40+
41+
By default the engine will use OpenAI's GPT-4o model (use the `llm` parameter to change that).
42+
"""
43+
44+
def __init__( # type: ignore[no-any-unimported]
45+
self,
46+
vector_store: "BasePydanticVectorStore",
47+
llm: Optional["LLM"] = None,
48+
file_reader_class: Optional["BaseReader"] = None,
49+
) -> None:
50+
"""
51+
Initializes the LlamaIndexQueryEngine with the given vector store.
52+
Args:
53+
vector_store: The vector store to use for indexing and querying documents.
54+
llm: LLM model used by LlamaIndex for query processing. You can find more supported LLMs at [LLM](https://docs.llamaindex.ai/en/stable/module_guides/models/llms/).
55+
file_reader_class: The file reader class to use for loading documents. Defaults to SimpleDirectoryReader.
56+
"""
57+
self.llm: LLM = llm or OpenAI(model="gpt-4o", temperature=0.0) # type: ignore[no-any-unimported]
58+
self.vector_store = vector_store
59+
self.file_reader_class = file_reader_class if file_reader_class else SimpleDirectoryReader
60+
61+
def init_db(
62+
self,
63+
new_doc_dir: Optional[Union[Path, str]] = None,
64+
new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,
65+
*args: Any,
66+
**kwargs: Any,
67+
) -> bool:
68+
"""Initialize the database with the input documents or records.
69+
70+
It takes the following steps:
71+
1. Set up LlamaIndex storage context.
72+
2. insert documents and build an index upon them.
73+
74+
Args:
75+
new_doc_dir: a dir of input documents that are used to create the records in database.
76+
new_doc_paths_or_urls: A sequence of input documents that are used to create the records in database. A document can be a Path to a file or a url.
77+
*args: Any additional arguments
78+
**kwargs: Any additional keyword arguments
79+
80+
Returns:
81+
bool: True if initialization is successful
82+
83+
"""
84+
85+
self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
86+
documents = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths_or_urls)
87+
self.index = VectorStoreIndex.from_documents(documents=documents, storage_context=self.storage_context)
88+
return True
89+
90+
def connect_db(self, *args: Any, **kwargs: Any) -> bool:
91+
"""Connect to the database.
92+
It sets up the LlamaIndex storage and create an index from the existing vector store.
93+
94+
Args:
95+
*args: Any additional arguments
96+
**kwargs: Any additional keyword arguments
97+
98+
Returns:
99+
bool: True if connection is successful
100+
"""
101+
self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
102+
self.index = VectorStoreIndex.from_vector_store(
103+
vector_store=self.vector_store, storage_context=self.storage_context
104+
)
105+
106+
return True
107+
108+
def add_docs(
109+
self,
110+
new_doc_dir: Optional[Union[Path, str]] = None,
111+
new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,
112+
*args: Any,
113+
**kwargs: Any,
114+
) -> None:
115+
"""Add new documents to the underlying database and add to the index.
116+
117+
Args:
118+
new_doc_dir: A dir of input documents that are used to create the records in database.
119+
new_doc_paths_or_urls: A sequence of input documents that are used to create the records in database. A document can be a Path to a file or a url.
120+
*args: Any additional arguments
121+
**kwargs: Any additional keyword arguments
122+
"""
123+
self._validate_query_index()
124+
documents = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths_or_urls)
125+
for doc in documents:
126+
self.index.insert(doc)
127+
128+
def query(self, question: str) -> str:
129+
"""
130+
Retrieve information from indexed documents by processing a query using the engine's LLM.
131+
132+
Args:
133+
question: A natural language query string used to search the indexed documents.
134+
135+
Returns:
136+
A string containing the response generated by LLM.
137+
"""
138+
self._validate_query_index()
139+
self.query_engine = self.index.as_query_engine(llm=self.llm)
140+
response = self.query_engine.query(question)
141+
142+
if str(response) == EMPTY_RESPONSE_TEXT:
143+
return EMPTY_RESPONSE_REPLY
144+
145+
return str(response)
146+
147+
def _validate_query_index(self) -> None:
148+
"""Ensures an index exists"""
149+
if not hasattr(self, "index"):
150+
raise Exception("Query index is not initialized. Please call init_db or connect_db first.")
151+
152+
def _load_doc( # type: ignore[no-any-unimported]
153+
self, input_dir: Optional[Union[Path, str]], input_docs: Optional[Sequence[Union[Path, str]]]
154+
) -> Sequence["LlamaDocument"]:
155+
"""
156+
Load documents from a directory and/or a sequence of file paths.
157+
158+
Default to uses LlamaIndex's SimpleDirectoryReader that supports multiple file[formats]((https://docs.llamaindex.ai/en/stable/module_guides/loading/simpledirectoryreader/#supported-file-types)).
159+
160+
Args:
161+
input_dir (Optional[Union[Path, str]]): The directory containing documents to be loaded.
162+
If provided, all files in the directory will be considered.
163+
input_docs (Optional[Sequence[Union[Path, str]]]): A sequence of individual file paths to load.
164+
Each path must point to an existing file.
165+
166+
Returns:
167+
A sequence of documents loaded as LlamaDocument objects.
168+
169+
Raises:
170+
ValueError: If the specified directory does not exist.
171+
ValueError: If any provided file path does not exist.
172+
ValueError: If neither input_dir nor input_docs is provided.
173+
"""
174+
loaded_documents = []
175+
if input_dir:
176+
logger.info(f"Loading docs from directory: {input_dir}")
177+
if not os.path.exists(input_dir):
178+
raise ValueError(f"Input directory not found: {input_dir}")
179+
loaded_documents.extend(self.file_reader_class(input_dir=input_dir).load_data())
180+
181+
if input_docs:
182+
for doc in input_docs:
183+
logger.info(f"Loading input doc: {doc}")
184+
if not os.path.exists(doc):
185+
raise ValueError(f"Document file not found: {doc}")
186+
loaded_documents.extend(self.file_reader_class(input_files=input_docs).load_data()) # type: ignore[arg-type]
187+
188+
if not input_dir and not input_docs:
189+
raise ValueError("No input directory or docs provided!")
190+
191+
return loaded_documents
192+
193+
194+
# mypy will fail if LlamaIndexQueryEngine does not implement RAGQueryEngine protocol
195+
if TYPE_CHECKING:
196+
from .query_engine import RAGQueryEngine
197+
198+
def _check_implement_protocol(o: LlamaIndexQueryEngine) -> RAGQueryEngine:
199+
return o

‎notebook/Chromadb_query_engine.ipynb‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,13 @@
44
"cell_type": "markdown",
55
"metadata": {},
66
"source": [
7-
"# Use ChromaDBQueryengine to query Markdown files \n",
7+
"# Use ChromaDBQueryEngine to query Markdown files \n",
88
"\n",
99
"This notebook demonstrates the use of the `ChromaDBQueryEngine` for retrieval-augmented question answering over documents. It shows how to set up the engine with Docling parsed Markdown files, and execute natural language queries against the indexed data. \n",
1010
"\n",
11-
"The `ChromaDBQueryEngine` integrates persistent ChromaDB vector storage with LlamaIndex for efficient document retrieval."
11+
"The `ChromaDBQueryEngine` integrates persistent ChromaDB vector storage with LlamaIndex for efficient document retrieval.\n",
12+
"\n",
13+
"You can create and add this ChromaDBQueryEngine to [DocAgent](https://docs.ag2.ai/docs/user-guide/reference-agents/docagent) to use."
1214
]
1315
},
1416
{
@@ -21,6 +23,15 @@
2123
"%pip install llama-index==0.12.16"
2224
]
2325
},
26+
{
27+
"cell_type": "markdown",
28+
"metadata": {},
29+
"source": [
30+
"### Load LLM configuration\n",
31+
"\n",
32+
"This demonstration requires an `OPENAI_API_KEY` to be in your environment variables. See [our documentation](https://docs.ag2.ai/docs/user-guide/basic-concepts/llm-configuration/llm-configuration#environment-variables) for guidance."
33+
]
34+
},
2435
{
2536
"cell_type": "code",
2637
"execution_count": null,
@@ -54,10 +65,10 @@
5465
"metadata": {},
5566
"outputs": [],
5667
"source": [
57-
"from autogen.agentchat.contrib.rag.chroma_db_query_engine import ChromaDBQueryEngine\n",
68+
"from autogen.agentchat.contrib.rag.chromadb_query_engine import ChromaDBQueryEngine\n",
5869
"\n",
5970
"query_engine = ChromaDBQueryEngine(\n",
60-
" host=\"172.17.0.3\", # Change this to the IP of the ChromaDB server\n",
71+
" host=\"host.docker.internal\", # Change this to the IP of the ChromaDB server\n",
6172
" port=8000, # Change this to the port of the ChromaDB server\n",
6273
")"
6374
]
@@ -82,9 +93,11 @@
8293
"cell_type": "markdown",
8394
"metadata": {},
8495
"source": [
85-
"Let's ingest a document and query it. \n",
86-
" ```init_db``` will overwrite the existing collection with the same name.\n",
87-
"\n"
96+
"Let's ingest a document and query it.\n",
97+
"\n",
98+
"If you don't have your documents ingested yet, follow the next two cells. Otherwise skip to the `connect_db` cell.\n",
99+
"\n",
100+
" `init_db` will overwrite the existing collection with the same name."
88101
]
89102
},
90103
{
@@ -149,7 +162,7 @@
149162
"outputs": [],
150163
"source": [
151164
"new_docs = [input_dir + \"Toast_financial_report.md\"]\n",
152-
"query_engine.add_records(new_doc_paths_or_urls=new_docs)"
165+
"query_engine.add_docs(new_doc_paths_or_urls=new_docs)"
153166
]
154167
},
155168
{
@@ -172,6 +185,17 @@
172185
}
173186
],
174187
"metadata": {
188+
"front_matter": {
189+
"description": "ChromaDB Query Engine for document queries",
190+
"tags": [
191+
"agents",
192+
"documents",
193+
"RAG",
194+
"docagent",
195+
"chroma",
196+
"chromadb"
197+
]
198+
},
175199
"kernelspec": {
176200
"display_name": "Python 3",
177201
"language": "python",
@@ -187,7 +211,7 @@
187211
"name": "python",
188212
"nbconvert_exporter": "python",
189213
"pygments_lexer": "ipython3",
190-
"version": "3.10.16"
214+
"version": "3.9.21"
191215
}
192216
},
193217
"nbformat": 4,

0 commit comments

Comments
 (0)