tinybird 0.0.1.dev25__py3-none-any.whl → 0.0.1.dev27__py3-none-any.whl

tinybird/tb/modules/create.py

@@ -4,9 +4,9 @@ from pathlib import Path
 from typing import Optional
 
 import click
-import requests
 
 from tinybird.client import TinyB
+from tinybird.prompts import create_prompt, mock_prompt
 from tinybird.tb.modules.cicd import init_cicd
 from tinybird.tb.modules.cli import cli
 from tinybird.tb.modules.common import _generate_datafile, check_user_token_with_client, coro, generate_datafile
@@ -15,15 +15,11 @@ from tinybird.tb.modules.datafile.fixture import build_fixture_name, persist_fix
 from tinybird.tb.modules.exceptions import CLIException
 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
 
 
 @cli.command()
-@click.option(
-    "--demo",
-    is_flag=True,
-    help="Demo data and files to get started",
-)
 @click.option(
     "--data",
     type=click.Path(exists=True),
@@ -39,13 +35,12 @@ from tinybird.tb.modules.local_common import get_tinybird_local_client
 @click.option(
     "--folder",
     default=".",
-    type=click.Path(exists=True, file_okay=False),
+    type=click.Path(exists=False, file_okay=False),
     help="Folder where datafiles will be placed",
 )
 @click.option("--rows", type=int, default=10, help="Number of events to send")
 @coro
 async def create(
-    demo: bool,
     data: Optional[str],
     prompt: Optional[str],
     folder: Optional[str],
@@ -53,6 +48,10 @@ async def create(
 ) -> None:
     """Initialize a new project."""
     folder = folder or getcwd()
+    folder_path = Path(folder)
+    if not folder_path.exists():
+        folder_path.mkdir()
+
     try:
         config = CLIConfig.get_project_config(folder)
         tb_client = config.get_client()
@@ -71,127 +70,92 @@ async def create(
                 )
                 return
         local_client = await get_tinybird_local_client(folder)
-        click.echo(FeedbackManager.gray(message="Creating new project structure..."))
-        await project_create(local_client, tb_client, user_token, data, prompt, folder)
-        click.echo(FeedbackManager.success(message="✓ Scaffolding completed!\n"))
 
-        click.echo(FeedbackManager.gray(message="\nCreating CI/CD files for GitHub and GitLab..."))
+        if not validate_project_structure(folder):
+            click.echo(FeedbackManager.highlight(message="\n» Creating new project structure..."))
+            create_project_structure(folder)
+            click.echo(FeedbackManager.success(message="✓ Scaffolding completed!\n"))
+
+        click.echo(FeedbackManager.highlight(message="\n» Creating resources..."))
+        await create_resources(local_client, tb_client, user_token, data, prompt, folder)
+        click.echo(FeedbackManager.success(message="✓ Done!\n"))
+
+        click.echo(FeedbackManager.highlight(message="\n» Creating CI/CD files for GitHub and GitLab..."))
         init_git(folder)
         await init_cicd(data_project_dir=os.path.relpath(folder))
         click.echo(FeedbackManager.success(message="✓ Done!\n"))
 
-        click.echo(FeedbackManager.gray(message="Building fixtures..."))
+        if validate_fixtures(folder):
+            click.echo(FeedbackManager.highlight(message="\n» Generating fixtures..."))
 
-        if demo:
-            # Users datasource
-            ds_name = "users"
-            datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
-            datasource_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/b48fb9c92825ed27c04e3104b9e871e1/raw/1f33c20eefbabc4903f38e234329e028d8ef9def/users.datasource"
-            )
-            datasource_path.write_text(datasource_content)
-            click.echo(FeedbackManager.info(message=f"✓ /datasources/{ds_name}.datasource"))
+            if data:
+                ds_name = os.path.basename(data.split(".")[0])
+                data_content = Path(data).read_text()
+                datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
+                fixture_name = build_fixture_name(
+                    datasource_path.absolute().as_posix(), ds_name, datasource_path.read_text()
+                )
+                click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
+                persist_fixture(fixture_name, data_content, folder)
+            elif prompt and user_token:
+                datasource_files = [f for f in os.listdir(Path(folder) / "datasources") if f.endswith(".datasource")]
+                for datasource_file in datasource_files:
+                    datasource_path = Path(folder) / "datasources" / datasource_file
+                    llm = LLM(user_token=user_token, client=tb_client)
+                    datasource_name = datasource_path.stem
+                    datasource_content = datasource_path.read_text()
+                    has_json_path = "`json:" in datasource_content
+                    if has_json_path:
+                        response = await llm.ask(prompt, system_prompt=mock_prompt(rows))
+                        sql = extract_xml(response, "sql")
+                        result = await local_client.query(f"{sql} FORMAT JSON")
+                        data = result.get("data", [])
+                        fixture_name = build_fixture_name(
+                            datasource_path.absolute().as_posix(), datasource_name, datasource_content
+                        )
+                        if data:
+                            persist_fixture(fixture_name, data, folder)
+                            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{datasource_name}"))
 
-            # Users fixtures
-            fixture_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/8e8f66a39d7576ce3a2529bf773334a8/raw/9cab636767990e97d44a141867e5f226e992de8c/users.ndjson"
-            )
-            fixture_name = build_fixture_name(
-                datasource_path.absolute().as_posix(), ds_name, datasource_path.read_text()
-            )
-            persist_fixture(fixture_name, fixture_content)
-            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
+            click.echo(FeedbackManager.success(message="✓ Done!\n"))
+    except Exception as e:
+        click.echo(FeedbackManager.error(message=f"Error: {str(e)}"))
 
