Merge pull request microsoft#37 from prateejain-linked/users/gbarrnsn…

…chez/add_graphdb_url_parameters

Add graphdb parameters for local emulator support

common/graph_db_client.py

@@ -23,7 +23,7 @@ def __init__(self,graph_db_params: GraphDBConfig|None,context_id: str|None):
         #    credential = ManagedIdentityCredential(client_id="295ce65c-28c6-4763-be6f-a5eb36c3ceb3")
         #    token = credential.get_token(COSMOS_DB_SCOPE)
         self._client=client.Client(
-            url=f"wss://{graph_db_params.account_name}.gremlin.cosmos.azure.com:443/",
+            url=f"{graph_db_params.gremlin_url}",
             traversal_source="g",
             username=self.username_prefix+"-contextid-"+context_id,
             password=token,

graphrag/config/create_graphrag_config.py

@@ -560,6 +560,8 @@ def hydrate_parallelization_params(
                 account_key=reader.str("account_key") or None,
                 username=reader.str("username") or None,
                 enabled=reader.bool("enabled") or False,
+                cosmos_url=reader.str("cosmos_url") or None,
+                gremlin_url=reader.str("gremlin_url") or None,
             )
 
         encoding_model = reader.str(Fragment.encoding_model) or defs.ENCODING_MODEL

graphrag/config/models/graphdb_config.py

@@ -24,4 +24,14 @@ class GraphDBConfig(BaseModel):
     enabled: bool = Field(
         description="Flag to enable querying into graphdb",
         default=False
+    )
+
+    cosmos_url: str|None = Field(
+        description="Cosmos account url",
+        default=None,
+    )
+
+    gremlin_url: str|None = Field(
+        description="Gremlin db url",
+        default=None,
     )

graphrag/index/context_switch/contextSwitcher.py

@@ -199,7 +199,7 @@ def _read_config_parameters(root: str, config: str | None):
         final_covariates = pd.DataFrame()
         if config.graphdb.enabled:
             cosmos_client = CosmosClient(
-                f"https://{config.graphdb.account_name}.documents.azure.com:443/",
+                f"{config.graphdb.cosmos_url}",
                 f"{config.graphdb.account_key}",
             )
             database_name = config.graphdb.username.split("/")[2]
@@ -252,6 +252,9 @@ def _read_config_parameters(root: str, config: str | None):
             if config.graphdb.enabled:
                 graph_db_client.write_vertices(final_entities)
                 graph_db_client.write_edges(final_relationships)
+
+        if config.graphdb.enabled:
+            graph_db_client._client.close()
 
     def deactivate(self):
         """DeActivate the context."""

graphrag/index/init_content.py

@@ -167,6 +167,8 @@
     account_key: ''
     username: ''
     enabled: false
+    cosmos_url: ''
+    gremlin_url: ''
 """
 
 INIT_DOTENV = """

graphrag/query/input/retrieval/relationships.py

@@ -3,6 +3,7 @@
 
 """Util functions to retrieve relationships from a collection."""
 
+import time
 from typing import Any, cast
 
 import pandas as pd
@@ -19,6 +20,8 @@ def get_relationships_from_graphdb(query:str,selected_entity_names:list[str],gra
             "prop_selected_entity_names": selected_entity_names,
         }
     )
+    time.sleep(5)
+    print(graphdb_client.result_to_df(relationships_result))
     return read_relationships(
         graphdb_client.result_to_df(relationships_result),
         short_id_col="human_readable_id"