mirror of
https://github.com/langchain-ai/langsmith-docs.git
synced 2026-08-27 01:41:19 -04:00
William Fu-Hinthorn
111 lines
3.5 KiB
Python
111 lines
3.5 KiB
Python
import re
|
|
from pathlib import Path
|
|
import ast
|
|
|
|
class ReplaceListRunsVisitor(ast.NodeTransformer):
|
|
def visit_Call(self, node):
|
|
self.generic_visit(node) # First, apply transformation to children
|
|
if isinstance(node.func, ast.Attribute) and node.func.attr == "list_runs":
|
|
# Wrap the original call in `list(...)`
|
|
return ast.Call(
|
|
func=ast.Name(id="list", ctx=ast.Load()),
|
|
args=[node],
|
|
keywords=[]
|
|
)
|
|
return node
|
|
|
|
def _call_list_on_list_runs(code_block: str):
|
|
tree = ast.parse(code_block)
|
|
transformer = ReplaceListRunsVisitor()
|
|
transformed_tree = transformer.visit(tree)
|
|
|
|
return ast.unparse(transformed_tree)
|
|
|
|
|
|
|
|
def extract_code_blocks(mdx_file: str) -> list[str]:
|
|
"""
|
|
Extract all code blocks from an MDX file.
|
|
|
|
:param mdx_file: Path to the MDX file
|
|
:return: A list of code blocks
|
|
"""
|
|
with open(mdx_file, "r") as file:
|
|
content = file.read()
|
|
# Pattern to match code inside PythonBlock(\``) from code tabs
|
|
code_blocks = re.findall(r"PythonBlock\(`(.*?)`(?:,|\))", content, re.DOTALL)
|
|
# Assuming code blocks are enclosed between triple backticks
|
|
code_blocks += [
|
|
c.strip() for c in re.findall(r"```python\n(.*?)\n```", content, re.DOTALL)
|
|
]
|
|
code_blocks = [
|
|
_call_list_on_list_runs(
|
|
code_block.replace("\\n", "\n").replace("\\\\", "\\")
|
|
)
|
|
for code_block in code_blocks
|
|
# Skip because we don't have the actual UUID to run in the test.
|
|
if "<run_id>" not in code_block
|
|
and "<your_project>" not in code_block
|
|
and "create_dataset" not in code_block
|
|
]
|
|
|
|
return code_blocks
|
|
|
|
|
|
def write_pytest_tests(
|
|
code_blocks: list[str], output_file: str, boilerplate: str = ""
|
|
) -> None:
|
|
"""
|
|
Write pytest unit tests to a Python file for each code block.
|
|
|
|
:param code_blocks: List of code blocks to test
|
|
:param output_file: Path to the output Python file
|
|
"""
|
|
with open(output_file, "w") as file:
|
|
file.write("import pytest\n")
|
|
file.write(boilerplate)
|
|
for index, code_block in enumerate(code_blocks):
|
|
indented_code = "\n".join(
|
|
[" " + line for line in code_block.split("\n")]
|
|
)
|
|
test_code = f"""
|
|
@pytest.mark.asyncio
|
|
async def test_code_block_{index}():
|
|
{indented_code}
|
|
"""
|
|
file.write(test_code)
|
|
|
|
|
|
def main(mdx_file: str, output_file: str, boilerplate: str = "") -> None:
|
|
"""
|
|
Main function to extract code blocks from MDX and write to pytest file.
|
|
|
|
:param mdx_file: Path to the MDX file
|
|
:param output_file: Path to the output Python file
|
|
"""
|
|
code_blocks = extract_code_blocks(mdx_file)
|
|
write_pytest_tests(code_blocks, output_file, boilerplate=boilerplate)
|
|
print(f"Tests written to {output_file}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
root = Path(__file__).parent.parent.absolute()
|
|
files = [
|
|
(
|
|
"tracing/faq/querying_traces.mdx",
|
|
"""from langsmith import Client
|
|
|
|
client = Client()
|
|
""",
|
|
),
|
|
]
|
|
for file, boilerplate in files:
|
|
stem = Path(file).stem
|
|
dest_folder = root / "tests/py_unit_tests" / Path(file).parent
|
|
dest_folder.mkdir(parents=True, exist_ok=True)
|
|
main(
|
|
str(root / "docs" / file),
|
|
str(dest_folder / f"test_{stem}.py"),
|
|
boilerplate=boilerplate,
|
|
)
|