fix: add type safety , validation, thread safety

Author: priyansh4320

autogen/agents/experimental/document_agent/document_agent.py

@@ -162,6 +162,7 @@ def create_summary_agent_prompt(agent: ConversableAgent, messages: list[dict[str
             """Create the summary agent prompt with context information."""
             update_ingested_documents()
 
+            # Safe type casting with defaults
             query_results = cast(list[dict[str, Any]], agent.context_variables.get("QueryResults", []))
             documents_ingested = cast(list[str], agent.context_variables.get("DocumentsIngested", []))
             documents_to_ingest = cast(list[Ingest], agent.context_variables.get("DocumentsToIngest", []))
@@ -209,86 +210,94 @@ def generate_inner_group_chat_reply(
         config: Any = None,
     ) -> tuple[bool, str | dict[str, Any] | None]:
         """Reply function that generates the inner group chat reply for the DocAgent."""
-        # Initialize or reuse context variables
-        if hasattr(self, "_group_chat_context_variables") and self._group_chat_context_variables is not None:
-            context_variables = self._group_chat_context_variables
-            # Reset pending tasks for new run
-            context_variables["DocumentsToIngest"] = []
-        else:
-            context_variables = ContextVariables(
-                data={
-                    "CompletedTaskCount": 0,
-                    "DocumentsToIngest": [],
-                    "DocumentsIngested": self.documents_ingested,
-                    "QueriesToRun": [],
-                    "QueryResults": [],
-                }
+        try:
+            # Initialize or reuse context variables
+            if hasattr(self, "_group_chat_context_variables") and self._group_chat_context_variables is not None:
+                context_variables = self._group_chat_context_variables
+                # Reset pending tasks for new run
+                context_variables["DocumentsToIngest"] = []
+            else:
+                context_variables = ContextVariables(
+                    data={
+                        "CompletedTaskCount": 0,
+                        "DocumentsToIngest": [],
+                        "DocumentsIngested": self.documents_ingested,
+                        "QueriesToRun": [],
+                        "QueryResults": [],
+                    }
+                )
+                self._group_chat_context_variables = context_variables
+
+            if messages and len(messages) > 0:
+                last_message = messages[-1]
+                if (
+                    isinstance(last_message, dict)
+                    and last_message.get("name") == "DocumentTriageAgent"
+                    and "content" in last_message
+                    and isinstance(last_message["content"], str)
+                ):
+                    try:
+                        import json
+
+                        document_task_data = json.loads(last_message["content"])
+
+                        # Extract ingestions and queries
+                        ingestions = [Ingest(**ing) for ing in document_task_data.get("ingestions", [])]
+                        queries = [Query(**q) for q in document_task_data.get("queries", [])]
+
+                        # Update context variables with new tasks
+                        existing_ingestions = context_variables.get("DocumentsToIngest", []) or []
+                        existing_queries = context_variables.get("QueriesToRun", []) or []
+                        documents_ingested = context_variables.get("DocumentsIngested", []) or []
+
+                        # Deduplicate and add new ingestions
+                        for ingestion in ingestions:
+                            if (
+                                ingestion.path_or_url not in [ing.path_or_url for ing in existing_ingestions]
+                                and ingestion.path_or_url not in documents_ingested
+                            ):
+                                existing_ingestions.append(ingestion)
+
+                        # Deduplicate and add new queries
+                        for query in queries:
+                            if query.query not in [q.query for q in existing_queries]:
+                                existing_queries.append(query)
+
+                        context_variables["DocumentsToIngest"] = existing_ingestions
+                        context_variables["QueriesToRun"] = existing_queries
+                        context_variables["TaskInitiated"] = True
+
+                        logger.info(f"Processed triage output: {len(ingestions)} ingestions, {len(queries)} queries")
+
+                    except json.JSONDecodeError as e:
+                        logger.warning(f"Failed to parse triage output JSON: {e}")
+                    except Exception as e:
+                        logger.warning(f"Failed to process triage output: {e}")
+
+            group_chat_agents = [
+                self._triage_agent,
+                self._task_manager_agent,
+                self._summary_agent,
+            ]
+
+            agent_pattern = DefaultPattern(
+                initial_agent=self._triage_agent,
+                agents=group_chat_agents,
+                context_variables=context_variables,
+                group_after_work=TerminateTarget(),
             )
-            self._group_chat_context_variables = context_variables
-
-        if messages and len(messages) > 0:
-            last_message = messages[-1]
-            if (
-                isinstance(last_message, dict)
-                and last_message.get("name") == "DocumentTriageAgent"
-                and "content" in last_message
-            ):
-                try:
-                    import json
-
-                    document_task_data = json.loads(last_message["content"])
-
-                    # Extract ingestions and queries
-                    ingestions = [Ingest(**ing) for ing in document_task_data.get("ingestions", [])]
-                    queries = [Query(**q) for q in document_task_data.get("queries", [])]
-
-                    # Update context variables with new tasks
-                    existing_ingestions = context_variables.get("DocumentsToIngest", []) or []
-                    existing_queries = context_variables.get("QueriesToRun", []) or []
-                    documents_ingested = context_variables.get("DocumentsIngested", []) or []
-
-                    # Deduplicate and add new ingestions
-                    for ingestion in ingestions:
-                        if (
-                            ingestion.path_or_url not in [ing.path_or_url for ing in existing_ingestions]
-                            and ingestion.path_or_url not in documents_ingested
-                        ):
-                            existing_ingestions.append(ingestion)
-
-                    # Deduplicate and add new queries
-                    for query in queries:
-                        if query.query not in [q.query for q in existing_queries]:
-                            existing_queries.append(query)
-
-                    context_variables["DocumentsToIngest"] = existing_ingestions
-                    context_variables["QueriesToRun"] = existing_queries
-                    context_variables["TaskInitiated"] = True
-
-                    logger.info(f"Processed triage output: {len(ingestions)} ingestions, {len(queries)} queries")
-
-                except Exception as e:
-                    logger.warning(f"Failed to process triage output: {e}")
-
-        group_chat_agents = [
-            self._triage_agent,
-            self._task_manager_agent,
-            self._summary_agent,
-        ]
-
-        agent_pattern = DefaultPattern(
-            initial_agent=self._triage_agent,
-            agents=group_chat_agents,
-            context_variables=context_variables,
-            group_after_work=TerminateTarget(),
-        )
 
-        chat_result, context_variables, last_speaker = initiate_group_chat(
-            pattern=agent_pattern,
-            messages=self._get_document_input_message(messages),
-        )
+            chat_result, context_variables, last_speaker = initiate_group_chat(
+                pattern=agent_pattern,
+                messages=self._get_document_input_message(messages),
+            )
+
+            # Always return the final result since we only have summary termination
+            return True, chat_result.summary
 
-        # Always return the final result since we only have summary termination
-        return True, chat_result.summary
+        except Exception as e:
+            logger.error(f"Critical error in DocAgent group chat: {e}")
+            return True, f"Error processing request: {str(e)}"
 
     def _get_document_input_message(self, messages: list[dict[str, Any]] | None) -> str:
         """Gets and validates the input message(s) for the document agent."""

autogen/agents/experimental/document_agent/task_manager.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import logging
+import threading
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Any
@@ -81,6 +82,7 @@ def __init__(
         self.parsed_docs_path = Path(parsed_docs_path) if parsed_docs_path else Path("./parsed_docs")
         self.executor = ThreadPoolExecutor(max_workers=max_workers)
         self._temp_citations_store: dict[str, list[dict[str, str]]] = {}
+        self._context_lock = threading.Lock()
 
         # Initialize RAG engines
         self.rag_engines = self._create_rag_engines(collection_name)
@@ -100,6 +102,27 @@ async def ingest_documents(
             Returns:
                 str: Status message about the ingestion process
             """
+            # Add input validation
+            if not documents_to_ingest:
+                return ReplyResult(
+                    message="No documents provided for ingestion",
+                    context_variables=context_variables,
+                )
+
+            # Validate document paths/URLs
+            valid_documents = []
+            for doc_path in documents_to_ingest:
+                if isinstance(doc_path, str) and doc_path.strip():
+                    valid_documents.append(doc_path.strip())
+                else:
+                    logger.warning(f"Invalid document path: {doc_path}")
+
+            if not valid_documents:
+                return ReplyResult(
+                    message="No valid documents found for ingestion",
+                    context_variables=context_variables,
+                )
+
             # Safely handle context variable initialization
             if "DocumentsToIngest" not in context_variables:
                 context_variables["DocumentsToIngest"] = []
@@ -111,7 +134,7 @@ async def ingest_documents(
                 context_variables["QueriesToRun"] = []
 
             # Add current batch to pending ingestions
-            context_variables["DocumentsToIngest"].append(documents_to_ingest)
+            context_variables["DocumentsToIngest"].append(valid_documents)
 
             try:
                 # Process documents concurrently using ThreadPoolExecutor
@@ -125,7 +148,7 @@ async def ingest_documents(
                         self.rag_config,
                         self.rag_engines,
                     )
-                    for doc_path in documents_to_ingest
+                    for doc_path in valid_documents
                 ]
 
                 # Wait for all documents to be processed
@@ -151,7 +174,7 @@ async def ingest_documents(
                 logger.info("=" * 80)
                 logger.info("TOOL: ingest_documents (CONCURRENT)")
                 logger.info("AGENT: TaskManagerAgent")
-                logger.info(f"DOCUMENTS: {documents_to_ingest}")
+                logger.info(f"DOCUMENTS: {valid_documents}")
                 logger.info(f"SUCCESSFULLY INGESTED: {successfully_ingested}")
                 logger.info("=" * 80)
 
@@ -175,11 +198,11 @@ async def ingest_documents(
                 logger.error("TOOL ERROR: ingest_documents (CONCURRENT)")
                 logger.error("AGENT: TaskManagerAgent")
                 logger.error(f"ERROR: {e}")
-                logger.error(f"DOCUMENTS: {documents_to_ingest}")
+                logger.error(f"DOCUMENTS: {valid_documents}")
                 logger.error("=" * 80)
 
                 # Preserve failed documents for retry
-                context_variables["DocumentsToIngest"] = [documents_to_ingest]
+                context_variables["DocumentsToIngest"] = [valid_documents]
                 return ReplyResult(
                     message=f"Documents ingestion failed: {e}",
                     context_variables=context_variables,
@@ -198,6 +221,11 @@ async def execute_query(queries_to_run: list[str], context_variables: ContextVar
             if not queries_to_run:
                 return "No queries to run"
 
+            # Validate queries
+            valid_queries = [q.strip() for q in queries_to_run if isinstance(q, str) and q.strip()]
+            if not valid_queries:
+                return "No valid queries provided"
+
             # Safely handle context variable initialization
             if "QueriesToRun" not in context_variables:
                 context_variables["QueriesToRun"] = []
@@ -209,7 +237,7 @@ async def execute_query(queries_to_run: list[str], context_variables: ContextVar
                 context_variables["Citations"] = []
 
             # Add current batch to pending queries
-            context_variables["QueriesToRun"].append(queries_to_run)
+            context_variables["QueriesToRun"].append(valid_queries)
 
             try:
                 # Clear temporary citations store before processing
@@ -219,7 +247,7 @@ async def execute_query(queries_to_run: list[str], context_variables: ContextVar
                 loop = asyncio.get_event_loop()
                 futures = [
                     loop.run_in_executor(self.executor, execute_single_query, query, self.rag_config, self.rag_engines)
-                    for query in queries_to_run
+                    for query in valid_queries
                 ]
 
                 # Wait for all queries to be processed
@@ -252,15 +280,15 @@ async def execute_query(queries_to_run: list[str], context_variables: ContextVar
                 logger.info("=" * 80)
                 logger.info("TOOL: execute_query (CONCURRENT)")
                 logger.info("AGENT: TaskManagerAgent")
-                logger.info(f"QUERIES: {queries_to_run}")
+                logger.info(f"QUERIES: {valid_queries}")
                 logger.info("=" * 80)
 
                 # Update context variables
                 context_variables["QueriesToRun"].pop(0)  # Remove processed batch
                 context_variables["CompletedTaskCount"] += 1
 
                 # Store query results with citations
-                query_result = {"query": queries_to_run, "answer": answers, "citations": all_citations}
+                query_result = {"query": valid_queries, "answer": answers, "citations": all_citations}
                 context_variables["QueryResults"].append(query_result)
                 # Clear temporary citations store after processing
                 self._temp_citations_store = {}
@@ -271,13 +299,13 @@ async def execute_query(queries_to_run: list[str], context_variables: ContextVar
                 )
 
             except Exception as e:
-                error_msg = f"Query failed for queries '{queries_to_run}': {str(e)}"
+                error_msg = f"Query failed for queries '{valid_queries}': {str(e)}"
 
                 # Enhanced error logging
                 logger.error("=" * 80)
                 logger.error("TOOL ERROR: execute_query (CONCURRENT)")
                 logger.error("AGENT: TaskManagerAgent")
-                logger.error(f"QUERIES: {queries_to_run}")
+                logger.error(f"QUERIES: {valid_queries}")
                 logger.error(f"ERROR: {e}")
                 logger.error("=" * 80)
 
@@ -298,8 +326,11 @@ async def execute_query(queries_to_run: list[str], context_variables: ContextVar
 
     def __del__(self) -> None:
         """Clean up the ThreadPoolExecutor when the agent is destroyed."""
-        if hasattr(self, "executor"):
-            self.executor.shutdown(wait=True)
+        if hasattr(self, "executor") and self.executor is not None:
+            try:
+                self.executor.shutdown(wait=False)  # Don't block in destructor
+            except Exception as e:
+                logger.warning(f"Error shutting down executor: {e}")
 
     def _create_rag_engines(self, collection_name: str | None = None) -> dict[str, Any]:
         """Create RAG engines based on rag_config."""
@@ -335,3 +366,8 @@ def _create_neo4j_engine(self, config: dict[str, Any]) -> Any:
         except ImportError as e:
             logger.warning(f"Neo4j dependencies not available: {e}. Skipping graph engine.")
             return None
+
+    def _safe_context_update(self, context_variables: ContextVariables, key: str, value: Any) -> None:
+        """Thread-safe context variable update."""
+        with self._context_lock:
+            context_variables[key] = value

autogen/agents/experimental/document_agent/task_manager_utils.py

@@ -28,11 +28,12 @@ def extract_text_from_pdf(doc_path: str) -> list[dict[str, str]]:
         text = ""
         # Save the PDF to a temporary file
         with tempfile.TemporaryDirectory() as temp_dir:
-            with open(temp_dir + "temp.pdf", "wb") as f:
+            temp_pdf_path = Path(temp_dir) / "temp.pdf"
+            with open(temp_pdf_path, "wb") as f:
                 f.write(response.content)
 
             # Open the PDF
-            with fitz.open(temp_dir + "temp.pdf") as doc:
+            with fitz.open(str(temp_pdf_path)) as doc:
                 # Read and extract text from each page
                 for page in doc:
                     text += page.get_text()
@@ -72,6 +73,9 @@ def compress_and_save_text(text: str, input_path: str, parsed_docs_path: Path) -
     text_compressor = TextMessageCompressor(text_compressor=llm_lingua)
     compressed_text = text_compressor.apply_transform([{"content": text}])
 
+    if not compressed_text or not compressed_text[0].get("content"):
+        raise ValueError("Text compression failed or returned empty result")
+
     # Create a markdown file with the extracted text
     output_file = parsed_docs_path / f"{Path(input_path).stem}.md"
     parsed_docs_path.mkdir(parents=True, exist_ok=True)
@@ -138,7 +142,7 @@ def process_single_document(
 
         if is_pdf:
             # Handle PDF with PyMuPDF
-            print("PDF found using PyMuPDF")
+            logger.info("PDF found, using PyMuPDF for extraction")
             if urllib3.util.url.parse_url(input_file_path).scheme:
                 # Download the PDF
                 response = requests.get(input_file_path)