-            # Events datasource
-            ds_name = "events"
-            datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
-            datasource_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/f8ca37b5b1f6707c75206b618de26bc9/raw/cd625da0dcd1ba8de29f12bc1c8600b9ff7c809c/events.datasource"
-            )
-            datasource_path.write_text(datasource_content)
-            click.echo(FeedbackManager.info(message=f"✓ /datasources/{ds_name}.datasource"))
 
-            # Events fixtures
-            fixture_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/859ab9439c17e77241d0c14a5a532809/raw/251f2f3f00a968f8759ec4068cebde915256b054/events.ndjson"
-            )
-            fixture_name = build_fixture_name(
-                datasource_path.absolute().as_posix(), ds_name, datasource_path.read_text()
-            )
-            persist_fixture(fixture_name, fixture_content)
-            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
+PROJECT_PATHS = ("datasources", "endpoints", "materializations", "copies", "sinks", "fixtures", "tests")
 
-            # Create sample endpoint
-            pipe_name = "api_token_usage"
-            pipe_path = Path(folder) / "endpoints" / f"{pipe_name}.pipe"
-            pipe_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/68ecc47472c2b754b0ae0c1187022963/raw/52cc3aa3afdf939e58d43355bfe4ddc739989ddd/api_token_usage.pipe"
-            )
-            pipe_path.write_text(pipe_content)
-            click.echo(FeedbackManager.info(message=f"✓ /endpoints/{pipe_name}.pipe"))
 
-            # Create sample test
-            test_name = "api_token_usage"
-            test_path = Path(folder) / "tests" / f"{test_name}.yaml"
-            test_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/e58620bbb977d6f42f1d0c2a7b46ac8f/raw/a3a1cd0ce3a90bcd2f6dfce00da51e6051443612/api_token_usage.yaml"
-            )
-            test_path.write_text(test_content)
-            click.echo(FeedbackManager.info(message=f"✓ /tests/{test_name}.yaml"))
+def validate_project_structure(folder: str) -> bool:
+    return all((Path(folder) / path).exists() for path in PROJECT_PATHS)
 
