Merge pull request #26 from semiotic-ai/ip

ip

Author: denverbaumgartner

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,

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))