test: task manager utils

Author: priyansh4320

autogen/agents/experimental/document_agent/document_agent.py

@@ -121,7 +121,7 @@ def __init__(
             query_engine = VectorChromaCitationQueryEngine(
                 collection_name=collection_name, enable_query_citations=True, citation_chunk_size=citation_chunk_size
             )
-        else:
+        elif query_engine is None:
             query_engine = VectorChromaQueryEngine(collection_name=collection_name)
 
         super().__init__(

pyproject.toml

@@ -144,8 +144,6 @@ rag = [
     "llama-index-embeddings-openai>=0.3,<0.6",
     "llama-index-llms-openai>=0.4,<0.6",
     "requests>=2.32.3,<3",
-    "pymupdf>=1.24.4",
-    "llmlingua>=0.2.2"
 ]
 
 
@@ -309,6 +307,8 @@ docs = [
     "termcolor==3.1.0",
     "nbclient==0.10.2",
     "mcp>=1.11.0",
+    "pymupdf>=1.24.4",
+    "llmlingua>=0.2.2"
 ]
 
 types = [

test/agents/experimental/document_agent/test_task_manager.py

@@ -215,7 +215,7 @@ def test_process_single_document_success(
         llm_config = credentials_gpt_4o_mini.llm_config
 
         with (
-            patch("autogen.agents.experimental.document_agent.task_manager.docling_parse_docs") as mock_parse,
+            patch("autogen.agents.experimental.document_agent.task_manager_utils.docling_parse_docs") as mock_parse,
             patch("autogen.agents.experimental.document_agent.task_manager.ThreadPoolExecutor"),
         ):
             # Mock successful document parsing
@@ -238,7 +238,7 @@ def test_process_single_document_parsing_failure(
         llm_config = credentials_gpt_4o_mini.llm_config
 
         with (
-            patch("autogen.agents.experimental.document_agent.task_manager.docling_parse_docs") as mock_parse,
+            patch("autogen.agents.experimental.document_agent.task_manager_utils.docling_parse_docs") as mock_parse,
             patch("autogen.agents.experimental.document_agent.task_manager.ThreadPoolExecutor"),
         ):
             # Mock document parsing failure
@@ -259,7 +259,7 @@ def test_process_single_document_no_output_files(
         llm_config = credentials_gpt_4o_mini.llm_config
 
         with (
-            patch("autogen.agents.experimental.document_agent.task_manager.docling_parse_docs") as mock_parse,
+            patch("autogen.agents.experimental.document_agent.task_manager_utils.docling_parse_docs") as mock_parse,
             patch("autogen.agents.experimental.document_agent.task_manager.ThreadPoolExecutor"),
         ):
             # Mock no output files
@@ -280,7 +280,7 @@ def test_process_single_document_non_markdown_output(
         llm_config = credentials_gpt_4o_mini.llm_config
 
         with (
-            patch("autogen.agents.experimental.document_agent.task_manager.docling_parse_docs") as mock_parse,
+            patch("autogen.agents.experimental.document_agent.task_manager_utils.docling_parse_docs") as mock_parse,
             patch("autogen.agents.experimental.document_agent.task_manager.ThreadPoolExecutor"),
         ):
             # Mock non-markdown output file