graphdoc package

class graphdoc.DocGeneratorModule(prompt: DocGeneratorPrompt | SinglePrompt, retry: bool = True, retry_limit: int = 1, rating_threshold: int = 3, fill_empty_descriptions: bool = True, token_tracker: TokenTracker | None = None)

Bases: Module

__init__(prompt: DocGeneratorPrompt | SinglePrompt, retry: bool = True, retry_limit: int = 1, rating_threshold: int = 3, fill_empty_descriptions: bool = True, token_tracker: TokenTracker | None = None) → None

Initialize the DocGeneratorModule. A module for generating documentation for a given GraphQL schema. Schemas are decomposed and individually used to generate documentation, with a quality check after each generation.

signature fields are:

• database_schema: str = dspy.InputField()

• documented_schema: str = dspy.OutputField()

Parameters:

• prompt (DocGeneratorPrompt) – The prompt to use for generating documentation.

• retry (bool) – Whether to retry the generation if the quality check fails.

• retry_limit (int) – The maximum number of retries.

• rating_threshold (int) – The minimum rating for a generated document to be considered valid.

• fill_empty_descriptions (bool) – Whether to fill empty descriptions with generated documentation.

_retry_by_rating(database_schema: str) → str

Retry the generation if the quality check fails. Rating threshold is determined at initialization.

Parameters:

database_schema (str) – The database schema to generate documentation for.

Returns:

The generated documentation.

Return type:

str

_predict(database_schema: str) → Prediction

Given a database schema, generate a documented schema. Performs the following steps:

• Check that the graphql is valid

• Fill the empty descriptions (if fill_empty_descriptions is True)

• Generate the documentation

• Check that the generated schema is valid

• Check that the generated schema matches the original schema

Parameters:

database_schema (str) – The database schema to generate documentation for.

Returns:

The generated documentation.

Return type:

dspy.Prediction

forward(database_schema: str) → Prediction

Given a database schema, generate a documented schema. If retry is True, the generation will be retried if the quality check fails.

Parameters:

database_schema (str) – The database schema to generate documentation for.

Returns:

The generated documentation.

Return type:

dspy.Prediction

document_full_schema(database_schema: str, trace: bool = False, client: MlflowClient | None = None, expirement_name: str | None = None, logging_id: str | None = None) → Prediction

Given a database schema, parse out the underlying components and document on a per-component basis.

Parameters:

• database_schema (str) – The database schema to generate documentation for.

• trace (bool) – Whether to trace the generation.

• client (mlflow.MlflowClient) – The mlflow client.

• expirement_name (str) – The name of the experiment.

• logging_id (str) – The id to use for logging. Maps back to the user request.

Returns:

The generated documentation.

Return type:

dspy.Prediction

class graphdoc.TokenTracker

Bases: object

A class to track the number of tokens used.

clear()

Clear the token tracker.

stats()

Get the stats of the token tracker.

global_token_callback(kwargs, response, start_time, end_time, **callback_kwargs)

A global callback to track the number of tokens used.

Intended to be used with the litellm ModelResponse object.

class graphdoc.DocGeneratorEvaluator(generator: DocGeneratorModule | Module | Any, evaluator: DocQualityPrompt | SinglePrompt | Any, evalset: List[Example] | Any, mlflow_helper: MlflowDataHelper, mlflow_experiment_name: str = 'doc_generator_eval', generator_prediction_field: str = 'documented_schema', evaluator_prediction_field: str = 'rating', readable_value: int = 25)

Bases: Module

__init__(generator: DocGeneratorModule | Module | Any, evaluator: DocQualityPrompt | SinglePrompt | Any, evalset: List[Example] | Any, mlflow_helper: MlflowDataHelper, mlflow_experiment_name: str = 'doc_generator_eval', generator_prediction_field: str = 'documented_schema', evaluator_prediction_field: str = 'rating', readable_value: int = 25)

A simple module for evaluating the quality of generated documentation. We will make this extensible to include more complex evaluation metrics in the future.

Important: we assume that the rating values returned by the evaluator are [1, 2, 3, 4]. We will make this more flexible in the future.

forward(database_schema: str) → dict[str, Any]

Takes a database schema, documents it, and then evaluates each component and the aggregate.

evaluate()

Batches the evaluation set and logs the results to mlflow.

class graphdoc.DocGeneratorTrainer(prompt: DocGeneratorPrompt, optimizer_type: str, optimizer_kwargs: Dict[str, Any], mlflow_model_name: str, mlflow_experiment_name: str, mlflow_tracking_uri: str, trainset: List[Example], evalset: List[Example])

Bases: SinglePromptTrainer

__init__(prompt: DocGeneratorPrompt, optimizer_type: str, optimizer_kwargs: Dict[str, Any], mlflow_model_name: str, mlflow_experiment_name: str, mlflow_tracking_uri: str, trainset: List[Example], evalset: List[Example])

Initialize the DocGeneratorTrainer.

Parameters:

• prompt (DocGeneratorPrompt) – The prompt to train.

• optimizer_type (str) – The type of optimizer to use.

• optimizer_kwargs (Dict[str, Any]) – The keyword arguments for the optimizer.

• mlflow_model_name (str) – The name of the model in mlflow.

• mlflow_experiment_name (str) – The name of the experiment in mlflow.

• mlflow_tracking_uri (str) – The uri of the mlflow tracking server.

• trainset (List[dspy.Example]) – The training set.

• evalset (List[dspy.Example]) – The evaluation set.

_calculate_average_score(evaluation: dict) → float

Given a dictionary of evaluation results, calculate the average score.

Parameters:

evaluation (Dict[str, Any]) – The evaluation results.

Returns:

The average score.

Return type:

float

evaluation_metrics(base_evaluation: Dict[str, Any], optimized_evaluation: Dict[str, Any]) → None

Log evaluation metrics to mlflow.

Parameters:

• base_evaluation (Dict[str, Any]) – The evaluation metrics of the base model.

• optimized_evaluation (Dict[str, Any]) – The evaluation metrics of the optimized model.

evaluate_training(base_model, optimized_model) → Tuple[Dict[str, Any], Dict[str, Any]]

Evaluate the training of the model. Comparing the base and optimized models.

Parameters:

• base_model (Any) – The base model.

• optimized_model (Any) – The optimized model.

train(load_model_args: Dict[str, Any] | None = None, save_model: bool = True)

Train the model. If load_model_args is provided, load the model from MLFlow.

Parameters:

• load_model_args (Optional[Dict[str, Any]]) – The arguments to load the model from mlflow.

• save_model (bool) – Whether to save the model to mlflow.

class graphdoc.DocQualityTrainer(prompt: DocQualityPrompt, optimizer_type: str, optimizer_kwargs: Dict[str, Any], mlflow_model_name: str, mlflow_experiment_name: str, mlflow_tracking_uri: str, trainset: List[Example], evalset: List[Example])

Bases: SinglePromptTrainer

__init__(prompt: DocQualityPrompt, optimizer_type: str, optimizer_kwargs: Dict[str, Any], mlflow_model_name: str, mlflow_experiment_name: str, mlflow_tracking_uri: str, trainset: List[Example], evalset: List[Example])

Initialize the DocQualityTrainer. This is the base class for implementing a trainer for a DocQualityPrompt.

Parameters:

• prompt (DocQualityPrompt) – The prompt to train.

• optimizer_type (str) – The type of optimizer to use.

• optimizer_kwargs (Dict[str, Any]) – The keyword arguments for the optimizer.

• mlflow_model_name (str) – The name of the model in mlflow.

• mlflow_experiment_name (str) – The name of the experiment in mlflow.

