ag2 0.9.1__py3-none-any.whl → 0.9.1.post0__py3-none-any.whl

autogen/agentchat/contrib/rag/chromadb_query_engine.py

@@ -0,0 +1,272 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
+
+from ....doc_utils import export_module
+from ....import_utils import optional_import_block, require_optional_import
+from ..vectordb.base import VectorDBFactory
+
+with optional_import_block():
+    from chromadb import HttpClient
+    from chromadb.api.types import EmbeddingFunction
+    from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings
+    from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex
+    from llama_index.core.llms import LLM
+    from llama_index.core.schema import Document as LlamaDocument
+    from llama_index.llms.openai import OpenAI
+    from llama_index.vector_stores.chroma import ChromaVectorStore
+
+__all__ = ["ChromaDBQueryEngine"]
+
+DEFAULT_COLLECTION_NAME = "docling-parsed-docs"
+EMPTY_RESPONSE_TEXT = "Empty Response"  # Indicates that the query did not return any results
+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
+
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("httpx").setLevel(logging.WARNING)
+logger = logging.getLogger(__name__)
+
+
+@require_optional_import(["chromadb", "llama_index"], "rag")
+@export_module("autogen.agentchat.contrib.rag")
+class ChromaDBQueryEngine:
+    """
+    This engine leverages Chromadb to persist document embeddings in a named collection
+    and LlamaIndex's VectorStoreIndex to efficiently index and retrieve documents, and generate an answer in response
+    to natural language queries. Collection can be regarded as an abstraction of group of documents in the database.
+
+    It expects a Chromadb server to be running and accessible at the specified host and port.
+    Refer to this [link](https://docs.trychroma.com/production/containers/docker) for running Chromadb in a Docker container.
+    If the host and port are not provided, the engine will create an in-memory ChromaDB client.
+
+
+    """
+
+    def __init__(  # type: ignore[no-any-unimported]
+        self,
+        host: Optional[str] = "localhost",
+        port: Optional[int] = 8000,
+        settings: Optional["Settings"] = None,
+        tenant: Optional[str] = None,
+        database: Optional[str] = None,
+        embedding_function: "Optional[EmbeddingFunction[Any]]" = None,
+        metadata: Optional[dict[str, Any]] = None,
+        llm: Optional["LLM"] = None,
+        collection_name: Optional[str] = None,
+    ) -> None:
+        """
+        Initializes the ChromaDBQueryEngine with db_path, metadata, and embedding function and llm.
+        Args:
+            host: The host address of the ChromaDB server. Default is localhost.
+            port: The port number of the ChromaDB server. Default is 8000.
+            settings: A dictionary of settings to communicate with the chroma server. Default is None.
+            tenant: The tenant to use for this client. Defaults to the default tenant.
+            database: The database to use for this client. Defaults to the default database.
+            embedding_function: A callable that converts text into vector embeddings. Default embedding uses Sentence Transformers model all-MiniLM-L6-v2.
+                For more embeddings that ChromaDB support, please refer to [embeddings](https://docs.trychroma.com/docs/embeddings/embedding-functions)
+            metadata: A dictionary containing configuration parameters for the Chromadb collection.
+                This metadata is typically used to configure the HNSW indexing algorithm. Defaults to `{"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32}`
+                For more details about the default metadata, please refer to [HNSW configuration](https://cookbook.chromadb.dev/core/configuration/#hnsw-configuration)
+            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/)
+            collection_name (str): The unique name for the Chromadb collection. If omitted, a constant name will be used. Populate this to reuse previous ingested data.
+        """
+        self.llm: LLM = llm or OpenAI(model="gpt-4o", temperature=0.0)  # type: ignore[no-any-unimported]
+        if not host or not port:
+            logger.warning(
+                "Can't connect to remote Chroma client without host or port not. Using an ephemeral, in-memory client."
+            )
+            self.client = None
+        else:
+            try:
+                self.client = HttpClient(
+                    host=host,
+                    port=port,
+                    settings=settings,
+                    tenant=tenant if tenant else DEFAULT_TENANT,  # type: ignore[arg-type, no-any-unimported]
+                    database=database if database else DEFAULT_DATABASE,  # type: ignore[arg-type, no-any-unimported]
+                )
+            except Exception as e:
+                raise ValueError(f"Failed to connect to the ChromaDB client: {e}")
+
+        self.db_config = {"client": self.client, "embedding_function": embedding_function, "metadata": metadata}
+        self.collection_name = collection_name if collection_name else DEFAULT_COLLECTION_NAME
+
+    def init_db(
+        self,
+        new_doc_dir: Optional[Union[Path, str]] = None,
+        new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> bool:
+        """Initialize the database with the input documents or records.
+        It overwrites the existing collection in the database.
+
+        It takes the following steps,
+        1. Set up ChromaDB and LlamaIndex storage.
+        2. insert documents and build indexes upon them.
+
+        Args:
+            new_doc_dir: a dir of input documents that are used to create the records in database.
+            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.
+            *args: Any additional arguments
+            **kwargs: Any additional keyword arguments
+
+        Returns:
+            bool: True if initialization is successful
+
+        """
+
+        self._set_up(overwrite=True)
+        documents = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths_or_urls)
+        self.index = VectorStoreIndex.from_documents(documents=documents, storage_context=self.storage_context)
+        return True
+
+    def connect_db(self, *args: Any, **kwargs: Any) -> bool:
+        """Connect to the database.
+        It does not overwrite the existing collection in the database.
+
+         It takes the following steps,
+        1. Set up ChromaDB and LlamaIndex storage.
+        2. Create the llamaIndex vector store index for querying or inserting docs later
+
+        Args:
+            *args: Any additional arguments
+            **kwargs: Any additional keyword arguments
+
+        Returns:
+            bool: True if connection is successful
+        """
+        self._set_up(overwrite=False)
+        self.index = VectorStoreIndex.from_vector_store(
+            vector_store=self.vector_store, storage_context=self.storage_context
+        )
+
+        return True
+
+    def add_docs(
+        self,
+        new_doc_dir: Optional[Union[Path, str]] = None,
+        new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Add new documents to the underlying database and add to the index.
+
+        Args:
+            new_doc_dir: A dir of input documents that are used to create the records in database.
+            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.
+            *args: Any additional arguments
+            **kwargs: Any additional keyword arguments
+        """
+        self._validate_query_index()
+        documents = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths_or_urls)
+        for doc in documents:
+            self.index.insert(doc)
+
+    def query(self, question: str) -> str:
+        """
+        Retrieve information from indexed documents by processing a query using the engine's LLM.
+
+        Args:
+            question: A natural language query string used to search the indexed documents.
+
+        Returns:
+            A string containing the response generated by LLM.
+        """
+        self._validate_query_index()
+        self.query_engine = self.index.as_query_engine(llm=self.llm)
+        response = self.query_engine.query(question)
+
+        if str(response) == EMPTY_RESPONSE_TEXT:
+            return EMPTY_RESPONSE_REPLY
+
+        return str(response)
+
+    def get_collection_name(self) -> str:
+        """
+        Get the name of the collection used by the query engine.
+
+        Returns:
+            The name of the collection.
+        """
+        if self.collection_name:
+            return self.collection_name
+        else:
+            raise ValueError("Collection name not set.")
+
+    def _validate_query_index(self) -> None:
+        """Ensures an index exists"""
+        if not hasattr(self, "index"):
+            raise Exception("Query index is not initialized. Please call init_db or connect_db first.")
+
+    def _set_up(self, overwrite: bool) -> None:
+        """
+        Set up ChromaDB and LlamaIndex storage by:
+        1. Initialize the ChromaDB using VectorDBFactory and create a collection with the given name.
+        2. Create the LlamaIndex vector store and storage context for the collection.
+        Args:
+            overwrite: If True, overwrite the existing collection with the same name.
+        """
+        self.vector_db = VectorDBFactory().create_vector_db(db_type="chroma", **self.db_config)
+        self.collection = self.vector_db.create_collection(collection_name=self.collection_name, overwrite=overwrite)
+        self.vector_store = ChromaVectorStore(chroma_collection=self.collection)
+        self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
+
+    def _load_doc(  # type: ignore[no-any-unimported]
+        self, input_dir: Optional[Union[Path, str]], input_docs: Optional[Sequence[Union[Path, str]]]
+    ) -> Sequence["LlamaDocument"]:
+        """
+        Load documents from a directory and/or a sequence of file paths.
+
+        It uses LlamaIndex's SimpleDirectoryReader that supports multiple file[formats]((https://docs.llamaindex.ai/en/stable/module_guides/loading/simpledirectoryreader/#supported-file-types)).
+
+        Args:
+            input_dir (Optional[Union[Path, str]]): The directory containing documents to be loaded.
+                If provided, all files in the directory will be considered.
+            input_docs (Optional[Sequence[Union[Path, str]]]): A sequence of individual file paths to load.
+                Each path must point to an existing file.
+
+        Returns:
+            A sequence of documents loaded as LlamaDocument objects.
+
+        Raises:
+            ValueError: If the specified directory does not exist.
+            ValueError: If any provided file path does not exist.
+            ValueError: If neither input_dir nor input_docs is provided.
+        """
+        loaded_documents = []
+        if input_dir:
+            logger.info(f"Loading docs from directory: {input_dir}")
+            if not os.path.exists(input_dir):
+                raise ValueError(f"Input directory not found: {input_dir}")
+            loaded_documents.extend(SimpleDirectoryReader(input_dir=input_dir).load_data())
+
+        if input_docs:
+            for doc in input_docs:
+                logger.info(f"Loading input doc: {doc}")
+                if not os.path.exists(doc):
+                    raise ValueError(f"Document file not found: {doc}")
+            loaded_documents.extend(SimpleDirectoryReader(input_files=input_docs).load_data())  # type: ignore[arg-type]
+
+        if not input_dir and not input_docs:
+            raise ValueError("No input directory or docs provided!")
+
+        return loaded_documents
+
+
+# mypy will fail if ChromaDBQueryEngine does not implement RAGQueryEngine protocol
+if TYPE_CHECKING:
+    from .query_engine import RAGQueryEngine
+
+    def _check_implement_protocol(o: ChromaDBQueryEngine) -> RAGQueryEngine:
+        return o

autogen/agentchat/contrib/rag/llamaindex_query_engine.py

@@ -0,0 +1,198 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
+
+from ....doc_utils import export_module
+from ....import_utils import optional_import_block, require_optional_import
+
+with optional_import_block():
+    from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex
+    from llama_index.core.llms import LLM
+    from llama_index.core.schema import Document as LlamaDocument
+    from llama_index.core.vector_stores.types import BasePydanticVectorStore
+    from llama_index.llms.openai import OpenAI
+
+__all__ = ["LlamaIndexQueryEngine"]
+
+
+EMPTY_RESPONSE_TEXT = "Empty Response"  # Indicates that the query did not return any results
+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
+
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("httpx").setLevel(logging.WARNING)
+logger = logging.getLogger(__name__)
+
+
+@require_optional_import("llama_index", "rag")
+@export_module("autogen.agentchat.contrib.rag")
+class LlamaIndexQueryEngine:
+    """
+    This engine leverages LlamaIndex's VectorStoreIndex to efficiently index and retrieve documents, and generate an answer in response
+    to natural language queries. It use any LlamaIndex [vector store](https://docs.llamaindex.ai/en/stable/module_guides/storing/vector_stores/).
+
+    By default the engine will use OpenAI's GPT-4o model (use the `llm` parameter to change that).
+    """
+
+    def __init__(  # type: ignore[no-any-unimported]
+        self,
+        vector_store: "BasePydanticVectorStore",
+        llm: Optional["LLM"] = None,
+        file_reader_class: Optional[type["SimpleDirectoryReader"]] = None,
+    ) -> None:
+        """
+        Initializes the LlamaIndexQueryEngine with the given vector store.
+        Args:
+            vector_store: The vector store to use for indexing and querying documents.
+            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/).
+            file_reader_class: The file reader class to use for loading documents. Only SimpleDirectoryReader is currently supported.
+        """
+        self.llm: LLM = llm or OpenAI(model="gpt-4o", temperature=0.0)  # type: ignore[no-any-unimported]
+        self.vector_store = vector_store
+        self.file_reader_class = file_reader_class if file_reader_class else SimpleDirectoryReader
+
+    def init_db(
+        self,
+        new_doc_dir: Optional[Union[Path, str]] = None,
+        new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> bool:
+        """Initialize the database with the input documents or records.
+
+        It takes the following steps:
+        1. Set up LlamaIndex storage context.
+        2. insert documents and build an index upon them.
+
+        Args:
+            new_doc_dir: a dir of input documents that are used to create the records in database.
+            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.
+            *args: Any additional arguments
+            **kwargs: Any additional keyword arguments
+
+        Returns:
+            bool: True if initialization is successful
+
+        """
+
+        self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
+        documents = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths_or_urls)
+        self.index = VectorStoreIndex.from_documents(documents=documents, storage_context=self.storage_context)
+        return True
+
+    def connect_db(self, *args: Any, **kwargs: Any) -> bool:
+        """Connect to the database.
+        It sets up the LlamaIndex storage and create an index from the existing vector store.
+
+        Args:
+            *args: Any additional arguments
+            **kwargs: Any additional keyword arguments
+
+        Returns:
+            bool: True if connection is successful
+        """
+        self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
+        self.index = VectorStoreIndex.from_vector_store(
+            vector_store=self.vector_store, storage_context=self.storage_context
+        )
+
+        return True
+
+    def add_docs(
+        self,
+        new_doc_dir: Optional[Union[Path, str]] = None,
+        new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Add new documents to the underlying database and add to the index.
+
+        Args:
+            new_doc_dir: A dir of input documents that are used to create the records in database.
+            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.
+            *args: Any additional arguments
+            **kwargs: Any additional keyword arguments
+        """
+        self._validate_query_index()
+        documents = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths_or_urls)
+        for doc in documents:
+            self.index.insert(doc)
+
+    def query(self, question: str) -> str:
+        """
+        Retrieve information from indexed documents by processing a query using the engine's LLM.
+
+        Args:
+            question: A natural language query string used to search the indexed documents.
+
+        Returns:
+            A string containing the response generated by LLM.
+        """
+        self._validate_query_index()
+        self.query_engine = self.index.as_query_engine(llm=self.llm)
+        response = self.query_engine.query(question)
+
+        if str(response) == EMPTY_RESPONSE_TEXT:
+            return EMPTY_RESPONSE_REPLY
+
+        return str(response)
+
+    def _validate_query_index(self) -> None:
+        """Ensures an index exists"""
+        if not hasattr(self, "index"):
+            raise Exception("Query index is not initialized. Please call init_db or connect_db first.")
+
+    def _load_doc(  # type: ignore[no-any-unimported]
+        self, input_dir: Optional[Union[Path, str]], input_docs: Optional[Sequence[Union[Path, str]]]
+    ) -> Sequence["LlamaDocument"]:
+        """
+        Load documents from a directory and/or a sequence of file paths.
+
+        Default to uses LlamaIndex's SimpleDirectoryReader that supports multiple file[formats](https://docs.llamaindex.ai/en/stable/module_guides/loading/simpledirectoryreader/#supported-file-types).
+
+        Args:
+            input_dir (Optional[Union[Path, str]]): The directory containing documents to be loaded.
+                If provided, all files in the directory will be considered.
+            input_docs (Optional[Sequence[Union[Path, str]]]): A sequence of individual file paths to load.
+                Each path must point to an existing file.
+
+        Returns:
+            A sequence of documents loaded as LlamaDocument objects.
+
+        Raises:
+            ValueError: If the specified directory does not exist.
+            ValueError: If any provided file path does not exist.
+            ValueError: If neither input_dir nor input_docs is provided.
+        """
+        loaded_documents: list["LlamaDocument"] = []  # type: ignore[no-any-unimported]
+        if input_dir:
+            logger.info(f"Loading docs from directory: {input_dir}")
+            if not os.path.exists(input_dir):
+                raise ValueError(f"Input directory not found: {input_dir}")
+            loaded_documents.extend(self.file_reader_class(input_dir=input_dir).load_data())  # type: ignore[operator]
+
+        if input_docs:
+            for doc in input_docs:
+                logger.info(f"Loading input doc: {doc}")
+                if not os.path.exists(doc):
+                    raise ValueError(f"Document file not found: {doc}")
+            loaded_documents.extend(self.file_reader_class(input_files=input_docs).load_data())  # type: ignore[operator, arg-type]
+
+        if not input_dir and not input_docs:
+            raise ValueError("No input directory or docs provided!")
+
+        return loaded_documents
+
+
+# mypy will fail if LlamaIndexQueryEngine does not implement RAGQueryEngine protocol
+if TYPE_CHECKING:
+    from .query_engine import RAGQueryEngine
+
+    def _check_implement_protocol(o: LlamaIndexQueryEngine) -> RAGQueryEngine:
+        return o