diff --git a/examples/python/code_graph_example.py b/examples/python/code_graph_example.py
index 44ab33aad..9cd9f99c4 100644
--- a/examples/python/code_graph_example.py
+++ b/examples/python/code_graph_example.py
@@ -8,9 +8,61 @@ async def main(repo_path, include_docs):
     async for result in run_code_graph_pipeline(repo_path, include_docs):
         print(result)
 
-if __name__ == "__main__":
+def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--repo_path", type=str, 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")
-    args = parser.parse_args()
-    asyncio.run(main(args.repo_path, args.include_docs))
\ No newline at end of file
+    parser.add_argument(
+        "--repo_path", 
+        type=str, 
+        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))
+        
\ No newline at end of file