• mlflow_tracking_uri (str) – The uri of the mlflow tracking server.

• trainset (List[dspy.Example]) – The training set.

• evalset (List[dspy.Example]) – The evaluation set.

evaluation_metrics(base_evaluation, optimized_evaluation)

Log evaluation metrics to mlflow. We will log the overall scores and the per category scores. Per category scores will be logged as a csv file.

Parameters:

• base_evaluation (Any) – The evaluation metrics of the base model.

• optimized_evaluation (Any) – The evaluation metrics of the optimized model.

evaluate_training(base_model, optimized_model) → Tuple[Dict[str, Any], Dict[str, Any]]

Evaluate the training of the model. Comparing the base and optimized models.

Parameters:

• base_model (Any) – The base model.

• optimized_model (Any) – The optimized model.

train(load_model_args: Dict[str, Any] | None = None, save_model: bool = True)

Train the model. If provided, we will load the model from mlflow. Otherwise, we will use the provided DocQualityPrompt as the base model.

Parameters:

• load_model_args (Dict[str, Any]) – The arguments to load the model.

• save_model (bool) – Whether to save the model.

class graphdoc.SinglePromptTrainer(prompt: SinglePrompt, optimizer_type: str, optimizer_kwargs: Dict[str, Any], mlflow_model_name: str, mlflow_experiment_name: str, mlflow_tracking_uri: str, trainset: List[Example], evalset: List[Example])

Bases: ABC

__init__(prompt: SinglePrompt, optimizer_type: str, optimizer_kwargs: Dict[str, Any], mlflow_model_name: str, mlflow_experiment_name: str, mlflow_tracking_uri: str, trainset: List[Example], evalset: List[Example])

Initialize the SinglePromptTrainer. This is the base class for implementing a trainer for a single prompt.

Parameters:

• prompt (SinglePrompt) – The prompt to train.

• optimizer_type (str) – The type of optimizer to use.

• optimizer_kwargs (Dict[str, Any]) – The keyword arguments for the optimizer.

• mlflow_model_name (str) – The name of the model in mlflow.

• mlflow_experiment_name (str) – The name of the experiment in mlflow.

• mlflow_tracking_uri (str) – The uri of the mlflow tracking server.

• trainset (List[dspy.Example]) – The training set.

• evalset (List[dspy.Example]) – The evaluation set.

abstract evaluation_metrics(base_evaluation, optimized_evaluation)

Log evaluation metrics to mlflow.

Parameters:

• base_evaluation (Any) – The evaluation metrics of the base model.

• optimized_evaluation (Any) – The evaluation metrics of the optimized model.

abstract evaluate_training(base_model, optimized_model) → Tuple[Dict[str, Any], Dict[str, Any]]

Evaluate the training of the model. Comparing the base and optimized models.

Parameters:

• base_model (Any) – The base model.

• optimized_model (Any) – The optimized model.

abstract train(load_model_args: Dict[str, Any] | None = None, save_model: bool = True)

Train the model.

Parameters:

• load_model_args (Dict[str, Any]) – The arguments to load the model.

• save_model (bool) – Whether to save the model.

graphdoc.optimizer_class(optimizer_type: str, optimizer_kwargs: Dict[str, Any])

graphdoc.optimizer_compile(optimizer_type: str, optimizer_kwargs: Dict[str, Any])

Compiles the optimizer given the optimizer type and optimizer kwargs.

Optimizer kwargs are optimizer specific, and must include a student field that maps to a dspy.ChainOfThought, dspy.Predict, etc.

class graphdoc.BadDocGeneratorSignature(*, database_schema: str, documented_schema: str)

Bases: Signature

### TASK: Given a GraphQL Schema, generate intentionally incorrect documentation for the columns of the tables in the database.

### Requirements: - Every table, entity, enum, etc. must have at least one column with a description that is obviosly incorrect. - The documentation must be incorrect and misleading. - The documentation should be scattered, with only some columns having documentation.

### Formatting - Ensure that the schema maintains proper documentation formatting, as is provided.

database_schema: str

documented_schema: str

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class graphdoc.DocGeneratorHelperSignature(*, database_schema: str, documented_schema: str)

Bases: Signature

### TASK: Analyze the provided GraphQL Schema and generate detailed yet concise descriptions for each field within the database tables and enums.

### Requirements: - If the field is unclear, and the documentation result is ambiguous, request additional information: “WARNING: Please provide additional information to avoid confusion”. - Utilize only the verified information from the schema to ensure accuracy. - Descriptions should be factual, straightforward, and avoid any speculative language. - Refrain from using the phrase “in the { table } table” within your descriptions. - Ensure that the documentation adheres to standard schema formatting without modifying the underlying schema structure.

### Formatting: - Maintain consistency with the existing documentation style and structure. - Focus on clarity and precision to aid developers and system architects in understanding the schema’s components effectively.

database_schema: str

documented_schema: str

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class graphdoc.DocGeneratorPrompt(prompt: str | Signature | SignatureMeta, prompt_type: Literal['predict', 'chain_of_thought'] | Callable, prompt_metric: DocQualityPrompt)

Bases: SinglePrompt

DocGeneratorPrompt class for generating documentation for GraphQL schemas.

evaluate_documentation_quality(schema: Example, pred: Prediction, trace=None, scalar=True) → int

Evaluate the quality of the documentation. Utilizes the instantiated metric type to evaluate the quality of the documentation.

Parameters:

• schema (dspy.Example) – The schema to evaluate the documentation for.

• pred (dspy.Prediction) – The predicted documentation.

• trace (Any) – The trace of the prediction.

• scalar (bool) – Whether to return a squared score or the full evaluation object.

Returns:

The squared score or the full evaluation object.

Return type:

int

evaluate_metric(example: Example, prediction: Prediction, trace=None) → Any

This is the metric used to evalaute the prompt.

Parameters:

• example (dspy.Example) – The example to evaluate the metric on.

• prediction (dspy.Prediction) – The prediction to evaluate the metric on.

• trace (Any) – The trace to evaluate the metric on. This is for DSPy.

format_metric(examples: List[Example], overall_score: float, results: List, scores: List) → Dict[str, Any]

Format the metric results into a dictionary.

Parameters:

• examples (List[dspy.Example]) – The examples used to evaluate the metric.

• overall_score (float) – The overall score of the metric.

• results (List) – The results of the metric.

• scores (List) – The scores of the metric.

compare_metrics(base_metrics: Any, optimized_metrics: Any, comparison_value: str = 'overall_score') → bool

Compare the base and optimized metrics.

Parameters:

• base_metrics (Any) – The base metrics.

• optimized_metrics (Any) – The optimized metrics.

• comparison_value (str) – The value to compare.

class graphdoc.DocGeneratorSignature(*, database_schema: str, documented_schema: str)

Bases: Signature

### TASK: Analyze the provided GraphQL Schema and generate detailed yet concise descriptions for each field within the database tables and enums.

### Requirements: - Utilize only the verified information from the schema to ensure accuracy. - Descriptions should be factual, straightforward, and avoid any speculative language. - Refrain from using the phrase “in the { table } table” within your descriptions. - Ensure that the documentation adheres to standard schema formatting without modifying the underlying schema structure. - Make sure that the entities themselves are documented.

### Formatting: - Maintain consistency with the existing documentation style and structure. - Focus on clarity and precision to aid developers and system architects in understanding the schema’s components effectively.

database_schema: str

documented_schema: str

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class graphdoc.DocQualityDemonstrationSignature(*, database_schema: str, category: Literal['perfect', 'almost perfect', 'poor but correct', 'incorrect'], rating: Literal[4, 3, 2, 1])

