Python Version and function references as additional context

[aseembits93]

PR Type

Enhancement

---

Description

• Pass function references across optimization flow

• Include python version in service payloads

• Extend models and APIs with new fields

• Adjust thread pool sizing for concurrency

---

Diagram Walkthrough

flowchart LR
  gen["Generate tests & refs"] -- "returns function_references" --> opt["Optimize function"]
  opt -- "selects best candidate" --> refine["Refine optimizations"]
  refine -- "payload includes function_references" --> aiRef["AI refine endpoint"]
  opt -- "request ranking with refs" --> rank["AI ranking endpoint"]
  review["Build optimization review"] -- "includes python_version, refs" --> aiRev["AI review endpoint"]

Loading

File Walkthrough

Relevant files

Enhancement

aiservice.py Enrich AI service payloads with refs and version                  codeflash/api/aiservice.py Add function_references to refinement payload Add python_version to multiple requests Extend get_new_explanation and generate_ranking signatures Forward function_references in payloads +14/-1    models.py Model updated for function references                                        codeflash/models/models.py Extend AIServiceRefinerRequest with function_references +1/-0      function_optimizer.py Propagate function references through optimizer workflow  codeflash/optimization/function_optimizer.py Compute and propagate function_references end-to-end Include refs in ranking, refinement, review, and explanation Return refs from generation results Increase thread pool max workers by 1 +37/-16

---

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

Codeflash Bot seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Type Consistency New parameter and return types for function_references vary between str and str | None; ensure upstream producers and downstream consumers agree to avoid runtime errors or None-handling gaps. def find_and_process_best_optimization( self, optimizations_set: OptimizationSet, code_context: CodeOptimizationContext, original_code_baseline: OriginalCodeBaseline, original_helper_code: dict[Path, str], file_path_to_helper_classes: dict[Path, set[str]], function_to_optimize_qualified_name: str, function_to_all_tests: dict[str, set[FunctionCalledInTest]], generated_tests: GeneratedTestsList, test_functions_to_remove: list[str], concolic_test_str: str | None, function_references: str, ) -> BestOptimization | None: Concurrency Change Increasing thread pool workers by +1 may affect resource usage and scheduling; validate that this does not oversubscribe environments with many parallel optimizations. self.function_benchmark_timings = function_benchmark_timings if function_benchmark_timings else {} self.total_benchmark_timings = total_benchmark_timings if total_benchmark_timings else {} self.replay_tests_dir = replay_tests_dir if replay_tests_dir else None self.generate_and_instrument_tests_results: ( tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet] | None ) = None n_tests = N_TESTS_TO_GENERATE_EFFECTIVE self.executor = concurrent.futures.ThreadPoolExecutor( max_workers=n_tests + 3 if self.experiment_id is None else n_tests + 4 ) API Contract Payloads now include python_version and function_references; confirm server endpoints accept these fields and that function_references optionality is handled to prevent 4xx on missing/None values. """ payload = { "trace_id": trace_id, "diffs": diffs, "speedups": speedups, "optimization_ids": optimization_ids, "python_version": platform.python_version(), "function_references": function_references, }

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Pass correct argument variable The get_optimization_review call was changed to pass function_references under calling_fn_details. Ensure the variable name on the left matches the new data (function_references) to avoid passing an undefined or stale symbol. This prevents runtime NameError or incorrect payload fields. codeflash/optimization/function_optimizer.py [1498-1500] -calling_fn_details = calling_fn_details, +calling_fn_details=function_references Suggestion importance[1-10]: 8 __ Why: The PR updates the review call to use function_references for calling_fn_details, and ensuring the named argument is calling_fn_details=function_references prevents runtime errors and incorrect payloads; this is a correctness fix with high impact.	Medium
	Align tuple unpacking order The order of values unpacked here must exactly match the tuple returned by generate_tests_and_optimizations. After adding function_references, ensure both producer and consumer use the same position to avoid misassigned variables at runtime. Align the unpacking order with the corrected generate_tests_and_optimizations return structure. codeflash/optimization/function_optimizer.py [327-334] +( + count_tests, + generated_tests, + function_to_concolic_tests, + concolic_test_str, + optimizations_set, + function_references, +) = generated_results.unwrap() - Suggestion importance[1-10]: 5 __ Why: Ensuring the unpacking order matches the producer is important; in this PR the order is already correct, so the suggestion is essentially a no-op reminder with limited impact but accurate.	Low
	Fix incorrect return annotation The declared return type places the error type in the middle of the type list, which does not match the actual Success tuple structure built later. Adjust the annotation so the Result success type includes both the OptimizationSet and function_references, and the error type is a single str. This prevents runtime type mismatches and static typing confusion. codeflash/optimization/function_optimizer.py [1097-1106] def generate_tests_and_optimizations( self, testgen_context: CodeStringsMarkdown, read_writable_code: CodeStringsMarkdown, read_only_context_code: str, helper_functions: list[FunctionSource], generated_test_paths: list[Path], generated_perf_test_paths: list[Path], run_experiment: bool = False, # noqa: FBT001, FBT002 -) -> Result[tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet], str, str]: +) -> Result[ + tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], str, OptimizationSet, str], + str +]: Suggestion importance[1-10]: 3 __ Why: The suggestion rightly flags a potential mismatch in the Result type annotation after adding function_references, but the proposed improved signature does not match the new Success tuple built in the PR and reorders elements incorrectly, risking further confusion.	Low

[KRRT7]

what's this needed for?

[aseembits93]

@KRRT7 for 3.9 unittests, we're using '|' for type hints which is not supported in 3.9, from the logs

TypeError: Unable to evaluate type annotation 'str | None'. If you are making use of the new typing syntax (unions using | since Python 3.10 or builtins subscripting since Python 3.9), you should either replace the use of new syntax with the existing typing constructs or install the eval_type_backport package.```

[aseembits93]

https://github.com/codeflash-ai/codeflash/actions/runs/19057649760/job/54441517008?pr=874 @KRRT7

[misrasaurabh1]

we don't use these types of type hints in our cli code for the same reason. use Union

[KRRT7]

a couple of your commits show up as codeflash bot, this happens when Claude or other LLMs commit and mess with the git config, that's why we got the CLA bot commenting, we need to fix that before being able to merge