gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: time code graph run and add mock support

Browse source
This commit is contained in:
lxobr 2025-01-06 11:25:04 +01:00
parent dbc33a6478
commit 5e79dc53c5
1 changed files with 57 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
examples/python/code_graph_example.py
View file

@ -8,9 +8,61 @@ async def main(repo_path, include_docs):
async for result in run_code_graph_pipeline(repo_path, include_docs): async for result in run_code_graph_pipeline(repo_path, include_docs):
print(result) print(result)
if __name__ == "__main__": def parse_args():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--repo_path", type=str, required=True, help="Path to the repository") parser.add_argument(
parser.add_argument("--include_docs", type=lambda x: x.lower() in ("true", "1"), default=True, help="Whether or not to process non-code files") "--repo_path",
args = parser.parse_args() type=str,
asyncio.run(main(args.repo_path, args.include_docs)) required=True,
help="Path to the repository"
)
parser.add_argument(
"--include_docs",
type=lambda x: x.lower() in ("true", "1"),
default=True,
help="Whether or not to process non-code files"
)
parser.add_argument(
"--mock_embedding",
type=lambda x: x.lower() in ("true", "1"),
default=True,
help="Whether or not to mock embedding and code summary"
)
parser.add_argument(
"--mock_code_summary",
type=lambda x: x.lower() in ("true", "1"),
default=True,
help="Whether or not to mock code summary"
)
parser.add_argument(
"--time",
type=lambda x: x.lower() in ("true", "1"),
default=True,
help="Whether or not to time the pipeline run"
)
return parser.parse_args()
if __name__ == "__main__":
import os
args = parse_args()
if args.mock_embedding:
os.environ["MOCK_EMBEDDING"] = "true"
print("Mocking embedding.")
if args.mock_code_summary:
os.environ["MOCK_CODE_SUMMARY"] = "true"
print("Mocking code summary.")
if args.time:
import time
start_time = time.time()
asyncio.run(main(args.repo_path, args.include_docs))
end_time = time.time()
print("\n" + "="*50)
print(f"Pipeline Execution Time: {end_time - start_time:.2f} seconds")
print("="*50 + "\n")
else:
asyncio.run(main(args.repo_path, args.include_docs))