graphdoc.modules package

class graphdoc.modules.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)[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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.modules.TokenTracker[source]

Bases: object

A class to track the number of tokens used.

clear()[source]

Clear the token tracker.

stats()[source]

Get the stats of the token tracker.

global_token_callback(kwargs, response, start_time, end_time, **callback_kwargs)[source]

A global callback to track the number of tokens used.

Intended to be used with the litellm ModelResponse object.

Submodules