Move third party non default module imports to optional block in auto…

…gen/agentchat/contrib/graph_rag files

autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py

@@ -4,18 +4,22 @@
 
 import os
 import warnings
+from typing import Optional
 
-from falkordb import FalkorDB, Graph
-from graphrag_sdk import KnowledgeGraph, Source
-from graphrag_sdk.model_config import KnowledgeGraphModelConfig
-from graphrag_sdk.models import GenerativeModel
-from graphrag_sdk.models.openai import OpenAiGenerativeModel
-from graphrag_sdk.ontology import Ontology
-
+from ....import_utils import optional_import_block, require_optional_import
 from .document import Document
 from .graph_query_engine import GraphStoreQueryResult
 
+with optional_import_block():
+    from falkordb import FalkorDB, Graph
+    from graphrag_sdk import KnowledgeGraph, Source
+    from graphrag_sdk.model_config import KnowledgeGraphModelConfig
+    from graphrag_sdk.models import GenerativeModel
+    from graphrag_sdk.models.openai import OpenAiGenerativeModel
+    from graphrag_sdk.ontology import Ontology
+
 
+@require_optional_import(["falkordb", "graphrag_sdk"], "graph-rag-falkor-db")
 class FalkorGraphQueryEngine:
     """This is a wrapper for FalkorDB KnowledgeGraph."""
 
@@ -26,8 +30,8 @@ def __init__(
         port: int = 6379,
         username: str | None = None,
         password: str | None = None,
-        model: GenerativeModel = OpenAiGenerativeModel("gpt-4o"),
-        ontology: Ontology | None = None,
+        model: Optional["GenerativeModel"] = None,
+        ontology: "Ontology" | None = None,
     ):
         """Initialize a FalkorDB knowledge graph.
         Please also refer to https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/kg.py
@@ -50,7 +54,7 @@ def __init__(
         self.port = port
         self.username = username
         self.password = password
-        self.model = model
+        self.model = model or OpenAiGenerativeModel("gpt-4o")
         self.model_config = KnowledgeGraphModelConfig.with_model(model)
         self.ontology = ontology
         self.knowledge_graph = None
@@ -149,17 +153,17 @@ def delete(self) -> bool:
             self.falkordb.select_graph(self.ontology_table_name).delete()
         return True
 
-    def __get_ontology_storage_graph(self) -> Graph:
+    def __get_ontology_storage_graph(self) -> "Graph":
         return self.falkordb.select_graph(self.ontology_table_name)
 
-    def _save_ontology_to_db(self, ontology: Ontology):
+    def _save_ontology_to_db(self, ontology: "Ontology"):
         """Save graph ontology to a separate table with {graph_name}_ontology"""
         if self.ontology_table_name in self.falkordb.list_graphs():
             raise ValueError(f"Knowledge graph {self.name} is already created.")
         graph = self.__get_ontology_storage_graph()
         ontology.save_to_graph(graph)
 
-    def _load_ontology_from_db(self) -> Ontology:
+    def _load_ontology_from_db(self) -> "Ontology":
         if self.ontology_table_name not in self.falkordb.list_graphs():
             raise ValueError(f"Knowledge graph {self.name} has not been created.")
         graph = self.__get_ontology_storage_graph()

autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py

@@ -4,8 +4,7 @@
 
 from typing import Any, Optional, Union
 
-from autogen import Agent, ConversableAgent, UserProxyAgent
-
+from .... import Agent, ConversableAgent, UserProxyAgent
 from .falkor_graph_query_engine import FalkorGraphQueryEngine
 from .graph_query_engine import GraphStoreQueryResult
 from .graph_rag_capability import GraphRagCapability

autogen/agentchat/contrib/graph_rag/graph_rag_capability.py

@@ -4,9 +4,8 @@
 #
 # Portions derived from https://github.com/microsoft/autogen are under the MIT License.
 # SPDX-License-Identifier: MIT
-from autogen.agentchat.contrib.capabilities.agent_capability import AgentCapability
-from autogen.agentchat.conversable_agent import ConversableAgent
-
+from ...conversable_agent import ConversableAgent
+from ..capabilities.agent_capability import AgentCapability
 from .graph_query_engine import GraphQueryEngine
 
 

autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py

@@ -2,26 +2,35 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 import os
-from typing import Optional, TypeAlias, Union
-
-from llama_index.core import PropertyGraphIndex, SimpleDirectoryReader
-from llama_index.core.base.embeddings.base import BaseEmbedding
-from llama_index.core.indices.property_graph import (
-    DynamicLLMPathExtractor,
-    SchemaLLMPathExtractor,
-)
-from llama_index.core.indices.property_graph.transformations.schema_llm import Triple
-from llama_index.core.llms import LLM
-from llama_index.core.readers.json import JSONReader
-from llama_index.core.schema import Document as LlamaDocument
-from llama_index.embeddings.openai import OpenAIEmbedding
-from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore
-from llama_index.llms.openai import OpenAI
+import sys
+from typing import Optional, Union
 
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+from ....import_utils import optional_import_block, require_optional_import
 from .document import Document, DocumentType
 from .graph_query_engine import GraphQueryEngine, GraphStoreQueryResult
 
-
+with optional_import_block():
+    from llama_index.core import PropertyGraphIndex, SimpleDirectoryReader
+    from llama_index.core.base.embeddings.base import BaseEmbedding
+    from llama_index.core.indices.property_graph import (
+        DynamicLLMPathExtractor,
+        SchemaLLMPathExtractor,
+    )
+    from llama_index.core.indices.property_graph.transformations.schema_llm import Triple
+    from llama_index.core.llms import LLM
+    from llama_index.core.readers.json import JSONReader
+    from llama_index.core.schema import Document as LlamaDocument
+    from llama_index.embeddings.openai import OpenAIEmbedding
+    from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore
+    from llama_index.llms.openai import OpenAI
+
+
+@require_optional_import("llama_index", "neo4j")
 class Neo4jGraphQueryEngine(GraphQueryEngine):
     """This class serves as a wrapper for a property graph query engine backed by LlamaIndex and Neo4j,
     facilitating the creating, connecting, updating, and querying of LlamaIndex property graphs.