Bases: Signature

You are evaluating the output of an LLM program, expect hallucinations. Given a GraphQL Schema, evaluate the quality of documentation for that schema and provide a category rating.

The categories are described as: - perfect (4): The documentation contains enough information so that the interpretation of the schema and its database content is completely free of ambiguity.

perfect (4) example: type Domain @entity {

“ The namehash (id) of the parent name. References the Domain entity that is the parent of the current domain. Type: Domain ” parent: Domain

}

• almost perfect (3): The documentation is almost perfect and free from ambiguity, but there is room for improvement.

  almost perfect (3) example: type Token @entity {

  “ Name of the token, mirrored from the smart contract ” name: String!

  }

• poor but correct (2): The documentation is poor but correct and has room for improvement due to missing information. The documentation is not incorrect.

  poor but correct (2) example: type InterestRate @entity {

  “Description for column: id” id: ID!

  }

• incorrect (1): The documentation is incorrect and contains inaccurate or misleading information. Any incorrect information automatically leads to an incorrect rating, even if some correct information is present.

  incorrect (1) example: type BridgeProtocol implements Protocol @entity {

  “ Social Security Number of the protocol’s main developer ” id: Bytes!

  }

Output a number rating that corresponds to the categories described above.

database_schema: str

category: Literal['perfect', 'almost perfect', 'poor but correct', 'incorrect']

rating: Literal[4, 3, 2, 1]

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class graphdoc.DocQualityPrompt(prompt: Literal['doc_quality', 'doc_quality_demo'] | Signature | SignatureMeta = 'doc_quality', prompt_type: Literal['predict', 'chain_of_thought'] | Callable = 'predict', prompt_metric: Literal['rating', 'category'] | Callable = 'rating')

Bases: SinglePrompt

DocQualityPrompt class for evaluating documentation quality.

This is a single prompt that can be used to evaluate the quality of the documentation for a given schema. This is a wrapper around the SinglePrompt class that implements the abstract methods.

__init__(prompt: Literal['doc_quality', 'doc_quality_demo'] | Signature | SignatureMeta = 'doc_quality', prompt_type: Literal['predict', 'chain_of_thought'] | Callable = 'predict', prompt_metric: Literal['rating', 'category'] | Callable = 'rating') → None

Initialize the DocQualityPrompt.

Parameters:

• prompt (Union[str, dspy.Signature]) – The prompt to use. Can either be a string that maps to a defined signature, as set in the doc_quality_factory, or a dspy.Signature.

• prompt_type (Union[Literal["predict", "chain_of_thought"], Callable]) – The type of prompt to use.

• prompt_metric (Union[Literal["rating", "category"], Callable]) – The metric to use. Can either be a string that maps to a defined metric, as set in the doc_quality_factory, or a custom callable function. Function must have the signature (example: dspy.Example, prediction: dspy.Prediction) -> bool.

evaluate_metric(example: Example, prediction: Prediction, trace=None) → bool

Evaluate the metric for the given example and prediction.

Parameters:

• example (dspy.Example) – The example to evaluate the metric on.

• prediction (dspy.Prediction) – The prediction to evaluate the metric on.

• trace (Any) – Used for DSPy.

Returns:

The result of the evaluation. A boolean for if the metric is correct.

Return type:

bool

format_metric(examples: List[Example], overall_score: float, results: List, scores: List) → Dict[str, Any]

Formats evaluation metrics into a structured report containing: - Overall score across all categories - Percentage correct per category - Detailed results for each evaluation

Parameters:

• examples (List[dspy.Example]) – The examples to evaluate the metric on.

• overall_score (float) – The overall score across all categories.

• results (List) – The results of the evaluation.

• scores (List) – The scores of the evaluation.

Returns:

A dictionary containing the overall score, per category scores, and details. { “overall_score”: 0, “per_category_scores”: {}, “details”: [], “results”: [] }

Return type:

Dict[str, Any]

compare_metrics(base_metrics: Any, optimized_metrics: Any, comparison_value: str = 'overall_score') → bool

Compare the metrics of the base and optimized models. Returns true if the optimized model is better than the base model.

Parameters:

• base_metrics (Any) – The metrics of the base model.

• optimized_metrics (Any) – The metrics of the optimized model.

• comparison_value (str) – The value to compare.

Returns:

True if the optimized model is better than the base model.

Return type:

bool

class graphdoc.DocQualitySignature(*, database_schema: str, category: Literal['perfect', 'almost perfect', 'poor but correct', 'incorrect'], rating: Literal[4, 3, 2, 1])

Bases: Signature

You are a documentation quality evaluator specializing in GraphQL schemas. Your task is to assess the quality of documentation provided for a given database schema. Carefully analyze the schema’s descriptions for clarity, accuracy, and completeness. Categorize the documentation into one of the following ratings based on your evaluation: - perfect (4): The documentation is comprehensive and leaves no room for ambiguity in understanding the schema and its database content. - almost perfect (3): The documentation is clear and mostly free of ambiguity, but there is potential for further improvement. - poor but correct (2): The documentation is correct but lacks detail, resulting in some ambiguity. It requires enhancement to be more informative. - incorrect (1): The documentation contains errors or misleading information, regardless of any correct segments present. Such inaccuracies necessitate an incorrect rating. Provide a step-by-step reasoning to support your evaluation, along with the appropriate category label and numerical rating.

database_schema: str

category: Literal['perfect', 'almost perfect', 'poor but correct', 'incorrect']

rating: Literal[4, 3, 2, 1]

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class graphdoc.PromptFactory

Bases: object

static single_prompt(prompt: str | Signature | SignatureMeta, prompt_class: str, prompt_type: str, prompt_metric: str | DocQualityPrompt | SinglePrompt) → SinglePrompt

Returns an instance of the specified prompt class. Allows for the user to pass in their own dspy signature.

Parameters:

• prompt (Union[str, dspy.Signature]) – The prompt to use.

• prompt_class (str) – The class of the prompt to use.

• prompt_type (str) – The type of the prompt to use.

• prompt_metric (Union[str, DocQualityPrompt, SinglePrompt]) – The metric to use for the prompt.

Returns:

An instance of the specified prompt class.

Return type:

SinglePrompt

class graphdoc.SinglePrompt(prompt: Signature | SignatureMeta, prompt_type: Literal['predict', 'chain_of_thought'] | Callable, prompt_metric: Any)

Bases: ABC

__init__(prompt: Signature | SignatureMeta, prompt_type: Literal['predict', 'chain_of_thought'] | Callable, prompt_metric: Any) → None

Initialize a single prompt.

Parameters:

• prompt (dspy.Signature) – The prompt to use.

• prompt_type (Union[Literal["predict", "chain_of_thought"], Callable]) – The type of prompt to use. Can be “predict” or “chain_of_thought”. Optionally, pass another dspy.Module.

• prompt_metric (Any) – The metric to use. Marked as Any for flexibility (as metrics can be other prompts).

abstract evaluate_metric(example: Example, prediction: Prediction, trace=None) → Any

This is the metric used to evalaute the prompt.

Parameters:

• example (dspy.Example) – The example to evaluate the metric on.

• prediction (dspy.Prediction) – The prediction to evaluate the metric on.

• trace (Any) – The trace to evaluate the metric on. This is for DSPy.

abstract format_metric(examples: List[Example], overall_score: float, results: List, scores: List) → Dict[str, Any]

This takes the results from the evaluate_evalset and does any necessary formatting, taking into account the metric type.

Parameters:

• examples (List[dspy.Example]) – The examples to evaluate the metric on.

• overall_score (float) – The overall score of the metric.

