tinybird 0.0.1.dev42__py3-none-any.whl → 0.0.1.dev44__py3-none-any.whl

tinybird/tb/modules/create.py

@@ -18,6 +18,7 @@ from tinybird.tb.modules.feedback_manager import FeedbackManager
 from tinybird.tb.modules.llm import LLM
 from tinybird.tb.modules.llm_utils import extract_xml, parse_xml
 from tinybird.tb.modules.local_common import get_tinybird_local_client
+from tinybird.tb.modules.project import Project
 
 
 @cli.command()
@@ -41,8 +42,10 @@ from tinybird.tb.modules.local_common import get_tinybird_local_client
 )
 @click.option("--rows", type=int, default=10, help="Number of events to send")
 @click.option("--cursor", default=False, is_flag=True, help="Create .cursorrules file with Tinybird rules")
+@click.pass_context
 @coro
 async def create(
+    ctx: click.Context,
     data: Optional[str],
     prompt: Optional[str],
     folder: Optional[str],
@@ -50,13 +53,14 @@ async def create(
     cursor: bool,
 ) -> None:
     """Initialize a new project."""
+    project: Project = ctx.ensure_object(dict)["project"]
     folder = folder or getcwd()
     folder_path = Path(folder)
     if not folder_path.exists():
         folder_path.mkdir()
 
     try:
-        config = CLIConfig.get_project_config(folder)
+        config = CLIConfig.get_project_config(str(project.path))
         tb_client = config.get_client()
         user_token: Optional[str] = None
         if prompt:
@@ -80,7 +84,7 @@ async def create(
             click.echo(FeedbackManager.success(message="✓ Scaffolding completed!\n"))
 
         click.echo(FeedbackManager.highlight(message="\n» Creating resources..."))
-        datasources_created = await create_resources(local_client, tb_client, user_token, data, prompt, folder)
+        result = await create_resources(local_client, tb_client, user_token, data, prompt, folder)
         click.echo(FeedbackManager.success(message="✓ Done!\n"))
 
         if not already_has_cicd(folder):
@@ -89,7 +93,7 @@ async def create(
             await init_cicd(data_project_dir=os.path.relpath(folder))
             click.echo(FeedbackManager.success(message="✓ Done!\n"))
 
-        if validate_fixtures(folder) and datasources_created:
+        if should_generate_fixtures(result):
             click.echo(FeedbackManager.highlight(message="\n» Generating fixtures..."))
 
             if data:
@@ -110,7 +114,7 @@ async def create(
                     datasource_content = datasource_path.read_text()
                     has_json_path = "`json:" in datasource_content
                     if has_json_path:
-                        prompt = f"<datasource_schema>{datasource_content}</datasource_schema>\n<user_input>{prompt}</user_input>"
+                        prompt = f"<datasource_schema>{datasource_content}</datasource_schema>"
                         response = llm.ask(system_prompt=mock_prompt(rows), prompt=prompt)
                         sql = extract_xml(response, "sql")
                         sql = sql.split("FORMAT")[0]
@@ -139,9 +143,10 @@ def validate_project_structure(folder: str) -> bool:
     return all((Path(folder) / path).exists() for path in PROJECT_PATHS)
 
 
-def validate_fixtures(folder: str) -> bool:
-    datasource_files = [f for f in os.listdir(Path(folder) / "datasources") if f.endswith(".datasource")]
-    return len(datasource_files) > 0
+def should_generate_fixtures(result: str) -> bool:
+    if not result:
+        return False
+    return "<type>datasource</type>" in result
 
 
 def already_has_cicd(folder: str) -> bool:
@@ -168,12 +173,12 @@ async def create_resources(
     prompt: Optional[str],
     folder: str,
 ):
-    force = True
+    result = ""
     folder_path = Path(folder)
     if data:
         path = folder_path / data
         format = path.suffix.lstrip(".")
-        await _generate_datafile(str(path), local_client, format=format, force=force)
+        await _generate_datafile(str(path), local_client, format=format, force=True)
         name = data.split(".")[0]
         generate_pipe_file(
             f"{name}_endpoint",
@@ -185,7 +190,10 @@ TYPE ENDPOINT
             """,
             folder,
         )
-        return True
+        result = (
+            f"<response><resource><type>datasource</type><name>{name}</name><content></content></resource></response>"
+        )
+
     elif prompt and user_token:
         datasource_paths = [
             Path(folder) / "datasources" / f
@@ -212,41 +220,74 @@ TYPE ENDPOINT
             ]
         )
         llm = LLM(user_token=user_token, host=tb_client.host)
-        result = llm.ask(system_prompt=create_prompt(resources_xml), prompt=prompt)
-        result = extract_xml(result, "response")
-        resources = parse_xml(result, "resource")
-        datasources = []
-        pipes = []
-        for resource_xml in resources:
-            resource_type = extract_xml(resource_xml, "type")
-            name = extract_xml(resource_xml, "name")
-            content = extract_xml(resource_xml, "content")
-            resource = {
-                "name": name,
-                "content": content,
-            }
-            if resource_type.lower() == "datasource":
-                datasources.append(resource)
-            elif resource_type.lower() == "pipe":
-                pipes.append(resource)
-
-        for ds in datasources:
-            content = ds["content"].replace("```", "")
-            filename = f"{ds['name']}.datasource"
-            generate_datafile(
-                content,
-                filename=filename,
-                data=None,
-                _format="ndjson",
-                force=force,
-                folder=folder,
-            )
-
-        for pipe in pipes:
-            content = pipe["content"].replace("```", "")
-            generate_pipe_file(pipe["name"], content, folder)
-
-        return len(datasources) > 0
+        result = ""
+        iterations = 0
+        history = ""
+        generated_paths: list[Path] = []
+
+        while iterations < 10:
+            feedback = ""
+            if iterations > 0:
+                feedback = click.prompt("\nFollow-up instructions or continue", default="continue")
+            if iterations > 0 and (not feedback or feedback in ("continue", "ok", "exit", "quit", "q")):
+                break
+            else:
+                if iterations > 0:
+                    click.echo(FeedbackManager.highlight(message="\n» Creating resources..."))
+                    for path in generated_paths:
+                        path.unlink()
+                    generated_paths = []
+
+                save_context(prompt, feedback)
+                result = llm.ask(system_prompt=create_prompt(resources_xml, feedback, history), prompt=prompt)
+                result = extract_xml(result, "response")
+                history = (
+                    history
+                    + f"""
+                <result_iteration_{iterations}>
+                    {result}
+                </result_iteration_{iterations}>
+                """
+                )
+                resources = parse_xml(result, "resource")
+                datasources = []
+                pipes = []
+                for resource_xml in resources:
+                    resource_type = extract_xml(resource_xml, "type")
+                    name = extract_xml(resource_xml, "name")
+                    content = extract_xml(resource_xml, "content")
+                    resource = {
+                        "name": name,
+                        "content": content,
+                    }
+                    if resource_type.lower() == "datasource":
+                        datasources.append(resource)
+                    elif resource_type.lower() == "pipe":
+                        pipes.append(resource)
+
+                for ds in datasources:
+                    content = ds["content"].replace("```", "")
+                    filename = f"{ds['name']}.datasource"
+                    datasource_path = generate_datafile(
+                        content,
+                        filename=filename,
+                        data=None,
+                        _format="ndjson",
+                        force=True,
+                        folder=folder,
+                    )
+                    generated_paths.append(datasource_path)
+                for pipe in pipes:
+                    content = pipe["content"].replace("```", "")
+                    pipe_path = generate_pipe_file(pipe["name"], content, folder)
+                    generated_paths.append(pipe_path)
+
+                iterations += 1
+
+        if iterations == 10:
+            click.echo(FeedbackManager.info(message="Too many iterations. Change the prompt and try again."))
+
+        return result
 
 
 def init_git(folder: str):
@@ -266,7 +307,7 @@ def init_git(folder: str):
         raise CLIException(f"Error initializing Git: {e}")
 
 
-def generate_pipe_file(name: str, content: str, folder: str):
+def generate_pipe_file(name: str, content: str, folder: str) -> Path:
     def is_copy(content: str) -> bool:
         return re.search(r"TYPE copy", content, re.IGNORECASE) is not None
 
@@ -292,6 +333,7 @@ def generate_pipe_file(name: str, content: str, folder: str):
     with open(f"{f}", "w") as file:
         file.write(content)
     click.echo(FeedbackManager.info_file_created(file=f.relative_to(folder)))
+    return f.relative_to(folder)
 
 
 def already_has_cursorrules(folder: str) -> bool:
@@ -301,3 +343,20 @@ def already_has_cursorrules(folder: str) -> bool:
 def create_cursorrules(folder: str):
     cursorrules_file = Path(folder) / ".cursorrules"
     cursorrules_file.write_text(cursorrules_prompt)
+
+
+def get_context_file() -> Path:
+    context_file = Path(os.path.expanduser("~/.tb_create_context"))
+    if not context_file.exists():
+        context_file.touch()
+    return context_file
+
+
+def get_context() -> str:
+    context_file = get_context_file()
+    return context_file.read_text()
+
+
+def save_context(prompt: str, feedback: str):
+    context_file = get_context_file()
+    context_file.write_text(f"- {prompt}\n{feedback}")