@@ -49,11 +58,11 @@ def __init__(
         database: str = "neo4j",
         username: str = "neo4j",
         password: str = "neo4j",
-        llm: LLM = OpenAI(model="gpt-4o", temperature=0.0),
-        embedding: BaseEmbedding = OpenAIEmbedding(model_name="text-embedding-3-small"),
-        entities: Optional[TypeAlias] = None,
-        relations: Optional[TypeAlias] = None,
-        schema: Optional[Union[dict[str, str], list[Triple]]] = None,
+        llm: Optional["LLM"] = None,
+        embedding: Optional["BaseEmbedding"] = None,
+        entities: Optional["TypeAlias"] = None,
+        relations: Optional["TypeAlias"] = None,
+        schema: Optional[Union[dict[str, str], list["Triple"]]] = None,
         strict: Optional[bool] = False,
     ):
         """Initialize a Neo4j Property graph.
@@ -78,8 +87,8 @@ def __init__(
         self.database = database
         self.username = username
         self.password = password
-        self.llm = llm
-        self.embedding = embedding
+        self.llm = llm or OpenAI(model="gpt-4o", temperature=0.0)
+        self.embedding = embedding or OpenAIEmbedding(model_name="text-embedding-3-small")
         self.entities = entities
         self.relations = relations
         self.schema = schema
@@ -188,7 +197,7 @@ def _clear(self) -> None:
         with self.graph_store._driver.session() as session:
             session.run("MATCH (n) DETACH DELETE n;")
 
-    def _load_doc(self, input_doc: list[Document]) -> list[LlamaDocument]:
+    def _load_doc(self, input_doc: list[Document]) -> list["LlamaDocument"]:
         """Load documents from the input files. Currently support the following file types:
         .csv - comma-separated values
         .docx - Microsoft Word

autogen/agentchat/contrib/graph_rag/neo4j_graph_rag_capability.py