• results (List) – The results from the evaluate_evalset.

• scores (List) – The scores from the evaluate_evalset.

abstract compare_metrics(base_metrics: Any, optimized_metrics: Any, comparison_value: str = 'overall_score') → bool

Compare the metrics of the base and optimized models. Return true if the optimized model is better than the base model.

Parameters:

• base_metrics (Any) – The metrics of the base model.

• optimized_metrics (Any) – The metrics of the optimized model.

• comparison_value (str) – The value to compare the metrics on. Determines which metric is used to compare the models.

Returns:

True if the optimized model is better than the base model, False otherwise.

Return type:

bool

evaluate_evalset(examples: List[Example], num_threads: int = 1, display_progress: bool = True, display_table: bool = True) → Dict[str, Any]

Take in a list of examples and evaluate the results.

Parameters:

• examples (List[dspy.Example]) – The examples to evaluate the results on.

• num_threads (int) – The number of threads to use for evaluation.

• display_progress (bool) – Whether to display the progress of the evaluation.

• display_table (bool) – Whether to display the table of the evaluation.

Returns:

A dictionary containing the overall score, results, and scores.

Return type:

Dict[str, Any]

class graphdoc.DspyDataHelper

Bases: ABC

Abstract class for creating data objects related to a given dspy.Signature.

prompt_signature() → Signature | SignatureMeta

Given a prompt, return a dspy.Signature object.

Parameters:

prompt (Any) – A prompt.

static _(prompt: Predict) → Signature | SignatureMeta

Given a dspy.Predict object, return a dspy.Signature object.

static formatted_signature(signature: Signature | SignatureMeta, example: Example) → str

Given a dspy.Signature and a dspy.Example, return a formatted signature as a string.

Parameters:

• signature (dspy.Signature) – A dspy.Signature object.

• example (dspy.Example) – A dspy.Example object.

Returns:

A formatted signature as a string.

Return type:

str

abstract static example(inputs: dict[str, Any]) → Example

Given a dictionary of inputs, return a dspy.Example object.

Parameters:

inputs (dict[str, Any]) – A dictionary of inputs.

Returns:

A dspy.Example object.

Return type:

dspy.Example

abstract static example_example() → Example

Return an example dspy.Example object with the inputs set to the example values.

Returns:

A dspy.Example object.

Return type:

dspy.Example

abstract static model_signature() → ModelSignature

Return a mlflow.models.ModelSignature object. Based on the example object, removes the output fields and utilizes the remaining fields to infer the model signature.

Returns:

A mlflow.models.ModelSignature object.

Return type:

mlflow.models.ModelSignature

abstract static prediction(inputs: dict[str, Any]) → Prediction

Given a dictionary of inputs, return a dspy.Prediction object.

Parameters:

inputs (dict[str, Any]) – A dictionary of inputs.

Returns:

A dspy.Prediction object.

Return type:

dspy.Prediction

abstract static prediction_example() → Prediction

Return an example dspy.Prediction object with the inputs set to the example values.

Returns:

A dspy.Prediction object.

Return type:

dspy.Prediction

abstract static trainset(inputs: dict[str, Any] | Dataset, filter_args: dict[str, Any] | None = None) → list[Example]

Given a dictionary of inputs or a datasets.Dataset object, return a list of dspy.Example objects.

Parameters:

• inputs (Union[dict[str, Any], datasets.Dataset]) – A dictionary of inputs or a datasets.Dataset object.

• filter_args (Optional[dict[str, Any]]) – A dictionary of filter arguments. These are instructions for how we will filter and / or transform the inputs.

Returns:

A list of dspy.Example objects.

Return type:

list[dspy.Example]

class graphdoc.GenerationDataHelper

Bases: DspyDataHelper

A helper class for creating data objects related to our Documentation Generation dspy.Signature.

The example signature is defined as: ` database_schema: str = dspy.InputField() documented_schema: str = dspy.OutputField() `

static example(inputs: dict[str, Any]) → Example

Given a dictionary of inputs, return a dspy.Example object.

Parameters:

inputs (dict[str, Any]) – A dictionary of inputs.

Returns:

A dspy.Example object.

Return type:

dspy.Example

static example_example() → Example

Return an example dspy.Example object with the inputs set to the example values.

Returns:

A dspy.Example object.

Return type:

dspy.Example

static model_signature() → ModelSignature

Return a mlflow.models.ModelSignature object. Based on the example object, removes the output fields and utilizes the remaining fields to infer the model signature.

Returns:

A mlflow.models.ModelSignature object.

Return type:

mlflow.models.ModelSignature

static prediction(inputs: dict[str, Any]) → Prediction

Given a dictionary of inputs, return a dspy.Prediction object.

Parameters:

inputs (dict[str, Any]) – A dictionary of inputs.

Returns:

A dspy.Prediction object.

Return type:

dspy.Prediction

static prediction_example() → Prediction

Return an example dspy.Prediction object with the inputs set to the example values.

Returns:

A dspy.Prediction object.

Return type:

dspy.Prediction

static trainset(inputs: dict[str, Any] | Dataset, filter_args: dict[str, Any] | None = None) → list[Example]

Given a dictionary of inputs or a datasets.Dataset object, return a list of dspy.Example objects.

Parameters:

• inputs (Union[dict[str, Any], datasets.Dataset]) – A dictionary of inputs or a datasets.Dataset object.

• filter_args (Optional[dict[str, Any]]) – A dictionary of filter arguments. These are instructions for how we will filter and / or transform the inputs.

Returns:

A list of dspy.Example objects.

Return type:

list[dspy.Example]

class graphdoc.LocalDataHelper(schema_directory_path: str | ~pathlib._local.Path | None = None, categories: ~typing.Type[~enum.Enum] = <enum 'SchemaCategory'>, ratings: ~typing.Type[~enum.Enum] = <enum 'SchemaRating'>, categories_ratings: ~typing.Callable = <function SchemaCategoryRatingMapping.get_rating>)

Bases: object

A helper class for loading data from a directory.

Parameters:

• schema_directory_path (Union[str, Path] Defaults to the path to the schemas in the graphdoc package.) – The path to the directory containing the schemas

• categories (Type[Enum]) – The categories of the schemas. Defaults to SchemaCategory.

• ratings (Type[Enum]) – The ratings of the schemas. Defaults to SchemaRating.

• categories_ratings – A callable that maps categories to ratings. Defaults to SchemaCategoryRatingMapping.get_rating.

schema_objects_from_folder(category: str, rating: int, folder_path: str | Path) → dict[str, SchemaObject]

Load schemas from a folder, keeping the difficulty tag.

Parameters:

• category (str) – The category of the schemas

• rating (int) – The rating of the schemas

• folder_path (Union[str, Path]) – The path to the folder containing the schemas

Returns:

A dictionary of schemas

Return type:

dict[str, SchemaObject]

schema_objects_from_folder_of_folders(folder_paths: ~typing.Type[~enum.Enum] | None = <enum 'SchemaCategoryPath'>) → Dict[str, SchemaObject] | None

Load a folder of folders containing schemas, keeping the difficulty tag.

Parameters:

folder_paths (Optional[Type[Enum]]) – Enum class defining folder paths, defaults to SchemaCategoryPath. Must have a get_path method.

Returns:

Dictionary of loaded schemas

Return type:

Union[Dict[str, SchemaObject], None]

folder_to_dataset(category: str, folder_path: str | Path, parse_objects: bool = True, type_mapping: dict[type, str] | None = None) → Dataset

Load a folder of schemas, keeping the difficulty tag.

Parameters:

• category (str) – The category of the schemas