-        elif data:
-            ds_name = os.path.basename(data.split(".")[0])
-            data_content = Path(data).read_text()
-            datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
-            fixture_name = build_fixture_name(
-                datasource_path.absolute().as_posix(), ds_name, datasource_path.read_text()
-            )
-            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
-            persist_fixture(fixture_name, data_content)
-        elif prompt and user_token:
-            datasource_files = [f for f in os.listdir(Path(folder) / "datasources") if f.endswith(".datasource")]
-            for datasource_file in datasource_files:
-                datasource_path = Path(folder) / "datasources" / datasource_file
-                llm = LLM(user_token=user_token, client=tb_client)
-                datasource_name = datasource_path.stem
-                datasource_content = datasource_path.read_text()
-                has_json_path = "`json:" in datasource_content
-                if has_json_path:
-                    sql = await llm.generate_sql_sample_data(schema=datasource_content, rows=rows, prompt=prompt)
-                    result = await local_client.query(f"{sql} FORMAT JSON")
-                    data = result.get("data", [])
-                    fixture_name = build_fixture_name(
-                        datasource_path.absolute().as_posix(), datasource_name, datasource_content
-                    )
-                    if data:
-                        persist_fixture(fixture_name, data)
-                        click.echo(FeedbackManager.info(message=f"✓ /fixtures/{datasource_name}"))
 
-        click.echo(FeedbackManager.success(message="✓ Done!\n"))
-    except Exception as e:
-        click.echo(FeedbackManager.error(message=f"Error: {str(e)}"))
+def validate_fixtures(folder: str) -> bool:
+    return (Path(folder) / "datasources").exists()
+
+
+def create_project_structure(folder: str):
+    folder_path = Path(folder)
+    for x in PROJECT_PATHS:
+        try:
+            f = folder_path / x
+            f.mkdir()
+        except FileExistsError:
+            pass
+        click.echo(FeedbackManager.info_path_created(path=x))
 
 
-async def project_create(
+async def create_resources(
     local_client: TinyB,
-    client: TinyB,
+    tb_client: TinyB,
     user_token: Optional[str],
     data: Optional[str],
     prompt: Optional[str],
     folder: str,
 ):
-    project_paths = ["datasources", "endpoints", "materializations", "copies", "sinks", "fixtures", "tests"]
     force = True
-    for x in project_paths:
-        try:
-            f = Path(folder) / x
-            f.mkdir()
-        except FileExistsError:
-            pass
-        click.echo(FeedbackManager.info_path_created(path=x))
-
+    folder_path = Path(folder)
     if data:
-        path = Path(folder) / data
+        path = folder_path / data
         format = path.suffix.lstrip(".")
         try:
             await _generate_datafile(str(path), local_client, format=format, force=force)
@@ -210,17 +174,64 @@ TYPE ENDPOINT
         )
     elif prompt and user_token:
         try:
-            llm = LLM(user_token=user_token, client=client)
-            result = await llm.create_project(prompt)
-            for ds in result.datasources:
-                content = ds.content.replace("```", "")
+            datasource_paths = [
+                Path(folder) / "datasources" / f
+                for f in os.listdir(Path(folder) / "datasources")
+                if f.endswith(".datasource")
+            ]
+            pipes_paths = [
+                Path(folder) / "endpoints" / f for f in os.listdir(Path(folder) / "endpoints") if f.endswith(".pipe")
+            ]
+            resources_xml = "\n".join(
+                [
+                    f"<resource><type>{resource_type}</type><name>{resource_name}</name><content>{resource_content}</content></resource>"
+                    for resource_type, resource_name, resource_content in [
+                        ("datasource", ds.stem, ds.read_text()) for ds in datasource_paths
+                    ]
+                    + [
+                        (
+                            "pipe",
+                            pipe.stem,
+                            pipe.read_text(),
+                        )
+                        for pipe in pipes_paths
+                    ]
+                ]
+            )
+            llm = LLM(user_token=user_token, client=tb_client)
+            result = await llm.ask(prompt, system_prompt=create_prompt(resources_xml))
+            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=f"{ds.name}.datasource", data=None, _format="ndjson", force=force, folder=folder
+                    content,
+                    filename=filename,
+                    data=None,
+                    _format="ndjson",
+                    force=force,
+                    folder=folder,
                 )
 
-            for pipe in result.pipes:
-                content = pipe.content.replace("```", "")
-                generate_pipe_file(pipe.name, content, folder)
+            for pipe in pipes:
+                content = pipe["content"].replace("```", "")
+                generate_pipe_file(pipe["name"], content, folder)
         except Exception as e:
             click.echo(FeedbackManager.error(message=f"Error: {str(e)}"))
 
@@ -250,9 +261,3 @@ 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)))
-
-
-def fetch_gist_content(url: str) -> str:  # TODO: replace this with a function that fetches the content from a repo
-    response = requests.get(url)
-    response.raise_for_status()
-    return response.text

