code generation chain added

Author: zilto

contrib/hamilton/contrib/user/zilto/llm_generate_code/README.md

@@ -0,0 +1,18 @@
+# Purpose of this module
+
+This module uses the OpenAI completion API to generate code.
+
+For any language, you can request `generated_code` to get the generated response. If you are generating Python code, you can execute it in a subprocess by requesting `execution_output` and `execution_error`.
+
+# Configuration Options
+## Config.when
+This module doesn't receive configurations.
+
+## Inputs
+- `query`: The query for which you want code generated.
+- `api_key`: Set the OpenAI API key to use. If None, read the environment variable `OPENAI_API_KEY`
+- `code_language`: Set the code language to generate the reponse in. Defaults to `python`
+
+## Overrides
+- `prompt_template_to_generate_code`: Create a new prompt template with the fields `query` and `code_language`.
+- `prompt_to_generate_code`: Manually provide a prompt to generate Python code

contrib/hamilton/contrib/user/zilto/llm_generate_code/__init__.py

@@ -0,0 +1,107 @@
+import logging
+import os
+import subprocess
+from typing import Optional
+
+from hamilton.function_modifiers import extract_fields
+
+logger = logging.getLogger(__name__)
+
+from hamilton import contrib
+
+with contrib.catch_import_errors(__name__, __file__, logger):
+    import openai
+
+
+def llm_client(api_key: Optional[str] = None) -> openai.OpenAI:
+    """Create an OpenAI client"""
+    if api_key is None:
+        api_key = os.environ.get("OPENAI_API_KEY")
+
+    return openai.OpenAI(api_key=api_key)
+
+
+def prompt_template_to_generate_code() -> str:
+    return """Write some {code_language} code to solve the user's problem.
+
+Return only python code in Markdown format, e.g.:
+
+```{code_language}
+....
+```
+
+user problem
+{query}
+
+{code_language} code
+"""
+
+
+def prompt_to_generate_code(
+    prompt_template_to_generate_code: str, query: str, code_language: str = "python"
+) -> str:
+    return prompt_template_to_generate_code.format(
+        query=query,
+        code_language=code_language,
+    )
+
+
+def response_generated_code(llm_client: openai.OpenAI, prompt_to_generate_code: str) -> str:
+    response = llm_client.completions.create(
+        model="gpt-3.5-turbo-instruct",
+        prompt=prompt_to_generate_code,
+    )
+    return response.choices[0].text
+
+
+def generated_code(response_generated_code: str) -> str:
+    _, _, lower_part = response_generated_code.partition("```python")
+    code_part, _, _ = lower_part.partition("```")
+    return code_part
+
+
+def code_prepared_for_execution(generated_code: str, code_language: str = "python") -> str:
+    if code_language != "python":
+        raise ValueError("Can only execute the generated code if `code_language` = 'python'")
+
+    code_to_get_vars = """
+excluded_vars = { 'excluded_vars', '__builtins__', '__annotations__'} | set(dir(__builtins__))
+local_vars = {k:v for k,v in locals().items() if k not in excluded_vars}
+print(local_vars)
+"""
+    return generated_code + code_to_get_vars
+
+
+@extract_fields(
+    dict(
+        execution_output=str,
+        execution_error=str,
+    )
+)
+def execute_output(code_prepared_for_execution: str) -> dict:
+    process = subprocess.Popen(
+        ["python", "-c", code_prepared_for_execution],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        universal_newlines=True,
+    )
+    output, errors = process.communicate()
+    return dict(execution_output=output, execution_error=errors)
+
+
+# run as a script to test dataflow
+if __name__ == "__main__":
+    import __init__ as llm_generate_code
+
+    from hamilton import driver
+
+    dr = driver.Builder().with_modules(llm_generate_code).build()
+
+    dr.display_all_functions("dag.png", orient="TB")
+
+    res = dr.execute(
+        ["execution_output", "execution_error"],
+        overrides=dict(generated_code="s = 'hello world'"),
+    )
+
+    print(res)

contrib/hamilton/contrib/user/zilto/llm_generate_code/dag.png

@@ -0,0 +1 @@
+openai

contrib/hamilton/contrib/user/zilto/llm_generate_code/requirements.txt

@@ -0,0 +1,5 @@
+{
+  "schema": "1.0",
+  "use_case_tags": ["LLM", "OpenAI"],
+  "secondary_tags": {}
+}