• folder_path (Union[str, Path]) – The path to the folder containing the schemas

• parse_objects (bool) – Whether to parse the objects from the schemas

• type_mapping (Optional[dict[type, str]]) – A dictionary mapping types to strings

Returns:

A dataset containing the schemas

Return type:

Dataset

folder_of_folders_to_dataset(folder_paths: ~typing.Type[~enum.Enum] = <enum 'SchemaCategoryPath'>, parse_objects: bool = True, type_mapping: dict[type, str] | None = None) → Dataset

Load a folder of folders containing schemas, keeping the difficulty tag.

Parameters:

• folder_paths (Type[Enum]) – Enum class defining folder paths, defaults to SchemaCategoryPath. Must have a get_path method.

• parse_objects (bool) – Whether to parse the objects from the schemas

• type_mapping (Optional[dict[type, str]]) – A dictionary mapping graphql-ast node values to strings

Returns:

A dataset containing the schemas

Return type:

Dataset

class graphdoc.MlflowDataHelper(mlflow_tracking_uri: str | Path, mlflow_tracking_username: str | None = None, mlflow_tracking_password: str | None = None)

Bases: object

__init__(mlflow_tracking_uri: str | Path, mlflow_tracking_username: str | None = None, mlflow_tracking_password: str | None = None)

A helper class for loading and saving models and metadata from mlflow.

Parameters:

• mlflow_tracking_uri (Union[str, Path]) – The uri of the mlflow tracking server.

• mlflow_tracking_username (Optional[str]) – The username for the mlflow tracking server.

• mlflow_tracking_password (Optional[str]) – The password for the mlflow tracking server.

update_auth_env_vars(mlflow_tracking_username: str, mlflow_tracking_password: str)

Update the authentication environment variables.

Parameters:

• mlflow_tracking_username (str) – The username for the mlflow tracking server.

• mlflow_tracking_password (str) – The password for the mlflow tracking server.

set_auth_env_vars()

Set the authentication environment variables.

latest_model_version(model_name: str)

Load the latest version of a model from mlflow.

Parameters:

model_name (str) – The name of the model to load.

Returns:

The loaded model.

model_by_name_and_version(model_name: str, model_version: str)

Load a model from mlflow by name and version.

Parameters:

• model_name (str) – The name of the model to load.

• model_version (str) – The version of the model to load.

Returns:

The loaded model.

model_by_uri(model_uri: str)

Load a model from mlflow by uri.

Parameters:

model_uri (str) – The uri of the model to load.

Returns:

The loaded model.

model_by_args(load_model_args: Dict[str, str])

Given a dictionary of arguments, load a model from mlflow. Ordering is model_by_uri, model_by_name_and_version, latest_model_version.

Parameters:

load_model_args (Dict[str, str]) – A dictionary of arguments.

Returns:

The loaded model.

save_model(model: Signature, model_signature: ModelSignature, model_name: str)

Save a model to mlflow.

Parameters:

• model (dspy.Signature) – The model to save.

• model_signature (ModelSignature) – The signature of the model.

• model_name (str) – The name of the model to save.

run_parameters(run_id: str) → dict[str, Any]

Load the parameters of a run from mlflow.

Parameters:

run_id (str) – The id of the run to load the parameters from.

Returns:

The parameters of the run.

class graphdoc.Parser(type_mapping: dict[type, str] | None = None)

Bases: object

A class for parsing and handling of GraphQL objects.

DEFAULT_NODE_TYPES = {<class 'graphql.language.ast.DocumentNode'>: 'full schema', <class 'graphql.language.ast.EnumTypeDefinitionNode'>: 'enum schema', <class 'graphql.language.ast.EnumValueDefinitionNode'>: 'enum value', <class 'graphql.language.ast.ObjectTypeDefinitionNode'>: 'table schema'}

static _check_node_type(node: Node, type_mapping: dict[type, str] | None = None) → str

Check the type of a schema node.

Parameters:

• node (Node) – The schema node to check

• type_mapping (Optional[dict[type, str]]) – Custom mapping of node types to strings. Defaults to DEFAULT_NODE_TYPES

Returns:

The type of the schema node

Return type:

str

static parse_schema_from_file(schema_file: str | Path, schema_directory_path: str | Path | None = None) → DocumentNode

Parse a schema from a file.

Parameters:

• schema_file (Union[str, Path]) – The name of the schema file

• schema_directory_path (Optional[Union[str, Path]]) – A path to a directory containing schemas

Returns:

The parsed schema

Return type:

DocumentNode

Raises:

Exception – If the schema cannot be parsed

static update_node_descriptions(node: Node, new_value: str | None = None) → Node

Given a GraphQL node, recursively traverse the node and its children, updating all descriptions with the new value. Can also be used to remove descriptions by passing None as the new value.

Parameters:

• node (Node) – The GraphQL node to update

• new_value (Optional[str]) – The new description value. If None, the description will be removed.

Returns:

The updated node

Return type:

Node

static count_description_pattern_matching(node: Node, pattern: str) → dict[str, int]

Counts the number of times a pattern matches a description in a node and its children.

Parameters:

• node (Node) – The GraphQL node to count the pattern matches in

• pattern (str) – The pattern to count the matches of

Returns:

A dictionary with the counts of matches

Return type:

dict[str, int]

static fill_empty_descriptions(node: Node, new_column_value: str = 'Description for column: {}', new_table_value: str = 'Description for table: {}', use_value_name: bool = True, value_name: str | None = None)

Recursively traverse the node and its children, filling in empty descriptions with the new column or table value. Do not update descriptions that already have a value. Default values are provided for the new column and table descriptions.

Parameters:

• node (Node) – The GraphQL node to update

• new_column_value (str) – The new column description value

• new_table_value (str) – The new table description value

• use_value_name (bool) – Whether to use the value name in the description

• value_name (Optional[str]) – The name of the value

Returns:

The updated node

Return type:

Node

static schema_equality_check(gold_node: Node, check_node: Node) → bool

A method to check if two schema nodes are equal. Only checks that the schemas structures are equal, not the descriptions.

Parameters:

• gold_node (Node) – The gold standard schema node

• check_node (Node) – The schema node to check

Returns:

Whether the schemas are equal

Return type:

bool

static schema_object_from_file(schema_file: str | Path, category: str | None = None, rating: int | None = None) → SchemaObject

Parse a schema object from a file.

static parse_objects_from_full_schema_object(schema: SchemaObject, type_mapping: dict[type, str] | None = None) → dict[str, SchemaObject] | None

Parse out all available tables from a full schema object.

Parameters:

• schema (SchemaObject) – The full schema object to parse

• type_mapping (Optional[dict[type, str]]) – Custom mapping of node types to strings. Defaults to DEFAULT_NODE_TYPES

Returns:

The parsed objects (tables and enums)

Return type:

Union[dict, None]

class graphdoc.QualityDataHelper

Bases: DspyDataHelper

A helper class for creating data objects related to our Documentation Quality dspy.Signature.

The example signature is defined as:

database_schema: str = dspy.InputField()
category: Literal["perfect", "almost perfect", "poor but correct", "incorrect"] = (
    dspy.OutputField()
)
rating: Literal[4, 3, 2, 1] = dspy.OutputField()

static example(inputs: dict[str, Any]) → Example

Given a dictionary of inputs, return a dspy.Example object.

Parameters:

inputs (dict[str, Any]) – A dictionary of inputs.

Returns:

A dspy.Example object.

Return type:

dspy.Example

static example_example() → Example

Return an example dspy.Example object with the inputs set to the example values.

Returns:

A dspy.Example object.

Return type:

dspy.Example

