diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py index 37eca45c0..a398af343 100644 --- a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +++ b/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() diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py index 4b09a2a58..2c7cb3b79 100644 --- a/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py +++ b/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 diff --git a/autogen/agentchat/contrib/graph_rag/graph_rag_capability.py b/autogen/agentchat/contrib/graph_rag/graph_rag_capability.py index 70823b78b..129a0417a 100644 --- a/autogen/agentchat/contrib/graph_rag/graph_rag_capability.py +++ b/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 diff --git a/autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py index afba87847..b86dbf287 100644 --- a/autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py +++ b/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 diff --git a/autogen/agentchat/contrib/graph_rag/neo4j_graph_rag_capability.py b/autogen/agentchat/contrib/graph_rag/neo4j_graph_rag_capability.py index 9362cb8fe..50a379b17 100644 --- a/autogen/agentchat/contrib/graph_rag/neo4j_graph_rag_capability.py +++ b/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 diff --git a/autogen/agentchat/contrib/graph_rag/neo4j_native_graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/neo4j_native_graph_query_engine.py index 7b85b0c2d..8f5a6921c 100644 --- a/autogen/agentchat/contrib/graph_rag/neo4j_native_graph_query_engine.py +++ b/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 diff --git a/autogen/agentchat/contrib/graph_rag/neo4j_native_graph_rag_capability.py b/autogen/agentchat/contrib/graph_rag/neo4j_native_graph_rag_capability.py index 31175f0a6..a91f14f69 100644 --- a/autogen/agentchat/contrib/graph_rag/neo4j_native_graph_rag_capability.py +++ b/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