ag2 0.10.2__py3-none-any.whl

autogen/agents/experimental/document_agent/chroma_query_engine.py

@@ -0,0 +1,301 @@
+# 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 collections.abc import Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+from pydantic import BaseModel
+
+from ....doc_utils import export_module
+from ....import_utils import optional_import_block, require_optional_import
+
+with optional_import_block():
+    import chromadb
+    from chromadb.api.models.Collection import Collection
+    from chromadb.api.types import EmbeddingFunction
+    from chromadb.utils.embedding_functions import DefaultEmbeddingFunction
+    from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex
+    from llama_index.core.llms import LLM
+    from llama_index.core.query_engine import CitationQueryEngine
+    from llama_index.core.schema import Document as LlamaDocument
+    from llama_index.core.schema import NodeWithScore
+    from llama_index.llms.openai import OpenAI
+    from llama_index.vector_stores.chroma import ChromaVectorStore
+
+__all__ = ["VectorChromaCitationQueryEngine", "VectorChromaQueryEngine"]
+
+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.agents.experimental")
+class VectorChromaQueryEngine:
+    """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. The Chromadb collection serves as the storage layer, while
+    the collection name uniquely identifies the set of documents within the persistent database.
+
+    This implements the autogen.agentchat.contrib.rag.RAGQueryEngine protocol.
+    """
+
+    def __init__(  # type: ignore[no-any-unimported]
+        self,
+        db_path: str | None = None,
+        embedding_function: "EmbeddingFunction[Any] | None" = None,
+        metadata: dict[str, Any] | None = None,
+        llm: Optional["LLM"] = None,
+        collection_name: str | None = None,
+    ) -> None:
+        """Initializes the VectorChromaQueryEngine with db_path, metadata, and embedding function and llm.
+
+        Args:
+            db_path: The file system path where Chromadb will store its persistent data.
+                If not specified, the default directory "./chroma" is used.
+            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.
+                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]
+        self.embedding_function: EmbeddingFunction[Any] = embedding_function or DefaultEmbeddingFunction()  # type: ignore[no-any-unimported,assignment]
+        self.metadata: dict[str, Any] = metadata or {
+            "hnsw:space": "ip",
+            "hnsw:construction_ef": 30,
+            "hnsw:M": 32,
+        }
+        self.client = chromadb.PersistentClient(path=db_path or "./chroma")
+        self.collection_name: str | None = collection_name
+
+        self.connect_db()
+
+    def connect_db(self, *args: Any, **kwargs: Any) -> bool:
+        """Establish a connection to the Chromadb database and initialize the collection."""
+        self.collection_name = self.collection_name or DEFAULT_COLLECTION_NAME
+
+        if self._collection_exists(self.collection_name):
+            logger.info(f"Using existing collection {self.collection_name} from the database.")
+        else:
+            logger.info(f"Creating new collection {self.collection_name} in the database.")
+
+        self.collection = self.client.create_collection(
+            name=self.collection_name,
+            embedding_function=self.embedding_function,
+            metadata=self.metadata,
+            get_or_create=True,  # If collection already exists, get the collection
+        )
+        self.index = self._create_index(self.collection)
+
+        return True
+
+    def query(self, question: str) -> str:
+        """Retrieve information from indexed documents by processing a natural language query.
+
+        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 add_docs(
+        self,
+        new_doc_dir: Path | str | None = None,
+        new_doc_paths_or_urls: Sequence[Path | str] | None = None,
+    ) -> None:
+        """Add additional documents to the existing vector index.
+
+        Loads new Docling-parsed Markdown files from a specified directory or a list of file paths
+        and inserts them into the current index for future queries.
+
+        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.
+        """
+        self.validate_query_index()
+        new_doc_dir = new_doc_dir or ""
+        new_doc_paths = new_doc_paths_or_urls or []
+        new_docs = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths)
+        for doc in new_docs:
+            self.index.insert(doc)
+
+    def _load_doc(  # type: ignore[no-any-unimported]
+        self, input_dir: Path | str | None, input_docs: Sequence[Path | str] | None
+    ) -> list["LlamaDocument"]:
+        """Load documents from a directory and/or a list of file paths.
+
+        This helper method reads Docling-parsed Markdown files using LlamaIndex's
+        SimpleDirectoryReader. It supports multiple file [formats]((https://docs.llamaindex.ai/en/stable/module_guides/loading/simpledirectoryreader/#supported-file-types)),
+          but the intended use is for documents processed by Docling.
+
+        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[list[Union[Path, str]]]): A list of individual file paths to load.
+                Each path must point to an existing file.
+
+        Returns:
+            A list 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=list(input_docs)).load_data())
+
+        if not input_dir and not input_docs:
+            raise ValueError("No input directory or docs provided!")
+
+        return loaded_documents
+
+    def _create_index(  # type: ignore[no-any-unimported]
+        self, collection: "Collection"
+    ) -> "VectorStoreIndex":
+        """Build a vector index for document retrieval using a Chromadb collection.
+
+        Wraps the provided Chromadb collection into a vector store and uses LlamaIndex's
+        StorageContext to create a VectorStoreIndex from the collection.
+
+        Args:
+            collection (Collection): A Chromadb Collection object that stores the embeddings of the documents.
+
+        Returns:
+            A VectorStoreIndex object built from the provided collection.
+        """
+        self.vector_store = ChromaVectorStore(chroma_collection=collection)
+        self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
+
+        index = VectorStoreIndex.from_vector_store(vector_store=self.vector_store, storage_context=self.storage_context)
+
+        return index
+
+    def _collection_exists(self, collection_name: str) -> bool:
+        """Check if a collection with the given name exists in the database.
+
+        Args:
+            collection_name (str): The name of the collection to check.
+
+        Returns:
+            True if a collection with the given name exists in the database, False otherwise.
+        """
+        existing_collections = self.client.list_collections()
+        return any(col == collection_name for col in existing_collections)
+
+    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 ingest some documents before querying.")
+
+    def init_db(
+        self,
+        new_doc_dir: Path | str | None = None,
+        new_doc_paths_or_urls: Sequence[Path | str] | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> bool:
+        """Not required nor implemented for VectorChromaQueryEngine"""
+        raise NotImplementedError("Method, init_db, not required nor implemented for VectorChromaQueryEngine")
+
+
+class AnswerWithCitations(BaseModel):  # type: ignore[no-any-unimported]
+    answer: str
+    citations: list["NodeWithScore"]  # type: ignore[no-any-unimported]
+
+
+@require_optional_import(["chromadb", "llama_index"], "rag")
+@export_module("autogen.agents.experimental")
+class VectorChromaCitationQueryEngine(VectorChromaQueryEngine):
+    """This engine leverages VectorChromaQueryEngine and CitationQueryEngine to answer queries with citations."""
+
+    def __init__(  # type: ignore[no-any-unimported]
+        self,
+        db_path: str | None = None,
+        embedding_function: "EmbeddingFunction[Any] | None" = None,
+        metadata: dict[str, Any] | None = None,
+        llm: Optional["LLM"] = None,
+        collection_name: str | None = None,
+        enable_query_citations: bool = False,
+        citation_chunk_size: int = 512,
+    ) -> None:
+        """See parent class VectorChromaQueryEngine."""
+        super().__init__(db_path, embedding_function, metadata, llm, collection_name)
+        self.enable_query_citations = enable_query_citations
+        self.citation_chunk_size = citation_chunk_size
+
+    def query_with_citations(self, query: str) -> AnswerWithCitations:
+        """Query the index with the given query and return the answer along with citations.
+
+        Args:
+            query (str): The query to be answered.
+            citation_chunk_size (int): The size of chunks to use for each citation. Default is 512.
+
+        Returns:
+            AnswerWithCitations: An object containing the answer and citations.
+        """
+        query_engine = CitationQueryEngine.from_args(
+            index=self.index,
+            citation_chunk_size=self.citation_chunk_size,
+        )
+
+        response = query_engine.query(query)
+
+        if hasattr(response, "response"):
+            return AnswerWithCitations(answer=response.response, citations=response.source_nodes)
+        else:
+            raise ValueError(f"Query response of type '{type(response)}' does not contain a response attribute.")
+
+
+# mypy will fail if ChromaDBQueryEngine does not implement RAGQueryEngine protocol
+if TYPE_CHECKING:
+    from ....agentchat.contrib.rag.query_engine import RAGQueryEngine
+
+    def _check_implement_protocol(o: VectorChromaQueryEngine) -> RAGQueryEngine:
+        return o

autogen/agents/experimental/document_agent/docling_doc_ingest_agent.py

@@ -0,0 +1,113 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from pathlib import Path
+from typing import Literal
+
+from .... import ConversableAgent
+from ....agentchat.contrib.rag.query_engine import RAGQueryEngine
+from ....agentchat.group.context_variables import ContextVariables
+from ....agentchat.group.reply_result import ReplyResult
+from ....agentchat.group.targets.transition_target import AgentNameTarget
+from ....doc_utils import export_module
+from ....llm_config import LLMConfig
+from ..document_agent.parser_utils import docling_parse_docs
+from .chroma_query_engine import VectorChromaQueryEngine
+from .document_utils import preprocess_path
+
+__all__ = ["DoclingDocIngestAgent"]
+
+logger = logging.getLogger(__name__)
+
+DOCLING_PARSE_TOOL_NAME = "docling_parse_docs"
+
+DEFAULT_DOCLING_PARSER_PROMPT = f"""
+You are an expert in parsing and understanding text. You can use {DOCLING_PARSE_TOOL_NAME} tool to parse various documents and extract information from them. You can only use the tool once per turn.
+"""
+
+
+@export_module("autogen.agents.experimental")
+class DoclingDocIngestAgent(ConversableAgent):
+    """A DoclingDocIngestAgent is a swarm agent that ingests documents using the docling_parse_docs tool."""
+
+    def __init__(
+        self,
+        name: str | None = None,
+        llm_config: LLMConfig | dict | Literal[False] | None = None,  # type: ignore[type-arg]
+        parsed_docs_path: Path | str | None = None,
+        query_engine: RAGQueryEngine | None = None,
+        return_agent_success: str = "TaskManagerAgent",
+        return_agent_error: str = "ErrorManagerAgent",
+        collection_name: str | None = None,
+    ):
+        """Initialize the DoclingDocIngestAgent.
+
+        Args:
+            name: The name of the DoclingDocIngestAgent.
+            llm_config: The configuration for the LLM.
+            parsed_docs_path: The path where parsed documents will be stored.
+            query_engine: The VectorChromaQueryEngine to use for querying documents.
+            return_agent_success: The agent to return on successful completion of the task.
+            return_agent_error: The agent to return on error.
+            collection_name: The unique name for the Chromadb collection. Set this to a value to reuse a collection. If a query_engine is provided, this will be ignored.
+        """
+        name = name or "DoclingDocIngestAgent"
+
+        parsed_docs_path = parsed_docs_path or Path("./parsed_docs")
+        parsed_docs_path = preprocess_path(str_or_path=parsed_docs_path, mk_path=True)
+
+        self._query_engine = query_engine or VectorChromaQueryEngine(collection_name=collection_name)
+
+        def data_ingest_task(context_variables: ContextVariables) -> ReplyResult:
+            """A tool for Swarm agent to ingests documents using the docling_parse_docs to parse documents to markdown
+            and add them to the docling_query_engine.
+
+            Args:
+            context_variables (ContextVariables): The context variables for the task.
+
+            Returns:
+            ReplyResult: The result of the task.
+            """
+            try:
+                input_file_path = ""
+                tasks = context_variables.get("DocumentsToIngest", [])
+                while tasks:
+                    task = tasks.pop()
+                    input_file_path = task.path_or_url
+                    output_files = docling_parse_docs(
+                        input_file_path=input_file_path, output_dir_path=parsed_docs_path, output_formats=["markdown"]
+                    )
+
+                    # Limit to one output markdown file for now.
+                    if output_files:
+                        output_file = output_files[0]
+                        if output_file.suffix == ".md":
+                            self._query_engine.add_docs(new_doc_paths_or_urls=[output_file])
+
+                    # Keep track of documents ingested
+                    context_variables["DocumentsIngested"].append(input_file_path)
+
+                context_variables["CompletedTaskCount"] += 1
+                logger.info(f"data_ingest_task context_variables: {context_variables.to_dict()}")
+
+            except Exception as e:
+                return ReplyResult(
+                    target=AgentNameTarget(agent_name=return_agent_error),
+                    message=f"Data Ingestion Task Failed, Error {e}: '{input_file_path}'",
+                    context_variables=context_variables,
+                )
+
+            return ReplyResult(
+                target=AgentNameTarget(agent_name=return_agent_success),
+                message=f"Data Ingestion Task Completed for {input_file_path}",
+                context_variables=context_variables,
+            )
+
+        super().__init__(
+            name=name,
+            llm_config=llm_config,
+            functions=[data_ingest_task],
+            system_message=DEFAULT_DOCLING_PARSER_PROMPT,
+        )