From e66808790d5be0a4c26719765e8ed365c22f29fe Mon Sep 17 00:00:00 2001 From: Brendan Date: Tue, 16 Apr 2024 17:57:16 -0700 Subject: [PATCH 1/3] feat: added functions to add command and retain order --- mirascope_cli/generic/prompt_template.j2 | 54 +++----- mirascope_cli/schemas.py | 17 ++- mirascope_cli/utils.py | 125 +++++++++++++----- .../golden/base_prompt/0001_base_prompt.py | 1 - .../golden/base_prompt/0002_base_prompt.py | 1 - .../0001_base_prompt_auto_tag.py | 1 - .../0002_base_prompt_auto_tag.py | 1 - .../0001_base_prompt_with_decorator.py | 1 - .../0002_base_prompt_with_decorator.py | 1 - .../0001_call_with_variables.py | 15 ++- .../0002_call_with_variables.py | 15 ++- .../call_with_variables.py | 10 +- tests/commands/test_add.py | 5 +- 13 files changed, 167 insertions(+), 80 deletions(-) diff --git a/mirascope_cli/generic/prompt_template.j2 b/mirascope_cli/generic/prompt_template.j2 index d518e27..28cf815 100644 --- a/mirascope_cli/generic/prompt_template.j2 +++ b/mirascope_cli/generic/prompt_template.j2 @@ -1,36 +1,15 @@ -{%- if comments -%} -"""{{ comments }}""" -{%- endif -%} -{%- if imports -%} -{%- for import, alias in imports -%} -{%- if alias %} -import {{ import }} as {{alias}} -{%- else %} -import {{ import }} -{%- endif -%} -{%- endfor -%} -{%- endif %} - -{% if from_imports -%} -{%- set from_import_groups = {} -%} -{%- for module, name, alias in from_imports -%} - {% if module not in from_import_groups %} - {%- set _ = from_import_groups.update({module: [(name, alias)]}) -%} - {%- else -%} - {%- set _ = from_import_groups[module].append((name, alias)) -%} - {%- endif -%} -{%- endfor -%} -{% for module, names in from_import_groups.items() -%} -from {{ module }} import {% for name, alias in names %}{% if alias %}{{ name }} as {{ alias }}{% else %}{{ name }}{% endif %}{% if not loop.last %}, {% endif %}{% endfor %} -{% endfor %} -{% endif -%} - -{% if variables %} -{%- for var_name, var_value in variables.items() -%} +{%- for item in order -%} +{%- if item.type == "comment" %} +"""{{ item.render }}""" +{%- elif item.type == "import" %} +import {{ item.render[0] }}{% if item.render[1] %} as {{ item.render[1] }}{% endif %} +{%- elif item.type == "from_import" %} +from {{ item.render[0] }} import {% if item.render[2] %}{{ item.render[1] }} as {{ item.render[2] }}{% else %}{{ item.render[1] }}{% endif %} +{%- elif item.type == "variable" %} +{%- set var_name, var_value = item.render %} {{ var_name }} = {{ var_value }} -{% endfor %} -{% endif -%} -{% for class in classes -%} +{%- elif item.type == "class" %} +{%- set class = item.render -%} {%- for decorator in class.decorators %} @{{ decorator }} {%- endfor %} @@ -50,4 +29,15 @@ class {{ class.name }}({{ class.bases | join(', ') }}): {%- endif %} {{ line }} {%- endfor %} +{%- elif item.type == "function" %} +{%- set function = item.render -%} +{%- for decorator in function.decorators %} + @{{ decorator }} +{%- endfor %} +{% if function.is_async %}async {% endif %}def {{ function.name }}({{ function.args | join(', ') }}){% if function.returns %} -> {{ function.returns }}{% endif %}: +{%- if function.docstring %} + """{{ function.docstring }}""" +{% endif %} + {{ function.body | indent(8) }} +{%- endif -%} {% endfor %} diff --git a/mirascope_cli/schemas.py b/mirascope_cli/schemas.py index d899356..4997c3b 100644 --- a/mirascope_cli/schemas.py +++ b/mirascope_cli/schemas.py @@ -1,5 +1,5 @@ """Contains the schema for files created by the mirascope cli.""" -from typing import Optional +from typing import Literal, Optional, Union from pydantic import BaseModel, ConfigDict, Field @@ -47,3 +47,18 @@ class FunctionInfo(BaseModel): decorators: list[str] docstring: Optional[str] is_async: bool + + +class ASTOrder(BaseModel): + type: Literal["class", "function", "import", "from_import", "variable", "comment"] + order: int + render: Optional[ + Union[ + ClassInfo, + FunctionInfo, + tuple[str, Optional[str]], + tuple[str, str, Optional[str]], + tuple[str, str], + str, + ] + ] = None diff --git a/mirascope_cli/utils.py b/mirascope_cli/utils.py index 8066307..241a7ea 100644 --- a/mirascope_cli/utils.py +++ b/mirascope_cli/utils.py @@ -1,5 +1,4 @@ """Utility functions for the mirascope library.""" - from __future__ import annotations import ast @@ -14,9 +13,11 @@ from jinja2 import Environment, FileSystemLoader -from .enums import MirascopeCommand +from mirascope_cli.enums import MirascopeCommand + from .constants import CURRENT_REVISION_KEY, LATEST_REVISION_KEY from .schemas import ( + ASTOrder, ClassInfo, FunctionInfo, MirascopeCliVariables, @@ -66,24 +67,37 @@ def __init__(self) -> None: self.classes: list[ClassInfo] = [] self.functions: list[FunctionInfo] = [] self.comments: str = "" + self.order: list[ASTOrder] = [] def visit_Import(self, node) -> None: """Extracts imports from the given node.""" for alias in node.names: - self.imports.append((alias.name, alias.asname)) + import_str = (alias.name, alias.asname) + self.imports.append(import_str) + self.order.append(ASTOrder(type="import", order=len(self.imports) - 1)) + self.generic_visit(node) def visit_ImportFrom(self, node) -> None: """Extracts from imports from the given node.""" for alias in node.names: - self.from_imports.append((node.module, alias.name, alias.asname)) + from_import_str = (node.module, alias.name, alias.asname) + self.from_imports.append(from_import_str) + self.order.append( + ASTOrder(type="from_import", order=len(self.from_imports) - 1) + ) self.generic_visit(node) def visit_Assign(self, node) -> None: """Extracts variables from the given node.""" target = node.targets[0] if isinstance(target, ast.Name): + already_exists = target.id in self.variables self.variables[target.id] = ast.unparse(node.value) + if not already_exists: + self.order.append( + ASTOrder(type="variable", order=len(self.variables.keys()) - 1) + ) self.generic_visit(node) def visit_ClassDef(self, node) -> None: @@ -127,6 +141,7 @@ def visit_ClassDef(self, node) -> None: class_info.body = "\n".join(body) self.classes.append(class_info) + self.order.append(ASTOrder(type="class", order=len(self.classes) - 1)) def visit_AsyncFunctionDef(self, node): """Extracts async functions from the given node.""" @@ -160,11 +175,13 @@ def _visit_Function(self, node, is_async): # Assuming you have a list to store functions self.functions.append(function_info) + self.order.append(ASTOrder(type="function", order=len(self.functions) - 1)) def visit_Module(self, node) -> None: """Extracts comments from the given node.""" comments = ast.get_docstring(node, False) self.comments = "" if comments is None else comments + self.order.append(ASTOrder(type="comment", order=0)) self.generic_visit(node) def check_function_changed(self, other: PromptAnalyzer) -> bool: @@ -475,25 +492,29 @@ def _update_tag_decorator_with_version( return import_name -def _update_mirascope_imports(imports: list[tuple[str, Optional[str]]]): +def _update_mirascope_imports(analyzer: PromptAnalyzer): """Updates the mirascope import. Args: imports: The imports from the PromptAnalyzer class """ + imports = analyzer.imports if not any(import_name == "mirascope" for import_name, _ in imports): imports.append(("mirascope", None)) + index = 0 + if analyzer.comments: + index = 1 + analyzer.order.insert(index, ASTOrder(type="import", order=len(imports) - 1)) -def _update_mirascope_from_imports( - member: str, from_imports: list[tuple[str, str, Optional[str]]] -): +def _update_mirascope_from_imports(member: str, analyzer: PromptAnalyzer): """Updates the mirascope from imports. Args: member: The member to import. from_imports: The from imports from the PromptAnalyzer class """ + from_imports = analyzer.from_imports if not any( ( module_name == "mirascope" @@ -504,6 +525,12 @@ def _update_mirascope_from_imports( for module_name, import_name, _ in from_imports ): from_imports.append(("mirascope", member, None)) + index = 0 + if analyzer.comments: + index = 1 + analyzer.order.insert( + index, ASTOrder(type="from_import", order=len(from_imports) - 1) + ) def write_prompt_to_template( @@ -537,20 +564,6 @@ def write_prompt_to_template( if variables is None: variables = MirascopeCliVariables() - if command == MirascopeCommand.ADD: - # double quote revision ids to match how `ast.unparse()` formats strings - new_variables = { - k: f"'{v}'" if isinstance(v, str) else None - for k, v in variables.__dict__.items() - } | analyzer.variables - else: # command == MirascopeCommand.USE - ignore_variable_keys = dict.fromkeys(ignore_variables, None) - new_variables = { - k: analyzer.variables[k] - for k in analyzer.variables - if k not in ignore_variable_keys - } - if auto_tag: import_tag_name: Optional[str] = None mirascope_alias = "mirascope" @@ -569,18 +582,70 @@ def write_prompt_to_template( import_tag_name = _update_tag_decorator_with_version( decorators, variables, mirascope_alias ) - if import_tag_name == "tags": - _update_mirascope_from_imports(import_tag_name, analyzer.from_imports) + _update_mirascope_from_imports(import_tag_name, analyzer) elif import_tag_name == f"{mirascope_alias}.tags": - _update_mirascope_imports(analyzer.imports) + _update_mirascope_imports(analyzer) + + for item in analyzer.order: + if item.type == "import": + item.render = analyzer.imports[item.order] + elif item.type == "from_import": + item.render = analyzer.from_imports[item.order] + elif item.type == "variable": + variable_name = list(analyzer.variables.keys())[item.order] + variable_value = analyzer.variables[variable_name] + item.render = (variable_name, variable_value) + elif item.type == "class": + item.render = analyzer.classes[item.order] + elif item.type == "function": + item.render = analyzer.functions[item.order] + elif item.type == "comment": + item.render = analyzer.comments + if command == MirascopeCommand.ADD: + # double quote revision ids to match how `ast.unparse()` formats strings + new_variables = { + k: f"'{v}'" if isinstance(v, str) else None + for k, v in variables.__dict__.items() + } + first_class_func_var_index = next( + ( + i + for i, item in enumerate(analyzer.order) + if item.type in ["class", "function", "variable"] + ), + None, + ) + new_variables_order = [ + ASTOrder(type="variable", order=i, render=(k, v)) + for i, (k, v) in enumerate(new_variables.items()) + ] + if first_class_func_var_index is not None: + analyzer.order[ + first_class_func_var_index:first_class_func_var_index + ] = new_variables_order + else: + analyzer.order += new_variables_order + else: # command == MirascopeCommand.USE + ignore_variable_keys = dict.fromkeys(ignore_variables, None) + analyzer.order = [ + item + for item in analyzer.order + if not ( + item.type == "variable" + and item.render is not None + and item.render[0] in ignore_variable_keys + ) + ] data = { - "comments": analyzer.comments, - "variables": new_variables, - "imports": analyzer.imports, - "from_imports": analyzer.from_imports, - "classes": analyzer.classes, + # "comments": analyzer.comments, + # "variables": new_variables, + # "imports": analyzer.imports, + # "from_imports": analyzer.from_imports, + # "classes": analyzer.classes, + # "functions": analyzer.functions, + "order": analyzer.order, } return template.render(**data) diff --git a/tests/commands/golden/base_prompt/0001_base_prompt.py b/tests/commands/golden/base_prompt/0001_base_prompt.py index 1d00320..834de9f 100644 --- a/tests/commands/golden/base_prompt/0001_base_prompt.py +++ b/tests/commands/golden/base_prompt/0001_base_prompt.py @@ -1,5 +1,4 @@ """A prompt for recommending movies of a particular genre.""" - from mirascope import BasePrompt prev_revision_id = None diff --git a/tests/commands/golden/base_prompt/0002_base_prompt.py b/tests/commands/golden/base_prompt/0002_base_prompt.py index ea09227..24eaee2 100644 --- a/tests/commands/golden/base_prompt/0002_base_prompt.py +++ b/tests/commands/golden/base_prompt/0002_base_prompt.py @@ -1,5 +1,4 @@ """A prompt for recommending movies of a particular genre.""" - from mirascope import BasePrompt prev_revision_id = "0001" diff --git a/tests/commands/golden/base_prompt_auto_tag/0001_base_prompt_auto_tag.py b/tests/commands/golden/base_prompt_auto_tag/0001_base_prompt_auto_tag.py index aad4112..9c9ec12 100644 --- a/tests/commands/golden/base_prompt_auto_tag/0001_base_prompt_auto_tag.py +++ b/tests/commands/golden/base_prompt_auto_tag/0001_base_prompt_auto_tag.py @@ -1,5 +1,4 @@ """A prompt for recommending movies of a particular genre.""" - from mirascope import BasePrompt, tags prev_revision_id = None diff --git a/tests/commands/golden/base_prompt_auto_tag/0002_base_prompt_auto_tag.py b/tests/commands/golden/base_prompt_auto_tag/0002_base_prompt_auto_tag.py index ee910b8..e98401f 100644 --- a/tests/commands/golden/base_prompt_auto_tag/0002_base_prompt_auto_tag.py +++ b/tests/commands/golden/base_prompt_auto_tag/0002_base_prompt_auto_tag.py @@ -1,5 +1,4 @@ """A prompt for recommending movies of a particular genre.""" - from mirascope import BasePrompt, tags prev_revision_id = "0001" diff --git a/tests/commands/golden/base_prompt_with_decorator/0001_base_prompt_with_decorator.py b/tests/commands/golden/base_prompt_with_decorator/0001_base_prompt_with_decorator.py index 5165971..ea4d7b5 100644 --- a/tests/commands/golden/base_prompt_with_decorator/0001_base_prompt_with_decorator.py +++ b/tests/commands/golden/base_prompt_with_decorator/0001_base_prompt_with_decorator.py @@ -1,5 +1,4 @@ """A prompt for recommending movies of a particular genre.""" - from mirascope import BasePrompt, tags prev_revision_id = None diff --git a/tests/commands/golden/base_prompt_with_decorator/0002_base_prompt_with_decorator.py b/tests/commands/golden/base_prompt_with_decorator/0002_base_prompt_with_decorator.py index 3869764..a14af12 100644 --- a/tests/commands/golden/base_prompt_with_decorator/0002_base_prompt_with_decorator.py +++ b/tests/commands/golden/base_prompt_with_decorator/0002_base_prompt_with_decorator.py @@ -1,5 +1,4 @@ """A prompt for recommending movies of a particular genre.""" - from mirascope import BasePrompt, tags prev_revision_id = "0001" diff --git a/tests/commands/golden/call_with_variables/0001_call_with_variables.py b/tests/commands/golden/call_with_variables/0001_call_with_variables.py index 26fafe8..165d9e8 100644 --- a/tests/commands/golden/call_with_variables/0001_call_with_variables.py +++ b/tests/commands/golden/call_with_variables/0001_call_with_variables.py @@ -1,5 +1,4 @@ """A call for recommending movies of a particular genre.""" - from mirascope import tags from mirascope.openai import OpenAICall, OpenAICallParams @@ -7,7 +6,12 @@ revision_id = "0001" number = 1 chat = OpenAICall() -a_list = [1, 2, 3] + + +def foo(a: int, b: str) -> int: + """ABC""" + + return a + int(b) @tags(["movie_project", "version:0001"]) @@ -23,9 +27,14 @@ class MovieRecommender(OpenAICall): include succinct and clear descriptions of the movie. You also make sure to pique their interest by mentioning any famous actors in the movie that might be of interest. - + + USER: Please recommend 3 movies in the {genre} cetegory. """ + genre: str call_params = OpenAICallParams(model="gpt-3.5-turbo") + + +a_list = [1, 2, 3] diff --git a/tests/commands/golden/call_with_variables/0002_call_with_variables.py b/tests/commands/golden/call_with_variables/0002_call_with_variables.py index fd414c2..d679107 100644 --- a/tests/commands/golden/call_with_variables/0002_call_with_variables.py +++ b/tests/commands/golden/call_with_variables/0002_call_with_variables.py @@ -1,5 +1,4 @@ """A call for recommending movies of a particular genre.""" - from mirascope import tags from mirascope.openai import OpenAICall, OpenAICallParams @@ -7,7 +6,12 @@ revision_id = "0002" number = 1 chat = OpenAICall() -a_list = [1, 2, 3] + + +def foo(a: int, b: str) -> int: + """ABC""" + + return a + int(b) @tags(["movie_project", "version:0001"]) @@ -23,9 +27,14 @@ class MovieRecommender(OpenAICall): include succinct and clear descriptions of the movie. You also make sure to pique their interest by mentioning any famous actors in the movie that might be of interest. - + + USER: Please recommend 3 movies in the {genre} cetegory. """ + genre: str call_params = OpenAICallParams(model="gpt-3.5-turbo") + + +a_list = [1, 2, 3] diff --git a/tests/commands/golden/call_with_variables/call_with_variables.py b/tests/commands/golden/call_with_variables/call_with_variables.py index 9e16d69..c0307c8 100644 --- a/tests/commands/golden/call_with_variables/call_with_variables.py +++ b/tests/commands/golden/call_with_variables/call_with_variables.py @@ -4,7 +4,12 @@ number = 1 chat = OpenAICall() -a_list = [1, 2, 3] + + +def foo(a: int, b: str) -> int: + """ABC""" + + return a + int(b) @tags(["movie_project", "version:0001"]) @@ -28,3 +33,6 @@ class MovieRecommender(OpenAICall): genre: str call_params = OpenAICallParams(model="gpt-3.5-turbo") + + +a_list = [1, 2, 3] diff --git a/tests/commands/test_add.py b/tests/commands/test_add.py index ab42547..d72728f 100644 --- a/tests/commands/test_add.py +++ b/tests/commands/test_add.py @@ -36,10 +36,7 @@ def _initialize_tmp_mirascope(tmp_path: Path, golden_prompt: str): @pytest.mark.parametrize( "golden_prompt,auto_tag", - [ - ("base_prompt", False), - ("base_prompt", True), - ], + [("base_prompt", False), ("base_prompt", True), ("call_with_variables", False)], ) @pytest.mark.parametrize( "version_text_file", From 94aff300556d748d5a893683dfc19f371c7a9d4c Mon Sep 17 00:00:00 2001 From: Brendan Date: Tue, 16 Apr 2024 17:58:28 -0700 Subject: [PATCH 2/3] remove unused variables --- mirascope_cli/utils.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/mirascope_cli/utils.py b/mirascope_cli/utils.py index 241a7ea..7604c66 100644 --- a/mirascope_cli/utils.py +++ b/mirascope_cli/utils.py @@ -639,12 +639,6 @@ def write_prompt_to_template( ) ] data = { - # "comments": analyzer.comments, - # "variables": new_variables, - # "imports": analyzer.imports, - # "from_imports": analyzer.from_imports, - # "classes": analyzer.classes, - # "functions": analyzer.functions, "order": analyzer.order, } return template.render(**data) From c592c6190040168bba63c2c6e39478df7e4d8da3 Mon Sep 17 00:00:00 2001 From: Brendan Date: Tue, 16 Apr 2024 18:19:01 -0700 Subject: [PATCH 3/3] ignored type --- mirascope_cli/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mirascope_cli/utils.py b/mirascope_cli/utils.py index 7604c66..9a5ebb5 100644 --- a/mirascope_cli/utils.py +++ b/mirascope_cli/utils.py @@ -635,7 +635,7 @@ def write_prompt_to_template( if not ( item.type == "variable" and item.render is not None - and item.render[0] in ignore_variable_keys + and item.render[0] in ignore_variable_keys # type: ignore ) ] data = {