81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
import dspy
|
|
import os
|
|
from typing import Any
|
|
from modaic import PrecompiledConfig, PrecompiledProgram
|
|
|
|
lm = dspy.LM("openai/gpt-5-mini", max_tokens=16000)
|
|
lm_mini = dspy.LM("openai/gpt-5-mini", max_tokens=16000)
|
|
dspy.configure(lm=lm)
|
|
|
|
|
|
class DocWriter(dspy.Signature):
|
|
"""
|
|
Write documentation for the provided Rails application source code.
|
|
"""
|
|
|
|
source_tree: dict[str, Any] = dspy.InputField()
|
|
documentation: str = dspy.OutputField(
|
|
description="Generated markdown documentation."
|
|
)
|
|
|
|
|
|
class DocWriterConfig(PrecompiledConfig):
|
|
max_iterations: int = 10
|
|
lm: str = "openai/gpt-5-mini"
|
|
sub_lm: str = "openai/gpt-5-mini"
|
|
|
|
|
|
class DocWriterProgram(PrecompiledProgram):
|
|
config: DocWriterConfig
|
|
|
|
def __init__(self, config: DocWriterConfig, **kwargs):
|
|
super().__init__(config=config, **kwargs)
|
|
self.doc_writer = dspy.RLM(
|
|
DocWriter,
|
|
max_iterations=config.max_iterations,
|
|
sub_lm=config.sub_lm,
|
|
verbose=True,
|
|
)
|
|
|
|
def forward(self, source_root: str) -> dspy.Prediction:
|
|
source_tree = load_source_tree(source_root)
|
|
print(f"Loaded source tree with {len(source_tree)} entries")
|
|
return self.doc_writer(source_tree=source_tree)
|
|
|
|
|
|
def load_source_tree(root_dir: str) -> dict[str, Any]:
|
|
"""Recursively load the folder into a nested dict."""
|
|
tree: dict[str, Any] = {}
|
|
for entry in os.listdir(root_dir):
|
|
path = os.path.join(root_dir, entry)
|
|
if os.path.isdir(path):
|
|
tree[entry] = load_source_tree(path)
|
|
else:
|
|
with open(path, "r", encoding="utf-8", errors="ignore") as f:
|
|
tree[entry] = f.read()
|
|
return tree
|
|
|
|
|
|
doc_writer = DocWriterProgram(DocWriterConfig())
|
|
SOURCE_ROOT = "."
|
|
|
|
|
|
def main():
|
|
"""
|
|
print("Starting documentation generation...")
|
|
result = doc_writer(source_root=SOURCE_ROOT)
|
|
|
|
with open("generated_documentation.md", "w", encoding="utf-8") as f:
|
|
print("Writing documentation to file...")
|
|
f.write(result.documentation)
|
|
print("Documentation written to generated_documentation.md")
|
|
|
|
print("Pushing documentation to hub...")
|
|
"""
|
|
doc_writer.push_to_hub("farouk1/ruby-rails-doc-generator")
|
|
print("Documentation generated and pushed to hub!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|