From 4bfae21bce78eddfed70207e6a397423d9aade35 Mon Sep 17 00:00:00 2001 From: Vidhu Mathur Date: Mon, 1 Jul 2024 01:08:58 +0530 Subject: [PATCH 1/6] chore: add support for more context. *WIP --- optimizeai/decorators/optimize.py | 29 ++++++++++++++++++++++++----- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/optimizeai/decorators/optimize.py b/optimizeai/decorators/optimize.py index b7ace73..6b67ac7 100644 --- a/optimizeai/decorators/optimize.py +++ b/optimizeai/decorators/optimize.py @@ -1,6 +1,7 @@ """This module contains the optimize decorator for optimizing Python functions using LLMs and performance profiling.""" +import os import inspect import functools from io import StringIO @@ -9,16 +10,32 @@ from optimizeai.llm_wrapper import LLMWrapper from optimizeai.config import Config -# Custom optimize decorator +def get_function_code(func): + """Get the source code of a function if it's defined in the current directory.""" + try: + source_lines, _ = inspect.getsourcelines(func) + source_file = inspect.getfile(func) + current_directory = os.path.dirname(os.path.abspath(__file__)) + if source_file.startswith(current_directory): + return ''.join(source_lines) + except (TypeError, OSError): + return None + def optimize(profiler_types, config: Config): """Decorator to optimize a Python function using LLMs and performance profiling.""" def decorator(func): @functools.wraps(func) def wrapper(*args, **kwargs): # Capture the executed code - source_lines, _ = inspect.getsourcelines(func) - # Remove the first line (which is the decorator line) - code = ''.join(source_lines[1:]) + code_lines, _ = inspect.getsourcelines(func) + code = ''.join(code_lines[1:]) # Remove the decorator line + + # Find and include the code of all called functions + called_functions = inspect.getmembers(func, predicate=inspect.isfunction) + for name, called_func in called_functions: + func_code = get_function_code(called_func) + if func_code: + code += f'\n\n# Function {name}\n' + func_code # Profile the function and capture the output with StringIO() as buf, redirect_stdout(buf): @@ -27,10 +44,12 @@ def wrapper(*args, **kwargs): # Initialize the LLMWrapper with the provided config llm_wrapper = LLMWrapper(config) - response= llm_wrapper.send_request(code = code, perf_metrics=captured_output) + response = llm_wrapper.send_request(code=code, perf_metrics=captured_output) print(response) return response return wrapper return decorator + + From dd5f4718f60c033968abc75d12ad9d3628f344b1 Mon Sep 17 00:00:00 2001 From: Vidhu Mathur Date: Sun, 7 Jul 2024 15:48:29 +0530 Subject: [PATCH 2/6] update: google genai package --- poetry.lock | 14 +++++++------- pyproject.toml | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/poetry.lock b/poetry.lock index 1b8d7fe..ee41d1f 100644 --- a/poetry.lock +++ b/poetry.lock @@ -701,13 +701,13 @@ beautifulsoup4 = "*" [[package]] name = "google-ai-generativelanguage" -version = "0.6.4" +version = "0.6.6" description = "Google Ai Generativelanguage API client library" optional = false python-versions = ">=3.7" files = [ - {file = "google-ai-generativelanguage-0.6.4.tar.gz", hash = "sha256:1750848c12af96cb24ae1c3dd05e4bfe24867dc4577009ed03e1042d8421e874"}, - {file = "google_ai_generativelanguage-0.6.4-py3-none-any.whl", hash = "sha256:730e471aa549797118fb1c88421ba1957741433ada575cf5dd08d3aebf903ab1"}, + {file = "google-ai-generativelanguage-0.6.6.tar.gz", hash = "sha256:1739f035caeeeca5c28f887405eec8690f3372daf79fecf26454a97a4f1733a8"}, + {file = "google_ai_generativelanguage-0.6.6-py3-none-any.whl", hash = "sha256:59297737931f073d55ce1268dcc6d95111ee62850349d2b6cde942b16a4fca5c"}, ] [package.dependencies] @@ -799,16 +799,16 @@ httplib2 = ">=0.19.0" [[package]] name = "google-generativeai" -version = "0.6.0" +version = "0.7.1" description = "Google Generative AI High level API client library and tools." optional = false python-versions = ">=3.9" files = [ - {file = "google_generativeai-0.6.0-py3-none-any.whl", hash = "sha256:ba1d3b826b872bffe330aaac0dc6de2f0e4610df861c8ce7ec6433771611b676"}, + {file = "google_generativeai-0.7.1-py3-none-any.whl", hash = "sha256:25017b1278c873b6db65ba41c70566ee9fce8f265498e8b2f6aac036528ba1c7"}, ] [package.dependencies] -google-ai-generativelanguage = "0.6.4" +google-ai-generativelanguage = "0.6.6" google-api-core = "*" google-api-python-client = "*" google-auth = ">=2.15.0" @@ -3564,4 +3564,4 @@ multidict = ">=4.0" [metadata] lock-version = "2.0" python-versions = "^3.11" -content-hash = "654112a79fb4fdbbfa173d280ec9b6f860b4329a035f614fa51c00cec4932c58" +content-hash = "7cfa6c5ad9e89f238504037464b234fa73542be986e5b8c7ce5403030ba9a938" diff --git a/pyproject.toml b/pyproject.toml index 231d76a..bb43ff1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ perfwatch = "^1.3.2" py-context = "^0.3.1" requests = "^2.32.3" transformers = "^4.41.2" -google-generativeai = "^0.6.0" +google-generativeai = "^0.7.1" dspy = "^0.1.5" anthropic = "^0.28.1" torch = "^2.3.1" From 5e868494ebaa08e04d3b5a0fd617574f2cb7b238 Mon Sep 17 00:00:00 2001 From: Vidhu Mathur Date: Sun, 7 Jul 2024 15:48:49 +0530 Subject: [PATCH 3/6] Add: context field to dspy --- optimizeai/cot.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/optimizeai/cot.py b/optimizeai/cot.py index 575d7ae..f9ac0ce 100644 --- a/optimizeai/cot.py +++ b/optimizeai/cot.py @@ -8,7 +8,7 @@ def __init__(self, signature: dspy.Signature): super().__init__() self.prog = dspy.ChainOfThought(signature) - def forward(self, code, perf_metrics): + def forward(self, code, context, perf_metrics): """Forward method to pass the code and performance metrics to the Chain of Thought model.""" - answer = self.prog(code=code, perf_metrics=perf_metrics) + answer = self.prog(code=code, context=context, perf_metrics=perf_metrics) return answer \ No newline at end of file From 3d5087f634a157da2623daa7537fd3028b8452da Mon Sep 17 00:00:00 2001 From: Vidhu Mathur Date: Sun, 7 Jul 2024 15:49:54 +0530 Subject: [PATCH 4/6] add: functionality to send relevant context with code --- optimizeai/decorators/optimize.py | 65 ++++++++++++++++++++----------- 1 file changed, 42 insertions(+), 23 deletions(-) diff --git a/optimizeai/decorators/optimize.py b/optimizeai/decorators/optimize.py index 6b67ac7..9e82dd4 100644 --- a/optimizeai/decorators/optimize.py +++ b/optimizeai/decorators/optimize.py @@ -1,25 +1,32 @@ """This module contains the optimize decorator for optimizing Python functions using LLMs and performance profiling.""" -import os import inspect import functools from io import StringIO from contextlib import redirect_stdout +import sys +import os from perfwatch import watch from optimizeai.llm_wrapper import LLMWrapper from optimizeai.config import Config +import types def get_function_code(func): - """Get the source code of a function if it's defined in the current directory.""" + """Retrieve the source code of a function.""" try: source_lines, _ = inspect.getsourcelines(func) - source_file = inspect.getfile(func) - current_directory = os.path.dirname(os.path.abspath(__file__)) - if source_file.startswith(current_directory): - return ''.join(source_lines) - except (TypeError, OSError): - return None + return ''.join(source_lines) + except (IOError, TypeError): + return f"Source code not available for {func.__name__}" + +def is_user_defined_function(func, base_folder): + """Check if a function is user-defined.""" + if isinstance(func, types.FunctionType): + func_file = inspect.getfile(func) + # Check if the function file path starts with the base folder path + return func_file.startswith(base_folder) + return False def optimize(profiler_types, config: Config): """Decorator to optimize a Python function using LLMs and performance profiling.""" @@ -27,29 +34,41 @@ def decorator(func): @functools.wraps(func) def wrapper(*args, **kwargs): # Capture the executed code - code_lines, _ = inspect.getsourcelines(func) - code = ''.join(code_lines[1:]) # Remove the decorator line - - # Find and include the code of all called functions - called_functions = inspect.getmembers(func, predicate=inspect.isfunction) - for name, called_func in called_functions: - func_code = get_function_code(called_func) - if func_code: - code += f'\n\n# Function {name}\n' + func_code + source_lines, _ = inspect.getsourcelines(func) + # Remove the first line (which is the decorator line) + code = ''.join(source_lines[1:]) # Profile the function and capture the output with StringIO() as buf, redirect_stdout(buf): - watch(profiler_types)(func)(*args, **kwargs) - captured_output = buf.getvalue() + # Create a dictionary to store called functions' source codes + called_funcs_code = {} + + def trace_calls(frame, event, arg): + if event == 'call': + called_func = frame.f_globals.get(frame.f_code.co_name) + base_folder = os.path.dirname(inspect.getfile(func)) + if is_user_defined_function(called_func, base_folder): + called_funcs_code[frame.f_code.co_name] = get_function_code(called_func) + return trace_calls + # Set the trace function + sys.settrace(trace_calls) + + try: + watch(profiler_types)(func)(*args, **kwargs) + finally: + # Remove the trace function + sys.settrace(None) + + captured_output = buf.getvalue() + # print(str(list(called_funcs_code.values())[1:])) + # print(code) + # print(captured_output) # Initialize the LLMWrapper with the provided config llm_wrapper = LLMWrapper(config) - response = llm_wrapper.send_request(code=code, perf_metrics=captured_output) - + response = llm_wrapper.send_request(code=str(code), context=str(list(called_funcs_code.values())[1:]), perf_metrics=captured_output) print(response) return response return wrapper return decorator - - From 4145e4363470b0a830394ba60108cb58d1d5cc36 Mon Sep 17 00:00:00 2001 From: Vidhu Mathur Date: Sun, 7 Jul 2024 15:50:32 +0530 Subject: [PATCH 5/6] update: dspy signature to support improved context --- optimizeai/llm_wrapper.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/optimizeai/llm_wrapper.py b/optimizeai/llm_wrapper.py index e3367ed..6ca806f 100644 --- a/optimizeai/llm_wrapper.py +++ b/optimizeai/llm_wrapper.py @@ -10,6 +10,7 @@ class ZeroShotQAWithCoT(dspy.Signature): Provide short detailed technical tips and refactor suggestions, such as algorithm improvements, data structure optimizations, or parallelization strategies if needed.""" code = dspy.InputField(desc="Code to be optimized") + context = dspy.InputField(desc="Code for all user defined functions in the original code provided") perf_metrics = dspy.InputField(desc="""Performance metrics of the code along with the output of the code execution""") optimization = dspy.OutputField(desc="""Detailed solution to how the code can be optimized with insights on performance metrics @@ -43,10 +44,10 @@ def __setup_llm(self): dspy.settings.configure(lm=self.llm) self.chain = CoT(ZeroShotQAWithCoT) - def send_request(self, code, perf_metrics): + def send_request(self, code, context, perf_metrics): """Send a request to the LLM model with the given prompt and performance metrics. Args: code (str): The code to send to the LLM model. perf_metrics (str): The performance metrics to send to the LLM model. Returns: answer (str): The answer generated by the LLM model.""" - answer = self.chain.forward(code=code, perf_metrics=perf_metrics) + answer = self.chain.forward(code=code, context=context, perf_metrics=perf_metrics) return answer.optimization From 0c8d8f43e7a480d6737331192753752da17e6952 Mon Sep 17 00:00:00 2001 From: Vidhu Mathur Date: Sun, 7 Jul 2024 15:51:18 +0530 Subject: [PATCH 6/6] version change --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index bb43ff1..f2b443f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "optimizeai" -version = "0.1.4" +version = "0.2.0" description = "OptimAI is a powerful Python module designed to optimize your code by analyzing its performance and providing actionable suggestions. It leverages a large language model (LLM) to give you detailed insights and recommendations based on the profiling data collected during the execution of your code." authors = ["Vidhu Mathur "] readme = "README.md"