From e1741e0880eecc8db08d9c5ce0b5c785016d3dfb Mon Sep 17 00:00:00 2001 From: denver Date: Mon, 17 Mar 2025 13:44:12 -0500 Subject: [PATCH 1/2] refactor: switch api_key -> logging_id for module tracing --- .../graphdoc/modules/doc_generator_module.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/graphdoc/graphdoc/modules/doc_generator_module.py b/graphdoc/graphdoc/modules/doc_generator_module.py index 6a73c28..9ea8bb8 100644 --- a/graphdoc/graphdoc/modules/doc_generator_module.py +++ b/graphdoc/graphdoc/modules/doc_generator_module.py @@ -279,13 +279,21 @@ def document_full_schema( trace: bool = False, client: Optional[mlflow.MlflowClient] = None, expirement_name: Optional[str] = None, - api_key: Optional[str] = None, + logging_id: Optional[str] = None, ) -> dspy.Prediction: """Given a database schema, parse out the underlying components and document on a per-component basis. :param database_schema: The database schema to generate documentation for. :type database_schema: str + :param trace: Whether to trace the generation. + :type trace: bool + :param client: The mlflow client. + :type client: mlflow.MlflowClient + :param expirement_name: The name of the experiment. + :type expirement_name: str + :param logging_id: The id to use for logging. Maps back to the user request. + :type logging_id: str :return: The generated documentation. :rtype: dspy.Prediction @@ -296,8 +304,8 @@ def document_full_schema( raise ValueError("client must be provided if trace is True") if expirement_name is None: raise ValueError("expirement_name must be provided if trace is True") - if api_key is None: - raise ValueError("api_key must be provided if trace is True") + if logging_id is None: + raise ValueError("logging_id must be provided if trace is True") # check that the graphql is valid try: @@ -323,7 +331,7 @@ def document_full_schema( # TODO: we should have better type handling, but we check at the top trace_name="document_full_schema", inputs={"database_schema": database_schema}, - attributes={"api_key": api_key}, + attributes={"logging_id": logging_id}, ) log.info("created trace: " + str(root_trace)) From e1979ab1325bf38bf0e04451cfea7837e79473cb Mon Sep 17 00:00:00 2001 From: denver Date: Mon, 17 Mar 2025 13:44:41 -0500 Subject: [PATCH 2/2] refactor: update doc gen eval to use logging_id in document_full_schema --- graphdoc/graphdoc/eval/doc_generator_eval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphdoc/graphdoc/eval/doc_generator_eval.py b/graphdoc/graphdoc/eval/doc_generator_eval.py index 76b3a3a..baa3c43 100644 --- a/graphdoc/graphdoc/eval/doc_generator_eval.py +++ b/graphdoc/graphdoc/eval/doc_generator_eval.py @@ -61,7 +61,7 @@ def forward(self, database_schema: str) -> dict[str, Any]: trace=True, client=self.mlflow_helper.mlflow_client, expirement_name=self.mlflow_experiment_name, - api_key="temp", + logging_id="temp", ) # TODO: let's decide if this is how we want to handle this in the future. # Alternatively, we could return the documented schema from forward,