tinybird/tb/modules/datafile/build.py

@@ -722,6 +722,34 @@ async def process(
     except Exception as e:
         raise click.ClickException(str(e))
 
+    # datasource
+    # {
+    #     "resource": "datasources",
+    #     "resource_name": name,
+    #     "version": doc.version,
+    #     "params": params,
+    #     "filename": filename,
+    #     "deps": deps,
+    #     "tokens": doc.tokens,
+    #     "shared_with": doc.shared_with,
+    #     "filtering_tags": doc.filtering_tags,
+    # }
+    # pipe
+    # {
+    #     "resource": "pipes",
+    #     "resource_name": name,
+    #     "version": doc.version,
+    #     "filename": filename,
+    #     "name": name + version,
+    #     "nodes": nodes,
+    #     "deps": [x for x in set(deps)],
+    #     "tokens": doc.tokens,
+    #     "description": description,
+    #     "warnings": doc.warnings,
+    #     "filtering_tags": doc.filtering_tags,
+    # }
+
+    # r is essentially a Datasource or a Pipe in dict shape, like in the comment above
     for r in res:
         resource_name = r["resource_name"]
         warnings = r.get("warnings", [])

tinybird/tb/modules/datafile/common.py

@@ -1340,6 +1340,7 @@ def parse(
         "export_compression": assign_var("export_compression"),
         "export_write_strategy": assign_var("export_write_strategy"),
         "export_kafka_topic": assign_var("export_kafka_topic"),
+        "forward_query": sql("forward_query"),
     }
 
     engine_vars = set()

tinybird/tb/modules/datafile/fixture.py

@@ -31,22 +31,26 @@ def build_fixture_name(filename: str, datasource_name: str, datasource_content:
     return f"{datasource_name}_{hash_str}"
 
 
-def get_fixture_dir() -> Path:
-    fixture_dir = Path("fixtures")
+def get_fixture_dir(folder: str) -> Path:
+    fixture_dir = Path(folder) / "fixtures"
     if not fixture_dir.exists():
         fixture_dir.mkdir()
     return fixture_dir
 
 
-def persist_fixture(fixture_name: str, data: Union[List[Dict[str, Any]], str], format="ndjson") -> Path:
-    fixture_dir = get_fixture_dir()
+def persist_fixture(fixture_name: str, data: Union[List[Dict[str, Any]], str], folder: str, format="ndjson") -> Path:
+    fixture_dir = get_fixture_dir(folder)
     fixture_file = fixture_dir / f"{fixture_name}.{format}"
     fixture_file.write_text(data if isinstance(data, str) else format_data_to_ndjson(data))
     return fixture_file
 
 
-def load_fixture(fixture_name: str, format="ndjson") -> Union[Path, None]:
-    fixture_dir = get_fixture_dir()
+def load_fixture(
+    fixture_name: str,
+    folder: str,
+    format="ndjson",
+) -> Union[Path, None]:
+    fixture_dir = get_fixture_dir(folder)
     fixture_file = fixture_dir / f"{fixture_name}.{format}"
     if not fixture_file.exists():
         return None

tinybird/tb/modules/datafile/parse_pipe.py

@@ -45,6 +45,8 @@ def parse_pipe(
         for node in doc.nodes:
             sql = node.get("sql", "")
             if sql.strip()[0] == "%":
+                # Note(eclbg): not sure what test_mode is for. I think it does something like using placeholder values
+                # for the variables in the template.
                 sql, _, variable_warnings = render_sql_template(sql[1:], test_mode=True, name=node["name"])
                 doc.warnings = variable_warnings
             # it'll fail with a ModuleNotFoundError when the toolset is not available but it returns the parsed doc

tinybird/tb/modules/datasource.py

@@ -453,7 +453,7 @@ async def generate_datasource(ctx: Context, connector: str, filenames, force: bo
     """Generate a data source file based on a sample CSV file from local disk or url"""
     client: TinyB = ctx.ensure_object(dict)["client"]
 
-    _connector: Optional["Connector"] = None
+    _connector: Optional[Connector] = None
     if connector:
         load_connector_config(ctx, connector, False, check_uninstalled=False)
         if connector not in ctx.ensure_object(dict):

tinybird/tb/modules/deploy.py

@@ -0,0 +1,160 @@
+import glob
+import json
+import logging
+import time
+from pathlib import Path
+from typing import List
+
+import click
+import requests
+
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.config import CLIConfig
+from tinybird.tb.modules.feedback_manager import FeedbackManager
+
+
+def project_files(project_path: Path) -> List[str]:
+    project_file_extensions = ("datasource", "pipe")
+    project_files = []
+    for extension in project_file_extensions:
+        for project_file in glob.glob(f"{project_path}/**/*.{extension}", recursive=True):
+            logging.debug(f"Found project file: {project_file}")
+            project_files.append(project_file)
+    return project_files
+
+
+def promote_deployment(host: str, headers: dict) -> None:
+    TINYBIRD_API_URL = host + "/v1/deployments"
+    r = requests.get(TINYBIRD_API_URL, headers=headers)
+    result = r.json()
+    logging.debug(json.dumps(result, indent=2))
+
+    deployments = result.get("deployments")
+    if not deployments:
+        click.echo(FeedbackManager.error(message="No deployments found"))
+        return
+
+    last_deployment, candidate_deployment = deployments[0], deployments[1]
+
+    if candidate_deployment.get("status") != "data_ready":
+        click.echo(FeedbackManager.error(message="Current deployment is not ready"))
+        return
+
+    if candidate_deployment.get("live"):
+        click.echo(FeedbackManager.error(message="Candidate deployment is already live"))
+    else:
+        click.echo(FeedbackManager.success(message="Promoting deployment"))
+
+        TINYBIRD_API_URL = host + f"/v1/deployments/{candidate_deployment.get('id')}/set-live"
+        r = requests.post(TINYBIRD_API_URL, headers=headers)
+        result = r.json()
+        logging.debug(json.dumps(result, indent=2))
+
+    click.echo(FeedbackManager.success(message="Removing old deployment"))
+
+    TINYBIRD_API_URL = host + f"/v1/deployments/{last_deployment.get('id')}"
+    r = requests.delete(TINYBIRD_API_URL, headers=headers)
+    result = r.json()
+    logging.debug(json.dumps(result, indent=2))
+
+    click.echo(FeedbackManager.success(message="Deployment promoted successfully"))
+
+
+@cli.command()
+@click.argument("project_path", type=click.Path(exists=True), default=Path.cwd())
+@click.option(
+    "--wait/--no-wait",
+    is_flag=True,
+    default=False,
+    help="Wait for deploy to finish. Disabled by default.",
+)
+@click.option(
+    "--auto/--no-auto",
+    is_flag=True,
+    default=False,
+    help="Auto-promote the deployment. Only works if --wait is enabled. Disabled by default.",
+)
+def deploy(project_path: Path, wait: bool, auto: bool) -> None:
+    """
+    Validate and deploy the project server side.
+    """
+    # TODO: This code is duplicated in build_server.py
+    # Should be refactored to be shared
+    MULTIPART_BOUNDARY_DATA_PROJECT = "data_project://"
+    DATAFILE_TYPE_TO_CONTENT_TYPE = {
+        ".datasource": "text/plain",
+        ".pipe": "text/plain",
+    }
+
+    config = CLIConfig.get_project_config(str(project_path))
+    TINYBIRD_API_URL = (config.get_host() or "") + "/v1/deploy"
+    TINYBIRD_API_KEY = config.get_token()
+
+    files = [
+        ("context://", ("cli-version", "1.0.0", "text/plain")),
+    ]
+    fds = []
+    for file_path in project_files(project_path):
+        relative_path = str(Path(file_path).relative_to(project_path))
+        fd = open(file_path, "rb")
+        fds.append(fd)
+        content_type = DATAFILE_TYPE_TO_CONTENT_TYPE.get(Path(file_path).suffix, "application/unknown")
+        files.append((MULTIPART_BOUNDARY_DATA_PROJECT, (relative_path, fd.read().decode("utf-8"), content_type)))
+
+    deployment = None
+    try:
+        HEADERS = {"Authorization": f"Bearer {TINYBIRD_API_KEY}"}
+
+        r = requests.post(TINYBIRD_API_URL, files=files, headers=HEADERS)
+        result = r.json()
+        logging.debug(json.dumps(result, indent=2))
+
+        deploy_result = result.get("result")
+        if deploy_result == "success":
+            click.echo(FeedbackManager.success(message="Deploy submitted successfully"))
+            deployment = result.get("deployment")
+        elif deploy_result == "failed":
+            click.echo(FeedbackManager.error(message="Deploy failed"))
+            deploy_errors = result.get("errors")
+            for deploy_error in deploy_errors:
+                if deploy_error.get("filename", None):
+                    click.echo(
+                        FeedbackManager.error(message=f"{deploy_error.get('filename')}\n\n{deploy_error.get('error')}")
+                    )
+                else:
+                    click.echo(FeedbackManager.error(message=f"{deploy_error.get('error')}"))
+        else:
+            click.echo(FeedbackManager.error(message=f"Unknown build result {deploy_result}"))
+    finally:
+        for fd in fds:
+            fd.close()
+
+    if deployment and wait:
+        while deployment.get("status") != "data_ready":
+            time.sleep(5)
+            TINYBIRD_API_URL = (config.get_host() or "") + f"/v1/deployments/{deployment.get('id')}"
+            r = requests.get(TINYBIRD_API_URL, headers=HEADERS)
+            result = r.json()
+            deployment = result.get("deployment")
+            if deployment.get("status") == "failed":
+                click.echo(FeedbackManager.error(message="Deployment failed"))
+                return
+
+        click.echo(FeedbackManager.success(message="Deployment is ready"))
+
+        if auto:
+            promote_deployment((config.get_host() or ""), HEADERS)
+
+
+@cli.command(name="release")
+@click.argument("project_path", type=click.Path(exists=True), default=Path.cwd())
+def deploy_promote(project_path: Path) -> None:
+    """
+    Promote last deploy to ready and remove old one.
+    """
+    config = CLIConfig.get_project_config(str(project_path))
+
+    TINYBIRD_API_KEY = config.get_token()
+    HEADERS = {"Authorization": f"Bearer {TINYBIRD_API_KEY}"}
+
+    promote_deployment((config.get_host() or ""), HEADERS)

tinybird/tb/modules/llm.py

@@ -1,10 +1,8 @@
-import asyncio
 import json
 import urllib.parse
 from copy import deepcopy
-from typing import Awaitable, Callable, List, Optional
+from typing import List
 
-from openai import OpenAI
 from pydantic import BaseModel
 
 from tinybird.client import TinyB
@@ -31,25 +29,43 @@ class TestExpectations(BaseModel):
 
 
 class LLM:
-    def __init__(self, user_token: str, client: TinyB, api_key: Optional[str] = None):
+    def __init__(
+        self,
+        user_token: str,
+        client: TinyB,
+    ):
         self.user_client = deepcopy(client)
         self.user_client.token = user_token
 
-        self.openai = OpenAI(api_key=api_key) if api_key else None
+    async def ask(self, prompt: str, system_prompt: str = "", model: str = "o1-mini") -> str:
+        """
+        Calls the model with the given prompt and returns the response.
 
-    async def _execute(self, action_fn: Callable[[], Awaitable[str]], checker_fn: Callable[[str], bool]):
-        is_valid = False
-        times = 0
+        Args:
+            prompt (str): The user prompt to send to the model.
 
-        while not is_valid and times < 5:
-            result = await action_fn()
-            if asyncio.iscoroutinefunction(checker_fn):
-                is_valid = await checker_fn(result)
-            else:
-                is_valid = checker_fn(result)
-            times += 1
+        Returns:
+            str: The response from the language model.
+        """
+        messages = []
 
-        return result
+        if system_prompt:
+            messages.append({"role": "user", "content": system_prompt})
+
+        if prompt:
+            messages.append({"role": "user", "content": prompt})
+
+        data = {
+            "model": model,
+            "messages": messages,
+        }
+        response = await self.user_client._req(
+            "/v0/llm",
+            method="POST",
+            data=json.dumps(data),
+            headers={"Content-Type": "application/json"},
+        )
+        return response.get("result", "")
 
     async def create_project(self, prompt: str) -> DataProject:
         try: