Only use third party modules in optional imports not autogen imports in test files

autogen/agentchat/contrib/vectordb/pgvectordb.py

@@ -10,7 +10,6 @@
 from typing import Callable, Optional, Union
 
 import numpy as np
-from sentence_transformers import SentenceTransformer
 
 from ....import_utils import optional_import_block, require_optional_import
 from .base import Document, ItemID, QueryResults, VectorDB
@@ -20,12 +19,13 @@
     import pgvector  # noqa: F401
     import psycopg
     from pgvector.psycopg import register_vector
+    from sentence_transformers import SentenceTransformer
 
 PGVECTOR_MAX_BATCH_SIZE = os.environ.get("PGVECTOR_MAX_BATCH_SIZE", 40000)
 logger = get_logger(__name__)
 
 
-@require_optional_import("psycopg", "retrievechat-pgvector")
+@require_optional_import(["psycopg", "sentence_transformers"], "retrievechat-pgvector")
 class Collection:
     """A Collection object for PGVector.
 
@@ -539,7 +539,7 @@ def create_collection(
         cursor.close()
 
 
-@require_optional_import(["pgvector", "psycopg"], "retrievechat-pgvector")
+@require_optional_import(["pgvector", "psycopg", "sentence_transformers"], "retrievechat-pgvector")
 class PGVectorDB(VectorDB):
     """A vector database that uses PGVector as the backend."""
 

test/agentchat/contrib/capabilities/test_teachable_agent.py

@@ -9,13 +9,15 @@
 import pytest
 
 from autogen import ConversableAgent
+from autogen.agentchat.contrib.capabilities.teachability import Teachability
 from autogen.formatting_utils import colored
 from autogen.import_utils import optional_import_block
 
 from ....conftest import Credentials
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.capabilities.teachability import Teachability
+    import chromadb  # noqa: F401
+
 
 skip = not result.is_successful
 

test/agentchat/contrib/capabilities/test_transforms.py

@@ -106,9 +106,11 @@ def get_messages_with_names_post_filtered() -> list[dict]:
 def get_text_compressors() -> list[TextCompressor]:
     compressors: list[TextCompressor] = [_MockTextCompressor()]
     with optional_import_block() as result:
-        from autogen.agentchat.contrib.capabilities.text_compressors import LLMLingua
+        import llmlingua  # noqa: F401
 
     if result.is_successful:
+        from autogen.agentchat.contrib.capabilities.text_compressors import LLMLingua
+
         compressors.append(LLMLingua())
 
     return compressors

test/agentchat/contrib/capabilities/test_vision_capability.py

@@ -9,13 +9,13 @@
 
 import pytest
 
+from autogen.agentchat.contrib.capabilities.vision_capability import VisionCapability
 from autogen.agentchat.conversable_agent import ConversableAgent
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
     from PIL import Image  # noqa: F401
 
-    from autogen.agentchat.contrib.capabilities.vision_capability import VisionCapability
 
 skip_test = not result.is_successful
 

test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py

@@ -7,16 +7,18 @@
 
 import pytest
 
+from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
+from autogen.agentchat.contrib.graph_rag.neo4j_native_graph_query_engine import (
+    GraphStoreQueryResult,
+    Neo4jNativeGraphQueryEngine,
+)
 from autogen.import_utils import optional_import_block
 
 from ....conftest import reason
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
-    from autogen.agentchat.contrib.graph_rag.neo4j_native_graph_query_engine import (
-        GraphStoreQueryResult,
-        Neo4jNativeGraphQueryEngine,
-    )
+    from neo4j import GraphDatabase  # noqa: F401
+    from neo4j_graphrag.embeddings import Embedder  # noqa: F401
 
 
 skip = not result.is_successful

test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py

@@ -10,16 +10,17 @@
 
 import pytest
 
+from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
+from autogen.agentchat.contrib.graph_rag.neo4j_graph_query_engine import (
+    GraphStoreQueryResult,
+    Neo4jGraphQueryEngine,
+)
 from autogen.import_utils import optional_import_block
 
 from ....conftest import reason
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
-    from autogen.agentchat.contrib.graph_rag.neo4j_graph_query_engine import (
-        GraphStoreQueryResult,
-        Neo4jGraphQueryEngine,
-    )
+    from llama_index.core import PropertyGraphIndex  # noqa: F401
 
 
 skip = not result.is_successful

test/agentchat/contrib/retrievechat/test_pgvector_retrievechat.py

@@ -12,16 +12,18 @@
 from sentence_transformers import SentenceTransformer
 
 from autogen import AssistantAgent
+from autogen.agentchat.contrib.retrieve_user_proxy_agent import (
+    RetrieveUserProxyAgent,
+)
 from autogen.import_utils import optional_import_block
 
 from ....conftest import Credentials
 
 with optional_import_block() as result:
+    import chromadb  # noqa: F401
     import pgvector  # noqa: F401
+    from IPython import get_ipython  # noqa: F401
 
-    from autogen.agentchat.contrib.retrieve_user_proxy_agent import (
-        RetrieveUserProxyAgent,
-    )
 
 skip = not result.is_successful
 

test/agentchat/contrib/retrievechat/test_qdrant_retrievechat.py

@@ -12,6 +12,11 @@
 import pytest
 
 from autogen import AssistantAgent
+from autogen.agentchat.contrib.qdrant_retrieve_user_proxy_agent import (
+    QdrantRetrieveUserProxyAgent,
+    create_qdrant_from_dir,
+    query_qdrant,
+)
 from autogen.import_utils import optional_import_block
 
 from ....conftest import Credentials
@@ -20,11 +25,6 @@
     import fastembed  # noqa: F401
     from qdrant_client import QdrantClient
 
-    from autogen.agentchat.contrib.qdrant_retrieve_user_proxy_agent import (
-        QdrantRetrieveUserProxyAgent,
-        create_qdrant_from_dir,
-        query_qdrant,
-    )
 
 QDRANT_INSTALLED = result.is_successful
 

test/agentchat/contrib/test_img_utils.py

@@ -11,25 +11,25 @@
 import unittest
 from unittest.mock import patch
 
+import numpy as np
 import pytest
 import requests
 
+from autogen.agentchat.contrib.img_utils import (
+    convert_base64_to_data_uri,
+    extract_img_paths,
+    get_image_data,
+    get_pil_image,
+    gpt4v_formatter,
+    llava_formatter,
+    message_formatter_pil_to_b64,
+    num_tokens_from_gpt_image,
+)
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
-    import numpy as np
     from PIL import Image
 
-    from autogen.agentchat.contrib.img_utils import (
-        convert_base64_to_data_uri,
-        extract_img_paths,
-        get_image_data,
-        get_pil_image,
-        gpt4v_formatter,
-        llava_formatter,
-        message_formatter_pil_to_b64,
-        num_tokens_from_gpt_image,
-    )
 
 skip = not result.is_successful
 

test/agentchat/contrib/test_llava.py

@@ -11,12 +11,13 @@
 
 import pytest
 
+from autogen.agentchat.contrib.llava_agent import LLaVAAgent, _llava_call_binary_with_config, llava_call
 from autogen.import_utils import optional_import_block
 
 from ...conftest import MOCK_OPEN_AI_API_KEY
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.llava_agent import LLaVAAgent, _llava_call_binary_with_config, llava_call
+    import replicate  # noqa: F401
 
 skip = not result.is_successful
 

test/agentchat/contrib/test_lmm.py

@@ -12,14 +12,16 @@
 import pytest
 
 import autogen
+from autogen.agentchat.contrib.img_utils import get_pil_image
+from autogen.agentchat.contrib.multimodal_conversable_agent import MultimodalConversableAgent
 from autogen.agentchat.conversable_agent import ConversableAgent
 from autogen.import_utils import optional_import_block
 
 from ...conftest import MOCK_OPEN_AI_API_KEY
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.img_utils import get_pil_image
-    from autogen.agentchat.contrib.multimodal_conversable_agent import MultimodalConversableAgent
+    from PIL import Image  # noqa: F401
+
 
 skip = not result.is_successful
 

test/agentchat/contrib/test_web_surfer.py

@@ -12,6 +12,7 @@
 import pytest
 
 from autogen import UserProxyAgent
+from autogen.agentchat.contrib.web_surfer import WebSurferAgent
 from autogen.import_utils import optional_import_block
 
 from ...conftest import MOCK_OPEN_AI_API_KEY, Credentials
@@ -21,7 +22,12 @@
 BING_QUERY = "Microsoft"
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.web_surfer import WebSurferAgent
+    import markdownify  # noqa: F401
+    import pathvalidate  # noqa: F401
+    import pdfminer  # noqa: F401
+    import requests  # noqa: F401
+    from bs4 import BeautifulSoup  # noqa: F401
+
 
 skip_all = not result.is_successful
 

test/agentchat/contrib/vectordb/test_chromadb.py

@@ -9,6 +9,7 @@
 
 import pytest
 
+from autogen.agentchat.contrib.vectordb.chromadb import ChromaVectorDB
 from autogen.import_utils import optional_import_block
 
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
@@ -18,7 +19,6 @@
     import chromadb.errors
     import sentence_transformers  # noqa: F401
 
-    from autogen.agentchat.contrib.vectordb.chromadb import ChromaVectorDB
 
 skip = not result.is_successful
 

test/agentchat/contrib/vectordb/test_mongodb.py

@@ -12,13 +12,13 @@
 import pytest
 
 from autogen.agentchat.contrib.vectordb.base import Document
+from autogen.agentchat.contrib.vectordb.mongodb import MongoDBAtlasVectorDB
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
     import pymongo  # noqa: F401
     import sentence_transformers  # noqa: F401
 
-    from autogen.agentchat.contrib.vectordb.mongodb import MongoDBAtlasVectorDB
 
 if not result.is_successful:
     # To display warning in pyproject.toml [tool.pytest.ini_options] set log_cli = true

test/agentchat/contrib/vectordb/test_pgvectordb.py

@@ -10,6 +10,7 @@
 
 import pytest
 
+from autogen.agentchat.contrib.vectordb.pgvectordb import PGVectorDB
 from autogen.import_utils import optional_import_block
 
 from ....conftest import reason
@@ -19,7 +20,6 @@
     import psycopg
     import sentence_transformers  # noqa: F401
 
-    from autogen.agentchat.contrib.vectordb.pgvectordb import PGVectorDB
 
 skip = not result.is_successful
 

test/agentchat/contrib/vectordb/test_qdrant.py

@@ -6,19 +6,19 @@
 # SPDX-License-Identifier: MIT
 import os
 import sys
+import uuid
 
 import pytest
 
+from autogen.agentchat.contrib.vectordb.qdrant import QdrantVectorDB
 from autogen.import_utils import optional_import_block
 
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
 
 with optional_import_block() as result:
-    import uuid
-
+    from fastembed import TextEmbedding  # noqa: F401
     from qdrant_client import QdrantClient
 
-    from autogen.agentchat.contrib.vectordb.qdrant import QdrantVectorDB
 
 skip = not result.is_successful
 

test/cache/test_cosmos_db_cache.py

@@ -10,12 +10,13 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
+from autogen.cache.cosmos_db_cache import CosmosDBCache
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
+    from azure.cosmos import CosmosClient  # noqa: F401
     from azure.cosmos.exceptions import CosmosResourceNotFoundError
 
-    from autogen.cache.cosmos_db_cache import CosmosDBCache
 
 skip_test = not result.is_successful
 

test/coding/test_embedded_ipython_code_executor.py

@@ -31,6 +31,12 @@
 
 classes_to_test = []
 with optional_import_block() as result:
+    import websocket  # noqa: F401
+    from jupyter_client import KernelManager  # noqa: F401
+    from jupyter_client.kernelspec import KernelSpecManager  # noqa: F401
+
+    # This has to be here
+    # because autogen.coding.jupyter.helpers is written in a way that it will raise an ImportError
     from autogen.coding.jupyter import (
         DockerJupyterServer,
         EmbeddedIPythonCodeExecutor,

test/oai/test_anthropic.py

@@ -11,11 +11,11 @@
 import pytest
 
 from autogen.import_utils import optional_import_block
+from autogen.oai.anthropic import AnthropicClient, _calculate_cost
 
 with optional_import_block() as result:
     from anthropic.types import Message, TextBlock
 
-    from autogen.oai.anthropic import AnthropicClient, _calculate_cost
 
 skip = not result.is_successful
 

test/oai/test_cohere.py

@@ -11,9 +11,10 @@
 import pytest
 
 from autogen.import_utils import optional_import_block
+from autogen.oai.cohere import CohereClient, calculate_cohere_cost
 
 with optional_import_block() as result:
-    from autogen.oai.cohere import CohereClient, calculate_cohere_cost
+    from cohere import Client as Cohere  # noqa: F401
 
 skip = not result.is_successful