@@ -4,8 +4,7 @@
 
 from typing import Any, Optional, Union
 
-from autogen import Agent, ConversableAgent, UserProxyAgent
-
+from .... import Agent, ConversableAgent, UserProxyAgent
 from .graph_query_engine import GraphStoreQueryResult
 from .graph_rag_capability import GraphRagCapability
 from .neo4j_graph_query_engine import Neo4jGraphQueryEngine

autogen/agentchat/contrib/graph_rag/neo4j_native_graph_query_engine.py

@@ -6,23 +6,26 @@
 import logging
 from typing import List, Optional, Union
 
-from neo4j import GraphDatabase
-from neo4j_graphrag.embeddings import Embedder, OpenAIEmbeddings
-from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
-from neo4j_graphrag.generation import GraphRAG
-from neo4j_graphrag.indexes import create_vector_index
-from neo4j_graphrag.llm.openai_llm import LLMInterface, OpenAILLM
-from neo4j_graphrag.retrievers import VectorRetriever
-
+from ....import_utils import optional_import_block, require_optional_import
 from .document import Document, DocumentType
 from .graph_query_engine import GraphQueryEngine, GraphStoreQueryResult
 
+with optional_import_block():
+    from neo4j import GraphDatabase
+    from neo4j_graphrag.embeddings import Embedder, OpenAIEmbeddings
+    from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
+    from neo4j_graphrag.generation import GraphRAG
+    from neo4j_graphrag.indexes import create_vector_index
+    from neo4j_graphrag.llm.openai_llm import LLMInterface, OpenAILLM
+    from neo4j_graphrag.retrievers import VectorRetriever
+
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logging.getLogger("httpx").setLevel(logging.WARNING)
 logger = logging.getLogger(__name__)
 
 
+@require_optional_import(["neo4j", "neo4j_graphrag"], "neo4j")
 class Neo4jNativeGraphQueryEngine(GraphQueryEngine):
     """A graph query engine implemented using the Neo4j GraphRAG SDK.
     Provides functionality to initialize a knowledge graph,
@@ -35,13 +38,10 @@ def __init__(
         port: int = 7687,
         username: str = "neo4j",
         password: str = "password",
-        embeddings: Optional[Embedder] = OpenAIEmbeddings(model="text-embedding-3-large"),
+        embeddings: Optional["Embedder"] = None,
         embedding_dimension: Optional[int] = 3072,
-        llm: Optional[LLMInterface] = OpenAILLM(
-            model_name="gpt-4o",
-            model_params={"response_format": {"type": "json_object"}, "temperature": 0},
-        ),
-        query_llm: Optional[LLMInterface] = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0}),
+        llm: Optional["LLMInterface"] = None,
+        query_llm: Optional["LLMInterface"] = None,
         entities: Optional[List[str]] = None,
         relations: Optional[List[str]] = None,
         potential_schema: Optional[List[tuple[str, str, str]]] = None,
@@ -64,10 +64,13 @@ def __init__(
         """
         self.uri = f"{host}:{port}"
         self.driver = GraphDatabase.driver(self.uri, auth=(username, password))
-        self.embeddings = embeddings
+        self.embeddings = embeddings or OpenAIEmbeddings(model="text-embedding-3-large")
         self.embedding_dimension = embedding_dimension
-        self.llm = llm
-        self.query_llm = query_llm
+        self.llm = llm or OpenAILLM(
+            model_name="gpt-4o",
+            model_params={"response_format": {"type": "json_object"}, "temperature": 0},
+        )
+        self.query_llm = query_llm or OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0})
         self.entities = entities
         self.relations = relations
         self.potential_schema = potential_schema

autogen/agentchat/contrib/graph_rag/neo4j_native_graph_rag_capability.py

@@ -4,8 +4,7 @@
 
 from typing import Any, Optional, Union
 
-from autogen import Agent, ConversableAgent
-
+from .... import Agent, ConversableAgent
 from .graph_query_engine import GraphStoreQueryResult
 from .graph_rag_capability import GraphRagCapability
 from .neo4j_native_graph_query_engine import Neo4jNativeGraphQueryEngine