static model_signature() → ModelSignature

Return a mlflow.models.ModelSignature object. Based on the example object, removes the output fields and utilizes the remaining fields to infer the model signature.

Returns:

A mlflow.models.ModelSignature object.

Return type:

mlflow.models.ModelSignature

static prediction(inputs: dict[str, Any]) → Prediction

Given a dictionary of inputs, return a dspy.Prediction object.

Parameters:

inputs (dict[str, Any]) – A dictionary of inputs.

Returns:

A dspy.Prediction object.

Return type:

dspy.Prediction

static prediction_example() → Prediction

Return an example dspy.Prediction object with the inputs set to the example values.

Returns:

A dspy.Prediction object.

Return type:

dspy.Prediction

static trainset(inputs: dict[str, Any] | Dataset, filter_args: dict[str, Any] | None = None) → list[Example]

Given a dictionary of inputs or a datasets.Dataset object, return a list of dspy.Example objects.

Parameters:

• inputs (Union[dict[str, Any], datasets.Dataset]) – A dictionary of inputs or a datasets.Dataset object.

• filter_args (Optional[dict[str, Any]]) – A dictionary of filter arguments. These are instructions for how we will filter and / or transform the inputs.

Returns:

A list of dspy.Example objects.

Return type:

list[dspy.Example]

class graphdoc.SchemaCategory(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Schema quality categories enumeration.

PERFECT = 'perfect'

ALMOST_PERFECT = 'almost perfect'

POOR_BUT_CORRECT = 'poor but correct'

INCORRECT = 'incorrect'

BLANK = 'blank'

classmethod from_str(value: str) → SchemaCategory | None

class graphdoc.SchemaCategoryPath(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Maps schema categories to their folder names.

PERFECT = 'perfect'

ALMOST_PERFECT = 'almost_perfect'

POOR_BUT_CORRECT = 'poor_but_correct'

INCORRECT = 'incorrect'

BLANK = 'blank'

classmethod get_path(category: SchemaCategory, folder_path: str | Path) → Path | None

Get the folder path for a given schema category and folder path.

Parameters:

category – The schema category

Returns:

The corresponding folder path

class graphdoc.SchemaCategoryRatingMapping

Bases: object

Mapping between schema categories and ratings.

static get_rating(category: SchemaCategory) → SchemaRating

Get the corresponding rating for a given schema category.

Parameters:

category – The schema category

Returns:

The corresponding rating

static get_category(rating: SchemaRating) → SchemaCategory

Get the corresponding category for a given schema rating.

Parameters:

rating – The schema rating

Returns:

The corresponding category

class graphdoc.SchemaObject(key: str, category: Enum | None = None, rating: Enum | None = None, schema_name: str | None = None, schema_type: Enum | None = None, schema_str: str | None = None, schema_ast: Node | None = None)

Bases: object

Schema object containing schema data and metadata.

key: str

category: Enum | None = None

rating: Enum | None = None

schema_name: str | None = None

schema_type: Enum | None = None

schema_str: str | None = None

schema_ast: Node | None = None

classmethod from_dict(data: dict, category_enum: ~typing.Type[~enum.Enum] = <enum 'SchemaCategory'>, rating_enum: ~typing.Type[~enum.Enum] = <enum 'SchemaRating'>, type_enum: ~typing.Type[~enum.Enum] = <enum 'SchemaType'>) → SchemaObject

Create SchemaObject from dictionary with validation.

Parameters:

• data – The data dictionary

• category_enum – Custom Enum class for categories

• rating_enum – Custom Enum class for ratings

• type_enum – Custom Enum class for schema types

to_dict() → dict

Convert the SchemaObject to a dictionary, excluding the key field.

Returns:

Dictionary representation of the SchemaObject without the key

Return type:

dict

static _hf_schema_object_columns() → Features

Return the columns for the graph_doc dataset, based on the SchemaObject fields.

Returns:

The columns for the graph_doc dataset

Return type:

Features

to_dataset() → Dataset

Convert the SchemaObject to a Hugging Face Dataset.

Returns:

The Hugging Face Dataset

Return type:

Dataset

class graphdoc.SchemaRating(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Schema quality ratings enumeration.

FOUR = '4'

THREE = '3'

TWO = '2'

ONE = '1'

ZERO = '0'

classmethod from_value(value: str | int) → SchemaRating | None

class graphdoc.SchemaType(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Schema type enumeration.

FULL_SCHEMA = 'full schema'

TABLE_SCHEMA = 'table schema'

ENUM_SCHEMA = 'enum schema'

classmethod from_str(value: str) → SchemaType | None

graphdoc._env_constructor(loader: SafeLoader, node: ScalarNode) → str

Custom constructor for environment variables.

Parameters:

• loader (yaml.SafeLoader) – The YAML loader.

• node (yaml.nodes.ScalarNode) – The node to construct.

Returns:

The environment variable value.

Return type:

str

Raises:

ValueError – If the environment variable is not set.

graphdoc.check_directory_path(directory_path: str | Path) → None

Check if the provided path resolves to a valid directory.

Parameters:

directory_path (Union[str, Path]) – The path to check.

Raises:

ValueError – If the path does not resolve to a valid directory.

Returns:

None

Return type:

None

graphdoc.check_file_path(file_path: str | Path) → None

Check if the provided path resolves to a valid file.

Parameters:

file_path (Union[str, Path]) – The path to check.

Raises:

ValueError – If the path does not resolve to a valid file.

Returns:

None

Return type:

None

graphdoc.load_yaml_config(file_path: str | Path, use_env: bool = True) → dict

Load a YAML configuration file.

Parameters:

• file_path (Union[str, Path]) – The path to the YAML file.

• use_env (bool) – Whether to use environment variables.

Returns:

The YAML configuration.

Return type:

dict

Raises:

ValueError – If the path does not resolve to a valid file or the environment variable is not set.

graphdoc.schema_objects_to_dataset(schema_objects: List[SchemaObject]) → Dataset

Convert a list of SchemaObjects to a Hugging Face Dataset.

Parameters:

schema_objects – The list of SchemaObjects

Returns:

The Hugging Face Dataset

graphdoc.setup_logging(log_level: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'])

Setup logging for the application.

Parameters:

log_level (Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]) – The log level.

graphdoc.mlflow_data_helper_from_dict(mlflow_config: dict) → MlflowDataHelper

Load a MLflow data helper from a dictionary of parameters.

The following keys are expected: - mlflow_tracking_uri - mlflow_tracking_username (optional) - mlflow_tracking_password (optional)

{
    "mlflow_tracking_uri": "http://localhost:5000",
    "mlflow_tracking_username": "admin",
    "mlflow_tracking_password": "password"
}

Parameters:

mlflow_config (dict) – Dictionary containing MLflow parameters.

Returns:

A MlflowDataHelper object.

Return type:

MlflowDataHelper

graphdoc.mlflow_data_helper_from_yaml(yaml_path: str | Path) → MlflowDataHelper

Load a mlflow data helper from a YAML file.

Parameters:

yaml_path (Union[str, Path]) – Path to the YAML file.

mlflow:
    mlflow_tracking_uri:      !env MLFLOW_TRACKING_URI      # The tracking URI for MLflow
    mlflow_tracking_username: !env MLFLOW_TRACKING_USERNAME # The username for the mlflow tracking server
    mlflow_tracking_password: !env MLFLOW_TRACKING_PASSWORD # The password for the mlflow tracking server

graphdoc.trainset_from_dict(trainset_dict: dict) → List[Example]

Load a trainset from a dictionary of parameters.

{
    "hf_api_key": !env HF_DATASET_KEY,          # Must be a valid Hugging
                                                # Face API key
                                                # (with permission to
                                                # access graphdoc)
                                                # TODO: we may make
                                                # this public in the future
    "load_from_hf": false,                      # Whether to load the dataset
                                                # from Hugging Face
    "load_from_local": true,                    # Whether to load the dataset
                                                # from a local directory
    "load_local_specific_category": false,      # Whether to load all categories
                                                # or a specific category
    "local_specific_category": perfect,         # The specific category
                                                # (if load_from_local is true)
    "local_parse_objects": true,                # Whether to parse the objects
                                                # in the dataset
                                                # (if load_from_local is true)
    "split_for_eval": true,                     # Whether to split the dataset
                                                # into trainset and evalset
    "trainset_size": 1000,                      # The size of the trainset
    "evalset_ratio": 0.1,                       # The proportionate size of evalset
    "data_helper_type": "quality"               # Type of data helper to use
                                                # (quality, generation)
}

Parameters:

trainset_dict (dict) – Dictionary containing trainset parameters.

Returns:

A trainset.

Return type:

List[dspy.Example]

graphdoc.trainset_from_yaml(yaml_path: str | Path) → List[Example]

Load a trainset from a YAML file.

data:
    hf_api_key: !env HF_DATASET_KEY         # Must be a valid Hugging Face API key
                                            # (with permission to access graphdoc)
                                            # TODO: we may make this public
    load_from_hf: false                     # Load the dataset from Hugging Face
    load_from_local: true                   # Load the dataset from a local directory
    load_local_specific_category: false     # Load all categories or a specific category
                                            # (if load_from_local is true)
    local_specific_category: perfect,       # Which category to load from the dataset
                                            # (if load_from_local is true)
    local_parse_objects: true,              # Whether to parse the objects
                                            # in the dataset
                                            # (if load_from_local is true)
    split_for_eval: true,                   # Whether to split the dataset
                                            # into trainset and evalset
    trainset_size: 1000,                    # The size of the trainset
    evalset_ratio: 0.1,                     # The proportionate size of evalset
    data_helper_type: quality               # Type of data helper to use
                                            # (quality, generation)

Parameters:

yaml_path (Union[str, Path]) – Path to the YAML file.

Returns:

A trainset.

Return type:

List[dspy.Example]

graphdoc.split_trainset(trainset: List[Example], evalset_ratio: float, seed: int = 42) → tuple[List[Example], List[Example]]

Split a trainset into a trainset and evalset.

Parameters:

• trainset (List[dspy.Example]) – The trainset to split.

• evalset_ratio (float) – The proportionate size of the evalset.

Returns:

A tuple of trainset and evalset.

Return type:

tuple[List[dspy.Example], List[dspy.Example]]

graphdoc.trainset_and_evalset_from_yaml(yaml_path: str | Path) → tuple[List[Example], List[Example]]

Load a trainset and evalset from a YAML file.

data:
    hf_api_key: !env HF_DATASET_KEY         # Must be a valid Hugging Face API key
                                            # (with permission to access graphdoc)
                                            # TODO: we may make this public
    load_from_hf: false                     # Load the dataset from Hugging Face
    load_from_local: true                   # Load the dataset from a local directory
    load_local_specific_category: false     # Load all categories or a specific category
                                            # (if load_from_local is true)
    local_specific_category: perfect,       # Which category to load from the dataset
                                            # (if load_from_local is true)
    local_parse_objects: true,              # Whether to parse the objects
                                            # in the dataset
                                            # (if load_from_local is true)
    split_for_eval: true,                   # Whether to split the dataset
                                            # into trainset and evalset
    trainset_size: 1000,                    # The size of the trainset
    evalset_ratio: 0.1,                     # The proportionate size of evalset
    data_helper_type: quality               # Type of data helper to use
                                            # (quality, generation)
    seed: 42                                # The seed for the random number generator

Parameters:

yaml_path (Union[str, Path]) – Path to the YAML file.

Returns:

A tuple of trainset and evalset.

Return type:

tuple[List[dspy.Example], List[dspy.Example]]

graphdoc.single_prompt_from_dict(prompt_dict: dict, prompt_metric: str | SinglePrompt, mlflow_dict: dict | None = None) → SinglePrompt

Load a single prompt from a dictionary of parameters.

{
    "prompt": "doc_quality",             # Which prompt signature to use
    "class": "SchemaDocQualityPrompt",   # Must be a child of SinglePrompt
    "type": "predict",                   # Must be one of predict, generate
    "metric": "rating",                  # The metric to use for evaluation
    "load_from_mlflow": false,           # Whether to load the prompt from MLflow
    "model_uri": null,                   # The tracking URI for MLflow
    "model_name": null,                  # The name of the model in MLflow
    "model_version": null                # The version of the model in MLflow
}

Parameters:

• prompt_dict (dict) – Dictionary containing prompt parameters.

• prompt_metric (Union[str, SinglePrompt]) – The prompt to use for the metric.

• mlflow_dict (Optional[dict]) – Dictionary containing MLflow parameters.

Returns:

A SinglePrompt object.

Return type:

SinglePrompt

graphdoc.single_prompt_from_yaml(yaml_path: str | Path) → SinglePrompt

Load a single prompt from a YAML file.

prompt:
    prompt: base_doc_gen        # Which prompt signature to use
    class: DocGeneratorPrompt   # Must be a child of SinglePrompt
                                # (we will use an enum to map this)
    type: chain_of_thought      # The type of prompt to use
                                # (predict, chain_of_thought)
    metric: rating              # The type of metric to use
                                # (rating, category)
    load_from_mlflow: false     # Whether to load the prompt
                                # from an MLFlow URI
    model_uri: null             # The tracking URI for MLflow
    model_name: null            # The name of the model in MLflow
    model_version: null         # The version of the model in MLflow
    prompt_metric: true         # Whether another prompt is used
                                # to calculate the metric
                                # (in which case we must also load that prompt)

prompt_metric:
    prompt: doc_quality         # The prompt to use to calculate
                                # the metric
    class: DocQualityPrompt     # The class of the prompt to use
                                # to calculate the metric
    type: predict               # The type of prompt to use
                                # to calculate the metric
    metric: rating              # The metric to use to calculate
                                # the metric
    load_from_mlflow: false     # Whether to load the prompt
                                # from an MLFlow URI

Parameters:

yaml_path (str) – Path to the YAML file.

Returns:

A SinglePrompt object.

Return type:

SinglePrompt

graphdoc.single_trainer_from_dict(trainer_dict: dict, prompt: SinglePrompt, trainset: List[Example] | None = None, evalset: List[Example] | None = None) → SinglePromptTrainer

Load a single trainer from a dictionary of parameters.

{
    "mlflow": {
        "mlflow_tracking_uri": "http://localhost:5000",
        "mlflow_tracking_username": "admin",
        "mlflow_tracking_password": "password",
    },
    "trainer": {
        "class": "DocQualityTrainer",
        "mlflow_model_name": "doc_quality_model",
        "mlflow_experiment_name": "doc_quality_experiment",
    },
    "optimizer": {
        "optimizer_type": "miprov2",
        "auto": "light",
        "max_labeled_demos": 2,
        "max_bootstrapped_demos": 4,
        "num_trials": 2,
        "minibatch": true
    },
}

Parameters:

• trainer_dict (dict) – Dictionary containing trainer parameters.

• prompt (SinglePrompt) – The prompt to use for this trainer.

Returns:

A SinglePromptTrainer object.

Return type:

SinglePromptTrainer

graphdoc.single_trainer_from_yaml(yaml_path: str | Path) → SinglePromptTrainer

Load a single prompt trainer from a YAML file.

trainer:
    hf_api_key: !env HF_DATASET_KEY         # Must be a valid Hugging Face API key
                                            # (with permission to access graphdoc)
                                            # TODO: we may make this public
    load_from_hf: false                     # Load the dataset from Hugging Face
    load_from_local: true                   # Load the dataset from a local directory
    load_local_specific_category: false     # Load all categories or a specific category
                                            # (if load_from_local is true)
    local_specific_category: perfect,       # Which category to load from the dataset
                                            # (if load_from_local is true)
    local_parse_objects: true,              # Whether to parse the objects
                                            # in the dataset
                                            # (if load_from_local is true)
    split_for_eval: true,                   # Whether to split the dataset
                                            # into trainset and evalset
    trainset_size: 1000,                    # The size of the trainset
    evalset_ratio: 0.1,                     # The proportionate size of evalset

prompt:
    prompt: base_doc_gen                    # Which prompt signature to use
    class: DocGeneratorPrompt               # Must be a child of SinglePrompt
                                            # (we will use an enum to map this)
    type: chain_of_thought                  # The type of prompt to use
                                            # (predict, chain_of_thought)
    metric: rating                          # The type of metric to use
                                            # (rating, category)
    load_from_mlflow: false                 # L oad the prompt from an MLFlow URI
    model_uri: null                         # The tracking URI for MLflow
    model_name: null                        # The name of the model in MLflow
    model_version: null                     # The version of the model in MLflow
    prompt_metric: true                     # Whether another prompt is used
                                            # to calculate the metric

Parameters:

yaml_path (Union[str, Path]) – Path to the YAML file.

Returns:

A SinglePromptTrainer object.

Return type:

SinglePromptTrainer

graphdoc.doc_generator_module_from_dict(module_dict: dict, prompt: DocGeneratorPrompt | SinglePrompt) → DocGeneratorModule

Load a single doc generator module from a dictionary of parameters.

{
    "retry": true,
    "retry_limit": 1,
    "rating_threshold": 3,
    "fill_empty_descriptions": true
}

Parameters:

• module_dict (dict) – Dictionary containing module parameters.

• prompt (DocGeneratorPrompt) – The prompt to use for this module.

Returns:

A DocGeneratorModule object.

Return type:

DocGeneratorModule

graphdoc.doc_generator_module_from_yaml(yaml_path: str | Path) → DocGeneratorModule

Load a doc generator module from a YAML file.

prompt:
    prompt: base_doc_gen            # Which prompt signature to use
    class: DocGeneratorPrompt       # Must be a child of SinglePrompt
                                    # (we will use an enum to map this)
    type: chain_of_thought          # The type of prompt to use
                                    # (predict, chain_of_thought)
    metric: rating                  # The type of metric to use
                                    # (rating, category)
    load_from_mlflow: false         # Whether to load the prompt
                                    # from an MLFlow URI
    model_uri: null                 # The tracking URI for MLflow
    model_name: null                # The name of the model in MLflow
    model_version: null             # The version of the model in MLflow
    prompt_metric: true             # Whether another prompt is used
                                    # to calculate the metric
                                    # (in which case we must load that prompt)

prompt_metric:
    prompt: doc_quality             # The prompt to use to calculate the metric
    class: DocQualityPrompt         # The class of the prompt to use
                                    # to calculate the metric
    type: predict                   # The type of prompt to use
                                    # to calculate the metric
    metric: rating                  # The metric to use to calculate the metric
    load_from_mlflow: false         # Whether to load the prompt
                                    # from an MLFlow URI
    model_uri: null                 # The tracking URI for MLflow
    model_name: null                # The name of the model in MLflow
    model_version: null             # The version of the model in MLflow

module:
    retry: true                     # Whether to retry the generation
                                    # if the quality check fails
    retry_limit: 1                  # The maximum number of retries
    rating_threshold: 3             # The rating threshold for the quality check
    fill_empty_descriptions: true   # Whether to fill empty descriptions with
                                    # generated documentation

Parameters:

yaml_path (Union[str, Path]) – Path to the YAML file.

Returns:

A DocGeneratorModule object.

Return type:

DocGeneratorModule

graphdoc.doc_generator_eval_from_yaml(yaml_path: str | Path) → DocGeneratorEvaluator

Load a doc generator evaluator from a YAML file.

mlflow:
    mlflow_tracking_uri:      !env MLFLOW_TRACKING_URI      # The tracking URI for MLflow
    mlflow_tracking_username: !env MLFLOW_TRACKING_USERNAME # The username for the mlflow tracking server
    mlflow_tracking_password: !env MLFLOW_TRACKING_PASSWORD # The password for the mlflow tracking server

prompt:
    prompt: base_doc_gen                                  # Which prompt signature to use
    class: DocGeneratorPrompt                             # Must be a child of SinglePrompt (we will use an enum to map this)
    type: chain_of_thought                                # The type of prompt to use (predict, chain_of_thought)
    metric: rating                                        # The type of metric to use (rating, category)
    load_from_mlflow: false                               # Whether to load the prompt from an MLFlow URI
    model_uri: null                                       # The tracking URI for MLflow
    model_name: null                                      # The name of the model in MLflow
    model_version: null                                   # The version of the model in MLflow
    prompt_metric: true                                   # Whether another prompt is used to calculate the metric (in which case we must also load that prompt)

prompt_metric:
    prompt: doc_quality                                   # The prompt to use to calculate the metric
    class: DocQualityPrompt                               # The class of the prompt to use to calculate the metric
    type: predict                                         # The type of prompt to use to calculate the metric
    metric: rating                                        # The metric to use to calculate the metric
    load_from_mlflow: false                               # Whether to load the prompt from an MLFlow URI
    model_uri: null                                       # The tracking URI for MLflow
    model_name: null                                      # The name of the model in MLflow
    model_version: null

module:
    retry: true                                           # Whether to retry the generation if the quality check fails
    retry_limit: 1                                        # The maximum number of retries
    rating_threshold: 3                                   # The rating threshold for the quality check
    fill_empty_descriptions: true                         # Whether to fill the empty descriptions in the schema

eval:
    mlflow_experiment_name: doc_generator_eval            # The name of the experiment in MLflow
    generator_prediction_field: documented_schema         # The field in the generator prediction to use
    evaluator_prediction_field: rating                    # The field in the evaluator prediction to use
    readable_value: 25

Parameters:

yaml_path (Union[str, Path]) – Path to the YAML file.

Returns:

A DocGeneratorEvaluator object.

Return type:

DocGeneratorEvaluator

Subpackages

• graphdoc.data package
  • Subpackages
    • graphdoc.data.dspy_data package
      • Submodules
    • graphdoc.data.mlflow_data package
      • Submodules
  • Submodules
    • graphdoc.data.helper module
    • graphdoc.data.local module
    • graphdoc.data.parser module
    • graphdoc.data.schema module
• graphdoc.eval package
  • Submodules
    • graphdoc.eval.doc_generator_eval module
• graphdoc.modules package
  • Submodules
    • graphdoc.modules.doc_generator_module module
• graphdoc.prompts package
  • Submodules
    • graphdoc.prompts.schema_doc_generation module
    • graphdoc.prompts.schema_doc_quality module
    • graphdoc.prompts.single_prompt module
• graphdoc.train package
  • Submodules
    • graphdoc.train.doc_generator_trainer module
    • graphdoc.train.doc_quality_trainer module
    • graphdoc.train.optimizers module
    • graphdoc.train.single_prompt_trainer module

Submodules

• graphdoc.config module
• graphdoc.main module