├── .env.sample ├── .gitignore ├── LICENSE ├── README.md ├── api-server ├── .env.sample ├── .gitignore ├── README.md ├── api │ ├── index.py │ └── modules │ │ ├── db.py │ │ ├── emb.py │ │ ├── file.py │ │ ├── instruments.py │ │ ├── llm.py │ │ ├── models.py │ │ └── turbo4.py ├── requirements.txt └── vercel.json ├── fe-clients ├── react-ts │ ├── .eslintrc.cjs │ ├── .gitignore │ ├── README.md │ ├── bun.lockb │ ├── index.html │ ├── package.json │ ├── public │ │ └── vite.svg │ ├── src │ │ ├── App.css │ │ ├── App.tsx │ │ ├── assets │ │ │ └── react.svg │ │ ├── index.css │ │ ├── main.tsx │ │ └── vite-env.d.ts │ ├── tsconfig.json │ ├── tsconfig.node.json │ └── vite.config.ts ├── svelte-ts │ ├── .gitignore │ ├── .vscode │ │ └── extensions.json │ ├── README.md │ ├── bun.lockb │ ├── index.html │ ├── package.json │ ├── public │ │ └── vite.svg │ ├── src │ │ ├── App.svelte │ │ ├── app.css │ │ ├── assets │ │ │ └── svelte.svg │ │ ├── lib │ │ │ └── Counter.svelte │ │ ├── main.ts │ │ └── vite-env.d.ts │ ├── svelte.config.js │ ├── tsconfig.json │ ├── tsconfig.node.json │ └── vite.config.ts └── vue-ts │ ├── .gitignore │ ├── .vscode │ └── extensions.json │ ├── README.md │ ├── bun.lockb │ ├── index.html │ ├── package.json │ ├── public │ └── vite.svg │ ├── src │ ├── App.vue │ ├── assets │ │ ├── ttydb.svg │ │ └── vue.svg │ ├── components │ │ └── HelloWorld.vue │ ├── main.ts │ ├── style.css │ └── vite-env.d.ts │ ├── tsconfig.json │ ├── tsconfig.node.json │ └── vite.config.ts ├── imgs ├── 1-prompt-engineering-postgres-ai-data-analytics-agent.png ├── 10-talk-to-your-database.png ├── 2-using-autogen-to-build-our-multi-agent-postgres-data-analytics-tool.png ├── 2024-predictions.png ├── 3-make-autogen-consistent-to-build-our-multi-agent-postgres-data-analytics-tool.png ├── 4-autogen-token-tactics-managing-llm-memory-and-costs-multi-agent-postgres-ai-data-analytics.png ├── 5-autogen-spyware-for-ai-agents-postgres-data-analytics-tool-ai.png ├── 6-autogen-and-guidance-for-autonomous-control-flow.png ├── 7-turbo4-assistants-threads-messages.png ├── 8-ccc-ai-engineering-with-aider.png ├── 9-self-correcting-openai-gpt4-turbo-assistant.png └── multi-agent-coding.png ├── poetry.lock ├── postgres_da_ai_agent ├── agents │ ├── agent_config.py │ ├── agents.py │ ├── instruments.py │ └── turbo4.py ├── main.py ├── modules │ ├── db.py │ ├── embeddings.py │ ├── file.py │ ├── llm.py │ ├── orchestrator.py │ └── rand.py ├── turbo_main.py └── types.py └── pyproject.toml /.env.sample: -------------------------------------------------------------------------------- 1 | DATABASE_URL= 2 | OPENAI_API_KEY= 3 | BASE_DIR=./agent_results -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .aider.ident.cache.v1 2 | .aider.tags.cache.v1 3 | .aider.chat.history.md 4 | .aider.input.history 5 | 6 | node_modules/ 7 | agent_results/ 8 | 9 | .env 10 | .venv 11 | 12 | *.pyc 13 | 14 | # Packages 15 | *.egg 16 | !/tests/**/*.egg 17 | /*.egg-info 18 | /dist/* 19 | build 20 | _build 21 | .cache 22 | *.so 23 | 24 | # Installer logs 25 | pip-log.txt 26 | 27 | # Unit test / coverage reports 28 | .coverage 29 | .pytest_cache 30 | 31 | .DS_Store 32 | .idea/* 33 | .python-version 34 | .vscode/* 35 | 36 | /test.py 37 | /test_*.* 38 | 39 | /setup.cfg 40 | MANIFEST.in 41 | /setup.py 42 | /docs/site/* 43 | /tests/fixtures/simple_project/setup.py 44 | /tests/fixtures/project_with_extras/setup.py 45 | .mypy_cache 46 | 47 | .venv 48 | /releases/* 49 | pip-wheel-metadata 50 | /poetry.toml 51 | 52 | poetry/core/* -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 disler 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Multi-Agent Postgres Data Analytics 2 | *The way we interact with our data is changing.* 3 | 4 | ![Multi-Agent Postgres Data Analytics](imgs/multi-agent-coding.png) 5 | 6 | # 💬 Read This First 💬 7 | > This repo is an **_experiment_** and **_learning tool_** for building multi-agent systems. 8 | > 9 | > It is **ONE** of **MANY** steps toward building fully autonomous, _agentic software_. 10 | > 11 | > It is **NOT** a framework, or library, or shortcut. 12 | > 13 | > It **IS** a **_stepping stone_** to help you internalize concepts, patterns and building blocks for your own multi-agent systems and applications. 14 | > 15 | > Code only tells a story at a moment in time. I highly recommend you watch the [video series](https://www.youtube.com/playlist?list=PLS_o2ayVCKvDzj2YxeFqMq9UbR1PkPEh0) to see the **how and the why** behind the structure of this experimental codebase. 16 | > 17 | > In the series we build this from scratch and dive deep into complexities, principles, patterns and ideas surrounding multi-agent software. The video order is linked below, mapping branches to videos. 18 | > 19 | > This repo will not be maintained or updated beyond the lifespan of the series. It is a snapshot in time of the code we built in the video series and is meant only to be a reference for you on your journey to building your own multi-agent systems, **_nothing more_**. 20 | > 21 | > When we complete the series will we freeze the codebase. We will then use it as a reference for experiments, products, and videos. 22 | 23 | 24 | ## 💻 Multi-Agent Postgres Data Analytics Tool 💻 25 | This is a multi-agent system that allows you to ask questions about your postgres database in natural language. 26 | 27 | The codebase is powered by GPT-4, Assistance API, AutoGen, Postgres, and Guidance. 28 | 29 | It's the first of many multi-agent applications that utilize LLMs (large language models) to enable reasoning and decision making with reduced need for explicit rules or logic. 30 | 31 | ## 💻 Setup 💻 32 | - **Read the codebase first**. Remember, this is an experiment and learning tool. It's not meant to be a framework or library. 33 | - Run `git branch -a` to view all branches. Each branch is a video in the series. 34 | - `git checkout ` you want to view. 35 | - `poetry install` 36 | - `cp .env.sample .env` 37 | - Fill out `.env` with your postgres url and openai api key 38 | - Run a prompt against your database 39 | - `poetry run start --prompt ""` 40 | - Start with something simple to get a feel for it and then build up to more complex questions. 41 | 42 | ## 🛠️ Core Tech Stack 🛠️ 43 | - [OpenAI](https://openai.com/) - GPT-4, GPT-4 Turbo, Assistance API 44 | - [AutoGen](https://microsoft.github.io/autogen/) - Multi-Agent Framework 45 | - [Postgres](https://www.postgresql.org/) - Database 46 | - [Guidance](https://github.com/guidance-ai/guidance) - Structured LLM Responses 47 | - [Aider](https://aider.chat/) - AI Pair Programming 48 | - [Poetry](https://python-poetry.org/) - Package Manager 49 | - [Python ^3.10](https://www.python.org/downloads/release/python-3100/) - Programming Language 50 | 51 | ## 🔵 Multi-Agent Patterns & Terminology 🔵 52 | Throughout the codebase we built up several existing and new patterns and terminology you've likely seen in some shape or form. Here's a quick overview of the most important ones. 53 | - **Agent** - An agent is LLM powered tool with a single purpose that can be assigned a function and/or prompt. 54 | - **Multi-Agent Team** - A collection of agents that exchange messages and work together to accomplish a goal. 55 | - **Conversations** - The exchange of messages between a multi-agent team. 56 | - **Conversation Flows** - The way agents communicate with each other. How you're agents communicate completely changes the way your application works. The conversation flow dictates which agent speaks, the order in which they speak, who they speak to and what they say. 57 | - **Orchestrator** - Manages a single agent team, their conversations and their output. Orchestrators contain different types of conversation flows. 58 | - **Instruments** - Instruments are the tools agents can use. Think of it like a front-end store. It contains state and functions that both agents and orchestrators can utilize throughout the lifecycle of the application. Agents and Orchestrators can consume and manipulate the state of instruments although typically, only agents update state. 59 | - **Decision Agents** - Agents that respond with concrete decisions which can dictate the flow of your applications. To build complex agentic systems you need agents to have the ability to make concrete decisions that then drive the flow of your application. 60 | - **Structured vs Unstructured Agents** - Structured agents are agents that respond with structured data. Unstructured agents are agents that respond with unstructured data. Structured agents are typically decision agents. 61 | 62 | ## 📺 Video Series - Learn By Watching 📺 63 | 64 | ### [Part 1 - Prompt Engineering an ENTIRE codebase: Postgres Data Analytics Al Agent](https://youtu.be/jmDMusirPKA) 65 | Branch: [v1-prompt-engineering-an-entire-codebase](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v1-prompt-engineering-an-entire-codebase) 66 | 67 | Video: [https://youtu.be/jmDMusirPKA](https://youtu.be/jmDMusirPKA) 68 | 69 | 70 | 71 | ### [Part 2 - One Prompt is NOT enough: Using AutoGen to code a Multi-Agent Postgres AI Tool](https://youtu.be/JjVvYDPVrAQ) 72 | Branch: [v2-using-autogen-to-build-our-multi-agent-tool](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v2-using-autogen-to-build-our-multi-agent-tool) 73 | 74 | Video: [https://youtu.be/JjVvYDPVrAQ](https://youtu.be/JjVvYDPVrAQ) 75 | 76 | 77 | 78 | ### [Part 3 - Make AutoGen Consistent: CONTROL your LLM agents for ACCURATE Postgres Al Data Analytics](https://youtu.be/4o8tymMQ5GM) 79 | Branch: [v3-make-autogen-consistent-control-your-llm](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v3-make-autogen-consistent-control-your-llm) 80 | 81 | Video: [https://youtu.be/4o8tymMQ5GM](https://youtu.be/4o8tymMQ5GM) 82 | 83 | 84 | 85 | ### [Part 4 - AutoGen Token Tactics: FIRING AI Agents, USELESS Vector Embeddings, GPT-4 Memory Tricks](https://youtu.be/CKo-czvxFkY) 86 | Branch: [v4-autogen-token-tactics-firing-ai-agents](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v4-autogen-token-tactics-firing-ai-agents) 87 | 88 | Video: [https://youtu.be/CKo-czvxFkY](https://youtu.be/CKo-czvxFkY) 89 | 90 | 91 | 92 | ### [Part 5 - AutoGen SPYWARE: Coding Systems for SUCCESSFUL AI Agents (Postgres Data Analytics)](https://youtu.be/UA6IVMDPuC8) 93 | Branch: [v5-autogen-spyware-coding-systems-for-successful-ai](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v5-autogen-spyware-coding-systems-for-successful-ai) 94 | 95 | Video: [https://youtu.be/UA6IVMDPuC8](https://youtu.be/UA6IVMDPuC8) 96 | 97 | 98 | 99 | ### [Part 6 - Using AUTOGEN & GUIDANCE to code LLM Control Flow & JSON Agents (No Prompt Engineering)](https://youtu.be/XGCWyfA3rgQ) 100 | Branch: [v6-control-flow-and-structured-response](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v6-control-flow-and-structured-response) 101 | 102 | Video: [https://youtu.be/XGCWyfA3rgQ](https://youtu.be/XGCWyfA3rgQ) 103 | 104 | 105 | 106 | ### [Part 7 - OpenAI Macro & Micro Strategy: Master Assistants API, Threads, Messages, and Runs](https://youtu.be/KwcrjP3vuy0) 107 | Branch: [v7-turbo4-assistants-threads-messages](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v7-turbo4-assistants-threads-messages) 108 | 109 | Video: [https://youtu.be/KwcrjP3vuy0](https://youtu.be/KwcrjP3vuy0) 110 | 111 | 112 | 113 | ### [Part 8 - Copilot Prompt Engineering: 3 UI Frameworks, 2 AI Agents, 1 Coding Assistant (AIDER CCC)](https://youtu.be/7EA19-D4-Zo) 114 | 115 | Branch: [v8-ccc-ai-engineering-with-aider](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v8-ccc-ai-engineering-with-aider) 116 | 117 | Video: [https://youtu.be/7EA19-D4-Zo](https://youtu.be/7EA19-D4-Zo) 118 | 119 | 120 | 121 | 122 | ### [Part 9 - Your AI Agents can SELF-CORRECT: Using Assistants API to AUTO FIX SQL Database Errors](https://youtu.be/Uf7cYAXe3eI) 123 | 124 | Branch: [v9-self-correcting-assistant](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v9-self-correcting-assistant) 125 | 126 | Video: [https://youtu.be/Uf7cYAXe3eI](https://youtu.be/Uf7cYAXe3eI) 127 | 128 | 129 | 130 | ### [Part 10 - Talk To Your Database - A GPT Multi Agent Postgres Data Analytics Tool](https://youtu.be/5wROK4lBoeo) 131 | 132 | Branch: [v10-talk-to-your-database-beta-launch](https://github.com/disler/multi-agent-postgres-data-analytics/tree/v10-talk-to-your-database-beta-launch) 133 | 134 | Video: [https://youtu.be/5wROK4lBoeo](https://youtu.be/5wROK4lBoeo) 135 | 136 | Talk To Your Database: [https://talktoyourdb.com](https://talktoyourdatabase.com) 137 | 138 | Exclusive Beta Launch Code: `9999` 139 | 140 | 141 | 142 | --- 143 | 144 | # 🧠 Major Learnings Throughout the Series 🧠 145 | 146 | ## 💡 Why are multi-agent applications important? 147 | - They're important because they allows us to create a more accurate model of the world. 148 | - We become orchestrators enabling less engineering level and more product level work. 149 | - They enable reasoning and decision making in a way that is more human like than ever before. 150 | - We can build systems that make decisions as we would while operating alongside us. 151 | - We can solve problems that previously required a dedicated hire or an entire team to solve. 152 | 153 | ## ✅ Multi-Agent Systems: The Good 154 | - Can assign functions & prompts to specific agents, enabling specialization yielding better results. 155 | - Agents can reflect on results to provide feedback thus improving the results. 156 | - Can role play real organizational structures, existing and new. 157 | - Ecosystem is evolving rapidly. New tools and frameworks are being built every day. 158 | - Upside potential is ridiculously massive. We're talking asymmetric ROI, max [leverage](https://www.navalmanack.com/almanack-of-naval-ravikant/find-a-position-of-leverage), [superlinear](http://www.paulgraham.com/superlinear.html) upside. The more agentic build blocks you have the more powerful your engineering and product potential becomes. 159 | - Multi-agent engineering is probably the most important thing happening in software right now (2023-2024). 160 | - The road to agentic software is clear. Solve small problems, create reusable building blocks, and then combine them to solve bigger problems. 161 | - GPT-4 can support multi-agent systems without a doubt. It is the best model by light-years and drives incredible reasoning readily available at your fingertips. 162 | 163 | ## ❌ Multi-Agent Systems: The Bad 164 | - It's an art to get the roles and the function of your agent right. How many do you need? What are they? How do you know? 165 | - Can get expensive in testing and scales with # of agents. The more agents the more expensive each query is. 166 | - Can be difficult to debug why a multi-agent system is not working as expected due to the non-deterministic nature of LLMs. 167 | - Memory management is a major issue. The context window is forcing a lot of weird, intricate code to manage memory. 168 | - Too much noise and hype in the AI Agent ecosystem. Lot's of clickbait hype with little follow through value. Hard to find good resources. 169 | - Very few are engineers are publicly building multi-agent systems. Most are toy examples or ripping from example codebases. 170 | - OpenAI is inadvertently killing startups with every new release. Risky to commit to building LLM powered applications. 171 | - At the current price, we cannot run a fully agentic system that runs 24/7 or even for an hour on GPT-4 without burning thousands per day. The price must come down WITHOUT sacrificing quality (looking at you open source models). 172 | - It's tricky to know when to write explicit code vs prompt engineer vs build a multi-agent team. This is a new skill that will take time to master. 173 | 174 | ## 🧠 2024 Multi-agent / LLM / Agentic Predictions 🧠 175 | ![2024-predictions](imgs/2024-predictions.png) 176 | - [2024 Predictions Video (Recommended)](https://youtu.be/UES89QRc3Sk) 177 | - [2024 Predictions Blog Post](https://indydevdan.com/blogs/2024-predictions) 178 | - [2024 Predictions Slides](https://firebasestorage.googleapis.com/v0/b/solopreneur-d8361.appspot.com/o/solopreneur%2Fblog%2Fimages%2F2024-predictions-for-advanced-ai-llm-engineers-small.pdf?alt=media&token=44c39c93-8e35-45c4-92c0-3198721f6081) -------------------------------------------------------------------------------- /api-server/.env.sample: -------------------------------------------------------------------------------- 1 | DATABASE_URL= 2 | OPENAI_API_KEY= 3 | -------------------------------------------------------------------------------- /api-server/.gitignore: -------------------------------------------------------------------------------- 1 | .aider.ident.cache.v1 2 | .aider.tags.cache.v1 3 | .aider.chat.history.md 4 | .aider.input.history 5 | 6 | agent_results/* 7 | 8 | .vercel 9 | *.log 10 | *.pyc 11 | __pycache__ 12 | 13 | # Environments 14 | .env 15 | .venv 16 | env/ 17 | venv/ 18 | ENV/ 19 | env.bak/ 20 | venv.bak/ -------------------------------------------------------------------------------- /api-server/README.md: -------------------------------------------------------------------------------- 1 | ### Setup 2 | - Copy .env.sample to .env and fill in the values 3 | - `cp .env.sample .env` 4 | - Create python environment 5 | - `python -m venv venv` 6 | - Activate python environment 7 | - `source venv/bin/activate` 8 | - Install dependencies 9 | - `pip install -r requirements.txt` 10 | - Run the server 11 | - `python api/index.py` -------------------------------------------------------------------------------- /api-server/api/index.py: -------------------------------------------------------------------------------- 1 | import json 2 | from flask import Flask, Request, Response, jsonify, request, make_response 3 | import dotenv 4 | from modules import db, llm, emb, instruments 5 | from modules.turbo4 import Turbo4 6 | 7 | import os 8 | 9 | from modules.models import TurboTool 10 | from psycopg2 import Error as PostgresError 11 | 12 | app = Flask(__name__) 13 | 14 | # ---------------- .Env Constants ---------------- 15 | 16 | dotenv.load_dotenv() 17 | 18 | assert os.environ.get("DATABASE_URL"), "POSTGRES_CONNECTION_URL not found in .env file" 19 | assert os.environ.get( 20 | "OPENAI_API_KEY" 21 | ), "POSTGRES_CONNECTION_URL not found in .env file" 22 | 23 | 24 | DB_URL = os.environ.get("DATABASE_URL") 25 | OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") 26 | 27 | # ---------------- Cors Helper ---------------- 28 | 29 | 30 | def make_cors_response(): 31 | # Set CORS headers for the preflight request 32 | response = make_response() 33 | response.headers.add("Access-Control-Allow-Origin", "*") 34 | response.headers.add("Access-Control-Allow-Headers", "Content-Type,Authorization") 35 | response.headers.add("Access-Control-Allow-Methods", "GET,PUT,POST,DELETE,OPTIONS") 36 | return response 37 | 38 | 39 | # ---------------- Self Correcting Assistant ---------------- 40 | 41 | 42 | def self_correcting_assistant( 43 | db: db.PostgresManager, 44 | agent_instruments: instruments.AgentInstruments, 45 | tools: TurboTool, 46 | error: PostgresError, 47 | ): 48 | # reset db - to unblock transactions 49 | db.roll_back() 50 | 51 | all_table_definitions = db.get_table_definitions_for_prompt() 52 | 53 | print(f"Loaded all table definitions") 54 | 55 | # ------ File prep 56 | 57 | file_path = agent_instruments.self_correcting_table_def_file 58 | 59 | # write all_table_definitions to file 60 | with open(file_path, "w") as f: 61 | f.write(all_table_definitions) 62 | 63 | files_to_upload = [file_path] 64 | 65 | sql_query = open(agent_instruments.sql_query_file).read() 66 | 67 | # ------ Prompts 68 | 69 | output_file_path = agent_instruments.run_sql_results_file 70 | 71 | diagnosis_prompt = f"Given the table_definitions.sql file, the following SQL_ERROR, and the SQL_QUERY, describe the most likely cause of the error. Think step by step.\n\nSQL_ERROR: {error}\n\nSQL_QUERY: {sql_query}" 72 | 73 | generation_prompt = ( 74 | f"Based on your diagnosis, generate a new SQL query that will run successfully." 75 | ) 76 | 77 | run_sql_prompt = "Use the run_sql function to run the SQL you've just generated." 78 | 79 | assistant_name = "SQL Self Correction" 80 | 81 | turbo4_assistant = Turbo4().get_or_create_assistant(assistant_name) 82 | 83 | print(f"Generated Assistant: {assistant_name}") 84 | 85 | file_ids = turbo4_assistant.upsert_files(files_to_upload) 86 | 87 | print(f"Uploaded files: {file_ids}") 88 | 89 | print(f"Running Self Correction Assistant...") 90 | 91 | ( 92 | turbo4_assistant.set_instructions( 93 | "You're an elite SQL developer. You generate the most concise and performant SQL queries. You review failed queries and generate new SQL queries to fix them." 94 | ) 95 | .enable_retrieval() 96 | .equip_tools(tools) 97 | .make_thread() 98 | # 1/3 STEP PATTERN: diagnose 99 | .add_message(diagnosis_prompt, file_ids=file_ids) 100 | .run_thread() 101 | .spy_on_assistant(agent_instruments.make_agent_chat_file(assistant_name)) 102 | # 2/3 STEP PATTERN: generate 103 | .add_message(generation_prompt) 104 | .run_thread() 105 | .spy_on_assistant(agent_instruments.make_agent_chat_file(assistant_name)) 106 | # 3/3 STEP PATTERN: execute 107 | .add_message(run_sql_prompt) 108 | .run_thread(toolbox=[tools[0].name]) 109 | .spy_on_assistant(agent_instruments.make_agent_chat_file(assistant_name)) 110 | # clean up, logging, reporting, cost 111 | .run_validation(agent_instruments.validate_file_exists(output_file_path)) 112 | .spy_on_assistant(agent_instruments.make_agent_chat_file(assistant_name)) 113 | .get_costs_and_tokens(agent_instruments.make_agent_cost_file(assistant_name)) 114 | ) 115 | 116 | pass 117 | 118 | 119 | # ---------------- Primary Endpoint ---------------- 120 | 121 | 122 | @app.route("/prompt", methods=["POST", "OPTIONS"]) 123 | def prompt(): 124 | # Set CORS headers for the main request 125 | response = make_cors_response() 126 | if request.method == "OPTIONS": 127 | return response 128 | 129 | # Get access to db, state, and functions 130 | with instruments.PostgresAgentInstruments(DB_URL, "prompt-endpoint") as ( 131 | agent_instruments, 132 | db, 133 | ): 134 | # ---------------- Build Prompt ---------------- 135 | 136 | base_prompt = request.json["prompt"] 137 | 138 | # simple word match for now - dropped embeddings for deployment size 139 | similar_tables = emb.DatabaseEmbedder(db).get_similar_table_defs_for_prompt( 140 | base_prompt 141 | ) 142 | 143 | if len(similar_tables) == 0: 144 | print(f"No similar tables found for prompt: {base_prompt}") 145 | response.status_code = 400 146 | response.data = "No similar tables found." 147 | return response 148 | 149 | print("similar_tables", similar_tables) 150 | 151 | print(f"base_prompt: {base_prompt}") 152 | 153 | prompt = f"Fulfill this database query: {base_prompt}. " 154 | prompt = llm.add_cap_ref( 155 | prompt, 156 | f"Use these TABLE_DEFINITIONS to satisfy the database query.", 157 | "TABLE_DEFINITIONS", 158 | similar_tables, 159 | ) 160 | 161 | # ---------------- Run 2 Agent Team - Generate SQL & Results ---------------- 162 | 163 | tools = [ 164 | TurboTool("run_sql", llm.run_sql_tool_config, agent_instruments.run_sql), 165 | ] 166 | 167 | sql_response = llm.prompt( 168 | prompt, 169 | model="gpt-4-1106-preview", 170 | instructions="You're an elite SQL developer. You generate the most concise and performant SQL queries.", 171 | ) 172 | try: 173 | llm.prompt_func( 174 | "Use the run_sql function to run the SQL you've just generated: " 175 | + sql_response, 176 | model="gpt-4-1106-preview", 177 | instructions="You're an elite SQL developer. You generate the most concise and performant SQL queries.", 178 | turbo_tools=tools, 179 | ) 180 | agent_instruments.validate_run_sql() 181 | except PostgresError as e: 182 | print( 183 | f"Received PostgresError -> Running Self Correction Team To Resolve: {e}" 184 | ) 185 | 186 | # ---------------- Run Self Correction Team - Diagnosis, Generate New SQL, Retry ---------------- 187 | self_correcting_assistant(db, agent_instruments, tools, e) 188 | 189 | print(f"Self Correction Team Complete.") 190 | 191 | # ---------------- Read result files and respond ---------------- 192 | 193 | sql_query = open(agent_instruments.sql_query_file).read() 194 | sql_query_results = open(agent_instruments.run_sql_results_file).read() 195 | 196 | response_obj = { 197 | "prompt": base_prompt, 198 | "results": sql_query_results, 199 | "sql": sql_query, 200 | } 201 | 202 | print("response_obj", response_obj) 203 | 204 | response.data = json.dumps(response_obj) 205 | 206 | return response 207 | 208 | 209 | if __name__ == "__main__": 210 | port = 3000 211 | print(f"Starting server on port {port}") 212 | app.run(debug=True, port=port) 213 | -------------------------------------------------------------------------------- /api-server/api/modules/db.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | import json 3 | import psycopg2 4 | from psycopg2.sql import SQL, Identifier 5 | 6 | 7 | # comm 8 | class PostgresManager: 9 | """ 10 | A class to manage postgres connections and queries 11 | """ 12 | 13 | def __init__(self): 14 | self.conn = None 15 | self.cur = None 16 | 17 | def __enter__(self): 18 | return self 19 | 20 | def __exit__(self, exc_type, exc_val, exc_tb): 21 | if self.cur: 22 | self.cur.close() 23 | if self.conn: 24 | self.conn.close() 25 | 26 | def connect_with_url(self, url): 27 | self.conn = psycopg2.connect(url) 28 | self.cur = self.conn.cursor() 29 | 30 | def close(self): 31 | if self.cur: 32 | self.cur.close() 33 | if self.conn: 34 | self.conn.close() 35 | 36 | def run_sql(self, sql) -> str: 37 | """ 38 | Run a SQL query against the postgres database 39 | """ 40 | self.cur.execute(sql) 41 | columns = [desc[0] for desc in self.cur.description] 42 | res = self.cur.fetchall() 43 | 44 | list_of_dicts = [dict(zip(columns, row)) for row in res] 45 | 46 | json_result = json.dumps(list_of_dicts, indent=4, default=self.datetime_handler) 47 | 48 | return json_result 49 | 50 | def datetime_handler(self, obj): 51 | """ 52 | Handle datetime objects when serializing to JSON. 53 | """ 54 | if isinstance(obj, datetime): 55 | return obj.isoformat() 56 | return str(obj) # or just return the object unchanged, or another default value 57 | 58 | def get_table_definition(self, table_name): 59 | """ 60 | Generate the 'create' definition for a table 61 | """ 62 | 63 | get_def_stmt = """ 64 | SELECT pg_class.relname as tablename, 65 | pg_attribute.attnum, 66 | pg_attribute.attname, 67 | format_type(atttypid, atttypmod) 68 | FROM pg_class 69 | JOIN pg_namespace ON pg_namespace.oid = pg_class.relnamespace 70 | JOIN pg_attribute ON pg_attribute.attrelid = pg_class.oid 71 | WHERE pg_attribute.attnum > 0 72 | AND pg_class.relname = %s 73 | AND pg_namespace.nspname = 'public' -- Assuming you're interested in public schema 74 | """ 75 | self.cur.execute(get_def_stmt, (table_name,)) 76 | rows = self.cur.fetchall() 77 | create_table_stmt = "CREATE TABLE {} (\n".format(table_name) 78 | for row in rows: 79 | create_table_stmt += "{} {},\n".format(row[2], row[3]) 80 | create_table_stmt = create_table_stmt.rstrip(",\n") + "\n);" 81 | return create_table_stmt 82 | 83 | def get_all_table_names(self): 84 | """ 85 | Get all table names in the database 86 | """ 87 | get_all_tables_stmt = ( 88 | "SELECT tablename FROM pg_tables WHERE schemaname = 'public';" 89 | ) 90 | self.cur.execute(get_all_tables_stmt) 91 | return [row[0] for row in self.cur.fetchall()] 92 | 93 | def get_table_definitions_for_prompt(self): 94 | """ 95 | Get all table 'create' definitions in the database 96 | """ 97 | table_names = self.get_all_table_names() 98 | definitions = [] 99 | for table_name in table_names: 100 | definitions.append(self.get_table_definition(table_name)) 101 | return "\n\n".join(definitions) 102 | 103 | def get_table_definition_map_for_embeddings(self): 104 | """ 105 | Creates a map of table names to table definitions 106 | """ 107 | table_names = self.get_all_table_names() 108 | definitions = {} 109 | for table_name in table_names: 110 | definitions[table_name] = self.get_table_definition(table_name) 111 | return definitions 112 | 113 | def get_related_tables(self, table_list, n=2): 114 | """ 115 | Get tables that have foreign keys referencing the given table 116 | """ 117 | 118 | related_tables_dict = {} 119 | 120 | for table in table_list: 121 | # Query to fetch tables that have foreign keys referencing the given table 122 | self.cur.execute( 123 | """ 124 | SELECT 125 | a.relname AS table_name 126 | FROM 127 | pg_constraint con 128 | JOIN pg_class a ON a.oid = con.conrelid 129 | WHERE 130 | confrelid = (SELECT oid FROM pg_class WHERE relname = %s) 131 | LIMIT %s; 132 | """, 133 | (table, n), 134 | ) 135 | 136 | related_tables = [row[0] for row in self.cur.fetchall()] 137 | 138 | # Query to fetch tables that the given table references 139 | self.cur.execute( 140 | """ 141 | SELECT 142 | a.relname AS referenced_table_name 143 | FROM 144 | pg_constraint con 145 | JOIN pg_class a ON a.oid = con.confrelid 146 | WHERE 147 | conrelid = (SELECT oid FROM pg_class WHERE relname = %s) 148 | LIMIT %s; 149 | """, 150 | (table, n), 151 | ) 152 | 153 | related_tables += [row[0] for row in self.cur.fetchall()] 154 | 155 | related_tables_dict[table] = related_tables 156 | 157 | # convert dict to list and remove dups 158 | related_tables_list = [] 159 | for table, related_tables in related_tables_dict.items(): 160 | related_tables_list += related_tables 161 | 162 | related_tables_list = list(set(related_tables_list)) 163 | 164 | return related_tables_list 165 | 166 | def roll_back(self): 167 | self.conn.rollback() 168 | -------------------------------------------------------------------------------- /api-server/api/modules/emb.py: -------------------------------------------------------------------------------- 1 | from modules.db import PostgresManager 2 | 3 | 4 | class DatabaseEmbedder: 5 | """ 6 | This class is responsible for embedding database table definitions and 7 | computing similarity between user queries and table definitions. 8 | """ 9 | 10 | def __init__(self, db: PostgresManager): 11 | self.map_name_to_embeddings = {} 12 | self.map_name_to_table_def = {} 13 | self.db = db 14 | 15 | def get_similar_table_defs_for_prompt(self, prompt: str, n_similar=5, n_foreign=0): 16 | map_table_name_to_table_def = self.db.get_table_definition_map_for_embeddings() 17 | for name, table_def in map_table_name_to_table_def.items(): 18 | self.add_table(name, table_def) 19 | 20 | similar_tables = self.get_similar_tables(prompt, n=n_similar) 21 | 22 | table_definitions = self.get_table_definitions_from_names(similar_tables) 23 | 24 | if n_foreign > 0: 25 | foreign_table_names = self.db.get_foreign_tables(similar_tables, n=3) 26 | 27 | table_definitions = self.get_table_definitions_from_names( 28 | foreign_table_names + similar_tables 29 | ) 30 | 31 | return table_definitions 32 | 33 | def add_table(self, table_name: str, text_representation: str): 34 | """ 35 | Add a table to the database embedder. 36 | Map the table name to its embedding and text representation. 37 | """ 38 | self.map_name_to_embeddings[table_name] = self.compute_embeddings( 39 | text_representation 40 | ) 41 | 42 | self.map_name_to_table_def[table_name] = text_representation 43 | 44 | def compute_embeddings(self, text): 45 | """ 46 | Compute embeddings for a given text using the BERT model. 47 | """ 48 | return "" 49 | 50 | def get_similar_tables_via_embeddings(self, query, n=3): 51 | """ 52 | Given a query, find the top 'n' tables that are most similar to it. 53 | 54 | Args: 55 | - query (str): The user's natural language query. 56 | - n (int, optional): Number of top tables to return. Defaults to 3. 57 | 58 | Returns: 59 | - list: Top 'n' table names ranked by their similarity to the query. 60 | """ 61 | return [] 62 | 63 | def get_similar_table_names_via_word_match(self, query: str): 64 | """ 65 | if any word in our query is a table name, add the table to a list 66 | """ 67 | 68 | tables = [] 69 | 70 | for table_name in self.map_name_to_table_def.keys(): 71 | if table_name.lower() in query.lower(): 72 | tables.append(table_name) 73 | 74 | return tables 75 | 76 | def get_similar_tables(self, query: str, n=3): 77 | """ 78 | combines results from get_similar_tables_via_embeddings and get_similar_table_names_via_word_match 79 | """ 80 | 81 | similar_tables_via_embeddings = self.get_similar_tables_via_embeddings(query, n) 82 | similar_tables_via_word_match = self.get_similar_table_names_via_word_match( 83 | query 84 | ) 85 | 86 | return similar_tables_via_embeddings + similar_tables_via_word_match 87 | 88 | def get_table_definitions_from_names(self, table_names: list) -> str: 89 | """ 90 | Given a list of table names, return their table definitions. 91 | """ 92 | table_defs = [ 93 | self.map_name_to_table_def[table_name] for table_name in table_names 94 | ] 95 | return "\n\n".join(table_defs) 96 | -------------------------------------------------------------------------------- /api-server/api/modules/file.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | 4 | def write_file(fname, content): 5 | with open(fname, "w") as f: 6 | f.write(content) 7 | 8 | 9 | def write_json_file(fname, json_str: str): 10 | # convert ' to " 11 | json_str = json_str.replace("'", '"') 12 | 13 | # Convert the string to a Python object 14 | data = json.loads(json_str) 15 | 16 | # Write the Python object to the file as JSON 17 | with open(fname, "w") as f: 18 | json.dump(data, f, indent=4) 19 | -------------------------------------------------------------------------------- /api-server/api/modules/instruments.py: -------------------------------------------------------------------------------- 1 | import json 2 | from modules.db import PostgresManager 3 | from modules import file 4 | import os 5 | 6 | BASE_DIR = os.environ.get("BASE_DIR", "./agent_results") 7 | 8 | 9 | class AgentInstruments: 10 | """ 11 | Base class for multli-agent instruments that are tools, state, and functions that an agent can use across the lifecycle of conversations 12 | """ 13 | 14 | def __init__(self) -> None: 15 | self.session_id = None 16 | self.messages = [] 17 | 18 | def __enter__(self): 19 | return self 20 | 21 | def __exit__(self, exc_type, exc_value, traceback): 22 | pass 23 | 24 | def sync_messages(self, messages: list): 25 | """ 26 | Syncs messages with the orchestrator 27 | """ 28 | raise NotImplementedError 29 | 30 | def make_agent_chat_file(self, team_name: str): 31 | return os.path.join(self.root_dir, f"agent_chats_{team_name}.json") 32 | 33 | def make_agent_cost_file(self, team_name: str): 34 | return os.path.join(self.root_dir, f"agent_cost_{team_name}.json") 35 | 36 | @property 37 | def root_dir(self): 38 | return os.path.join(BASE_DIR, self.session_id) 39 | 40 | 41 | class PostgresAgentInstruments(AgentInstruments): 42 | """ 43 | Unified Toolset for the Postgres Data Analytics Multi-Agent System 44 | 45 | Advantages: 46 | - All agents have access to the same state and functions 47 | - Gives agent functions awareness of changing context 48 | - Clear and concise capabilities for agents 49 | - Clean database connection management 50 | 51 | Guidelines: 52 | - Agent Functions should not call other agent functions directly 53 | - Instead Agent Functions should call external lower level modules 54 | - Prefer 1 to 1 mapping of agents and their functions 55 | - The state lifecycle lives between all agent orchestrations 56 | """ 57 | 58 | def __init__(self, db_url: str, session_id: str) -> None: 59 | super().__init__() 60 | 61 | self.db_url = db_url 62 | self.db = None 63 | self.session_id = session_id 64 | self.messages = [] 65 | self.innovation_index = 0 66 | 67 | def __enter__(self): 68 | """ 69 | Support entering the 'with' statement 70 | """ 71 | self.reset_files() 72 | self.db = PostgresManager() 73 | self.db.connect_with_url(self.db_url) 74 | return self, self.db 75 | 76 | def __exit__(self, exc_type, exc_val, exc_tb): 77 | """ 78 | Support exiting the 'with' statement 79 | """ 80 | self.db.close() 81 | 82 | def sync_messages(self, messages: list): 83 | """ 84 | Syncs messages with the orchestrator 85 | """ 86 | self.messages = messages 87 | 88 | def reset_files(self): 89 | """ 90 | Clear everything in the root_dir 91 | """ 92 | 93 | # if it does not exist create it 94 | if not os.path.exists(self.root_dir): 95 | os.makedirs(self.root_dir) 96 | 97 | for fname in os.listdir(self.root_dir): 98 | os.remove(os.path.join(self.root_dir, fname)) 99 | 100 | def get_file_path(self, fname: str): 101 | """ 102 | Get the full path to a file in the root_dir 103 | """ 104 | return os.path.join(self.root_dir, fname) 105 | 106 | # -------------------------- Agent Properties -------------------------- # 107 | 108 | @property 109 | def run_sql_results_file(self): 110 | return self.get_file_path("run_sql_results.json") 111 | 112 | @property 113 | def sql_query_file(self): 114 | return self.get_file_path("sql_query.sql") 115 | 116 | @property 117 | def self_correcting_table_def_file(self): 118 | return self.get_file_path("table_definitions.sql") 119 | 120 | # -------------------------- Agent Functions -------------------------- # 121 | 122 | def run_sql(self, sql: str) -> str: 123 | """ 124 | Run a SQL query against the postgres database 125 | """ 126 | 127 | with open(self.sql_query_file, "w") as f: 128 | f.write(sql) 129 | 130 | results_as_json = self.db.run_sql(sql) 131 | 132 | fname = self.run_sql_results_file 133 | 134 | # dump these results to a file 135 | with open(fname, "w") as f: 136 | f.write(results_as_json) 137 | 138 | return "Successfully delivered results to json file" 139 | 140 | def validate_run_sql(self): 141 | """ 142 | validate that the run_sql results file exists and has content 143 | """ 144 | fname = self.run_sql_results_file 145 | 146 | with open(fname, "r") as f: 147 | content = f.read() 148 | 149 | if not content: 150 | return False, f"File {fname} is empty" 151 | 152 | return True, "" 153 | 154 | def write_file(self, content: str): 155 | fname = self.get_file_path(f"write_file.txt") 156 | return file.write_file(fname, content) 157 | 158 | def write_json_file(self, json_str: str): 159 | fname = self.get_file_path(f"write_json_file.json") 160 | return file.write_json_file(fname, json_str) 161 | 162 | def write_innovation_file(self, content: str): 163 | fname = self.get_file_path(f"{self.innovation_index}_innovation_file.json") 164 | file.write_file(fname, content) 165 | self.innovation_index += 1 166 | return f"Successfully wrote innovation file. You can check my work." 167 | 168 | def validate_innovation_files(self): 169 | """ 170 | loop from 0 to innovation_index and verify file exists with content 171 | """ 172 | for i in range(self.innovation_index): 173 | fname = self.get_file_path(f"{i}_innovation_file.json") 174 | with open(fname, "r") as f: 175 | content = f.read() 176 | if not content: 177 | return False, f"File {fname} is empty" 178 | 179 | return True, "" 180 | 181 | def validate_file_exists(self, file: str): 182 | def file_exists(): 183 | if not os.path.exists(file): 184 | raise Exception(f"File {file} does not exist") 185 | 186 | return file_exists 187 | -------------------------------------------------------------------------------- /api-server/api/modules/llm.py: -------------------------------------------------------------------------------- 1 | """ 2 | Purpose: 3 | Interact with the OpenAI API. 4 | Provide supporting prompt engineering functions. 5 | """ 6 | 7 | import json 8 | import sys 9 | from dotenv import load_dotenv 10 | import os 11 | from typing import Any, Dict, List 12 | import openai 13 | 14 | from modules.models import TurboTool 15 | 16 | # load .env file 17 | load_dotenv() 18 | 19 | assert os.environ.get("OPENAI_API_KEY") 20 | 21 | # get openai api key 22 | openai.api_key = os.environ.get("OPENAI_API_KEY") 23 | 24 | 25 | run_sql_tool_config = { 26 | "type": "function", 27 | "function": { 28 | "name": "run_sql", 29 | "description": "Run a SQL query against the postgres database", 30 | "parameters": { 31 | "type": "object", 32 | "properties": { 33 | "sql": { 34 | "type": "string", 35 | "description": "The SQL query to run", 36 | } 37 | }, 38 | "required": ["sql"], 39 | }, 40 | }, 41 | } 42 | 43 | 44 | # ------------------ helpers ------------------ 45 | 46 | 47 | def safe_get(data, dot_chained_keys): 48 | """ 49 | {'a': {'b': [{'c': 1}]}} 50 | safe_get(data, 'a.b.0.c') -> 1 51 | """ 52 | keys = dot_chained_keys.split(".") 53 | for key in keys: 54 | try: 55 | if isinstance(data, list): 56 | data = data[int(key)] 57 | else: 58 | data = data[key] 59 | except (KeyError, TypeError, IndexError): 60 | return None 61 | return data 62 | 63 | 64 | def response_parser(response: Dict[str, Any]): 65 | return safe_get(response, "choices.0.message.content") 66 | 67 | 68 | # ------------------ content generators ------------------ 69 | 70 | 71 | def prompt( 72 | prompt: str, 73 | model: str = "gpt-4-1106-preview", 74 | instructions: str = "You are a helpful assistant.", 75 | ) -> str: 76 | """ 77 | Generate a response from a prompt using the OpenAI API. 78 | """ 79 | 80 | if not openai.api_key: 81 | sys.exit( 82 | """ 83 | ERORR: OpenAI API key not found. Please export your key to OPENAI_API_KEY 84 | Example bash command: 85 | export OPENAI_API_KEY= 86 | """ 87 | ) 88 | 89 | response = openai.chat.completions.create( 90 | model=model, 91 | messages=[ 92 | { 93 | "role": "system", 94 | "content": instructions, # Added instructions as a system message 95 | }, 96 | { 97 | "role": "user", 98 | "content": prompt, 99 | }, 100 | ], 101 | ) 102 | 103 | return response_parser(response.model_dump()) 104 | 105 | 106 | def prompt_func( 107 | prompt: str, 108 | turbo_tools: List[TurboTool], 109 | model: str = "gpt-4-1106-preview", 110 | instructions: str = "You are a helpful assistant.", 111 | ) -> str: 112 | """ 113 | Generate a response from a prompt using the OpenAI API. 114 | Force function calls to the provided turbo tools. 115 | 116 | :param prompt: The prompt to send to the model. 117 | :param turbo_tools: List of TurboTool objects each containing the tool's name, configuration, and function. 118 | :param model: The model version to use, default is 'gpt-4-1106-preview'. 119 | :return: The response generated by the model. 120 | """ 121 | 122 | messages = [{"role": "user", "content": prompt}] 123 | tools = [turbo_tool.config for turbo_tool in turbo_tools] 124 | 125 | tool_choice = ( 126 | "auto" 127 | if len(turbo_tools) > 1 128 | else {"type": "function", "function": {"name": turbo_tools[0].name}} 129 | ) 130 | 131 | messages.insert( 132 | 0, {"role": "system", "content": instructions} 133 | ) # Insert instructions as the first system message 134 | response = openai.chat.completions.create( 135 | model=model, messages=messages, tools=tools, tool_choice=tool_choice 136 | ) 137 | 138 | response_message = response.choices[0].message 139 | tool_calls = response_message.tool_calls 140 | 141 | func_responses = [] 142 | 143 | if tool_calls: 144 | messages.append(response_message) 145 | 146 | for tool_call in tool_calls: 147 | for turbo_tool in turbo_tools: 148 | if tool_call.function.name == turbo_tool.name: 149 | function_response = turbo_tool.function( 150 | **json.loads(tool_call.function.arguments) 151 | ) 152 | 153 | func_responses.append(function_response) 154 | 155 | message_to_append = { 156 | "tool_call_id": tool_call.id, 157 | "role": "tool", 158 | "name": turbo_tool.name, 159 | "content": function_response, 160 | } 161 | messages.append(message_to_append) 162 | break 163 | 164 | return func_responses 165 | 166 | 167 | def prompt_json_response( 168 | prompt: str, 169 | model: str = "gpt-4-1106-preview", 170 | instructions: str = "You are a helpful assistant.", 171 | ) -> str: 172 | """ 173 | Generate a response from a prompt using the OpenAI API. 174 | 175 | Example: 176 | res = llm.prompt_json_response( 177 | f"You're a data innovator. You analyze SQL databases table structure and generate 3 novel insights for your team to reflect on and query. 178 | Generate insights for this this prompt: {prompt}. 179 | Format your insights in JSON format. Respond in this json format [{{insight, sql, actionable_business_value}}, ...]", 180 | ) 181 | """ 182 | 183 | if not openai.api_key: 184 | sys.exit( 185 | """ 186 | ERORR: OpenAI API key not found. Please export your key to OPENAI_API_KEY 187 | Example bash command: 188 | export OPENAI_API_KEY= 189 | """ 190 | ) 191 | 192 | response = openai.chat.completions.create( 193 | model=model, 194 | messages=[ 195 | { 196 | "role": "system", 197 | "content": instructions, # Added instructions as a system message 198 | }, 199 | { 200 | "role": "user", 201 | "content": prompt, 202 | }, 203 | ], 204 | response_format={"type": "json_object"}, 205 | ) 206 | 207 | return response_parser(response.model_dump()) 208 | 209 | 210 | def add_cap_ref( 211 | prompt: str, prompt_suffix: str, cap_ref: str, cap_ref_content: str 212 | ) -> str: 213 | """ 214 | Attaches a capitalized reference to the prompt. 215 | Example 216 | prompt = 'Refactor this code.' 217 | prompt_suffix = 'Make it more readable using this EXAMPLE.' 218 | cap_ref = 'EXAMPLE' 219 | cap_ref_content = 'def foo():\n return True' 220 | returns 'Refactor this code. Make it more readable using this EXAMPLE.\n\nEXAMPLE\n\ndef foo():\n return True' 221 | """ 222 | 223 | new_prompt = f"""{prompt} {prompt_suffix}\n\n{cap_ref}\n\n{cap_ref_content}""" 224 | 225 | return new_prompt 226 | 227 | 228 | def count_tokens(text: str): 229 | """ 230 | Count the number of tokens in a string. 231 | """ 232 | return len(text) * 1.3 233 | 234 | 235 | map_model_to_cost_per_1k_tokens = { 236 | "gpt-4": 0.075, # ($0.03 Input Tokens + $0.06 Output Tokens) / 2 237 | "gpt-4-1106-preview": 0.02, # ($0.01 Input Tokens + $0.03 Output Tokens) / 2 238 | "gpt-4-1106-vision-preview": 0.02, # ($0.01 Input Tokens + $0.03 Output Tokens) / 2 239 | "gpt-3.5-turbo-1106": 0.0015, # ($0.001 Input Tokens + $0.002 Output Tokens) / 2 240 | } 241 | 242 | 243 | def estimate_price_and_tokens(text, model="gpt-4"): 244 | """ 245 | Conservative estimate the price and tokens for a given text. 246 | """ 247 | # round up to the output tokens 248 | COST_PER_1k_TOKENS = map_model_to_cost_per_1k_tokens[model] 249 | 250 | tokens = count_tokens(text) 251 | 252 | estimated_cost = (tokens / 1000) * COST_PER_1k_TOKENS 253 | 254 | # round 255 | estimated_cost = round(estimated_cost, 2) 256 | 257 | return estimated_cost, tokens 258 | -------------------------------------------------------------------------------- /api-server/api/modules/models.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass, field 2 | import time 3 | from typing import Callable 4 | 5 | 6 | @dataclass 7 | class TurboTool: 8 | name: str 9 | config: dict 10 | function: Callable 11 | 12 | 13 | @dataclass 14 | class Chat: 15 | from_name: str 16 | to_name: str 17 | message: str 18 | created: int = field(default_factory=time.time) 19 | -------------------------------------------------------------------------------- /api-server/api/modules/turbo4.py: -------------------------------------------------------------------------------- 1 | """ 2 | Clone of postgres_da_ai_agent/agents/turbo4.py 3 | """ 4 | 5 | import json 6 | import os 7 | import openai 8 | import time 9 | from typing import Callable, Dict, Any, List, Optional, Union, Tuple 10 | import dotenv 11 | from dataclasses import dataclass, asdict 12 | from openai import OpenAI 13 | from openai.types.beta import Thread, Assistant 14 | from openai.types import FileObject 15 | from openai.types.beta.threads.thread_message import ThreadMessage 16 | from openai.types.beta.threads.run_submit_tool_outputs_params import ToolOutput 17 | from modules import llm 18 | from modules.models import Chat, TurboTool 19 | 20 | dotenv.load_dotenv() 21 | 22 | 23 | class Turbo4: 24 | """ 25 | Simple, chainable class for the OpenAI's GPT-4 Assistant APIs. 26 | """ 27 | 28 | def __init__(self): 29 | openai.api_key = os.environ.get("OPENAI_API_KEY") 30 | self.client: openai = OpenAI() 31 | 32 | self.map_function_tools: Dict[str, TurboTool] = {} 33 | self.current_thread_id = None 34 | self.thread_messages: List[ThreadMessage] = [] 35 | self.local_messages = [] 36 | self.file_ids = [] 37 | self.assistant_id = None 38 | self.polling_interval = ( 39 | 0.5 # Interval in seconds to poll the API for thread run completion 40 | ) 41 | self.model = "gpt-4-1106-preview" 42 | 43 | @property 44 | def chat_messages(self) -> List[Chat]: 45 | return [ 46 | Chat( 47 | from_name=msg.role, 48 | to_name="assistant" if msg.role == "user" else "user", 49 | message=llm.safe_get(msg.model_dump(), "content.0.text.value"), 50 | created=msg.created_at, 51 | ) 52 | for msg in self.thread_messages 53 | ] 54 | 55 | @property 56 | def tool_config(self): 57 | return [tool.config for tool in self.map_function_tools.values()] 58 | 59 | # ------------- Additional Utility Functions ----------------- 60 | 61 | def run_validation(self, validation_func: Callable): 62 | print(f"run_validation({validation_func.__name__})") 63 | validation_func() 64 | return self 65 | 66 | def spy_on_assistant(self, output_file: str): 67 | sorted_messages = sorted( 68 | self.chat_messages, key=lambda msg: msg.created, reverse=False 69 | ) 70 | messages_as_json = [asdict(msg) for msg in sorted_messages] 71 | with open(output_file, "w") as f: 72 | json.dump(messages_as_json, f, indent=2) 73 | 74 | return self 75 | 76 | def get_costs_and_tokens(self, output_file: str) -> Tuple[float, float]: 77 | """ 78 | Get the estimated cost and token usage for the current thread. 79 | 80 | https://openai.com/pricing 81 | 82 | Open questions - how to calculate retrieval and code interpreter costs? 83 | """ 84 | 85 | retrival_costs = 0 86 | code_interpreter_costs = 0 87 | 88 | msgs = [ 89 | llm.safe_get(msg.model_dump(), "content.0.text.value") 90 | for msg in self.thread_messages 91 | ] 92 | joined_msgs = " ".join(msgs) 93 | 94 | msg_cost, tokens = llm.estimate_price_and_tokens(joined_msgs, self.model) 95 | 96 | with open(output_file, "w") as f: 97 | json.dump( 98 | { 99 | "cost": msg_cost, 100 | "tokens": tokens, 101 | }, 102 | f, 103 | indent=2, 104 | ) 105 | 106 | return self 107 | 108 | # ------------- CORE ASSISTANTS API FUNCTIONS ----------------- 109 | 110 | def get_or_create_assistant(self, name: str, model: str = "gpt-4-1106-preview"): 111 | print(f"get_or_create_assistant({name}, {model})") 112 | # Retrieve the list of existing assistants 113 | assistants: List[Assistant] = self.client.beta.assistants.list().data 114 | 115 | # Check if an assistant with the given name already exists 116 | for assistant in assistants: 117 | if assistant.name == name: 118 | self.assistant_id = assistant.id 119 | 120 | # update model if different 121 | if assistant.model != model: 122 | print(f"Updating assistant model from {assistant.model} to {model}") 123 | self.client.beta.assistants.update( 124 | assistant_id=self.assistant_id, model=model 125 | ) 126 | break 127 | else: # If no assistant was found with the name, create a new one 128 | assistant = self.client.beta.assistants.create(model=model, name=name) 129 | self.assistant_id = assistant.id 130 | 131 | self.model = model 132 | 133 | return self 134 | 135 | def set_instructions(self, instructions: str): 136 | print(f"set_instructions()") 137 | if self.assistant_id is None: 138 | raise ValueError( 139 | "No assistant has been created or retrieved. Call get_or_create_assistant() first." 140 | ) 141 | # Update the assistant with the new instructions 142 | updated_assistant = self.client.beta.assistants.update( 143 | assistant_id=self.assistant_id, instructions=instructions 144 | ) 145 | return self 146 | 147 | def equip_tools( 148 | self, turbo_tools: List[TurboTool], equip_on_assistant: bool = False 149 | ): 150 | print(f"equip_tools({turbo_tools}, {equip_on_assistant})") 151 | if self.assistant_id is None: 152 | raise ValueError( 153 | "No assistant has been created or retrieved. Call get_or_create_assistant() first." 154 | ) 155 | 156 | # Update the functions dictionary with the new tools 157 | self.map_function_tools = {tool.name: tool for tool in turbo_tools} 158 | 159 | if equip_on_assistant: 160 | # Update the assistant with the new list of tools, replacing any existing tools 161 | updated_assistant = self.client.beta.assistants.update( 162 | tools=self.tool_config, assistant_id=self.assistant_id 163 | ) 164 | 165 | return self 166 | 167 | def make_thread(self): 168 | print(f"make_thread()") 169 | 170 | if self.assistant_id is None: 171 | raise ValueError( 172 | "No assistant has been created. Call create_assistant() first." 173 | ) 174 | 175 | response = self.client.beta.threads.create() 176 | self.current_thread_id = response.id 177 | self.thread_messages = [] 178 | return self 179 | 180 | def add_message( 181 | self, 182 | message: str, 183 | file_ids: Optional[List[str]] = None, 184 | refresh_threads: bool = False, 185 | ): 186 | print(f"add_message(message={message}, file_ids={file_ids})") 187 | self.local_messages.append(message) 188 | self.client.beta.threads.messages.create( 189 | thread_id=self.current_thread_id, 190 | content=message, 191 | role="user", 192 | file_ids=file_ids or [], 193 | ) 194 | if refresh_threads: 195 | self.load_threads() 196 | return self 197 | 198 | def load_threads(self): 199 | self.thread_messages = self.client.beta.threads.messages.list( 200 | thread_id=self.current_thread_id 201 | ).data 202 | 203 | def list_steps(self): 204 | print(f"list_steps()") 205 | steps = self.client.beta.threads.runs.steps.list( 206 | thread_id=self.current_thread_id, 207 | run_id=self.run_id, 208 | ) 209 | print("steps", steps) 210 | return steps 211 | 212 | def run_thread(self, toolbox: Optional[List[str]] = None): 213 | print(f"run_thread(toolbox={toolbox})") 214 | if self.current_thread_id is None: 215 | raise ValueError("No thread has been created. Call make_thread() first.") 216 | if self.local_messages == []: 217 | raise ValueError("No messages have been added to the thread.") 218 | 219 | if toolbox is None: 220 | tools = None 221 | else: 222 | # get tools from toolbox 223 | tools = [self.map_function_tools[tool_name].config for tool_name in toolbox] 224 | 225 | # throw if tool not found 226 | if len(tools) != len(toolbox): 227 | raise ValueError( 228 | f"Tool not found in toolbox. toolbox={toolbox}, tools={tools}. Make sure all tools are equipped on the assistant." 229 | ) 230 | 231 | # refresh current thread 232 | self.load_threads() 233 | 234 | # Start the thread running 235 | run = self.client.beta.threads.runs.create( 236 | thread_id=self.current_thread_id, 237 | assistant_id=self.assistant_id, 238 | tools=tools, 239 | ) 240 | self.run_id = run.id 241 | 242 | # Polling mechanism to wait for thread's run completion or required actions 243 | while True: 244 | # self.list_steps() 245 | 246 | run_status = self.client.beta.threads.runs.retrieve( 247 | thread_id=self.current_thread_id, run_id=self.run_id 248 | ) 249 | if run_status.status == "requires_action": 250 | tool_outputs: List[ToolOutput] = [] 251 | for ( 252 | tool_call 253 | ) in run_status.required_action.submit_tool_outputs.tool_calls: 254 | tool_function = tool_call.function 255 | tool_name = tool_function.name 256 | 257 | # Check if tool_arguments is already a dictionary, if so, proceed directly 258 | if isinstance(tool_function.arguments, dict): 259 | tool_arguments = tool_function.arguments 260 | else: 261 | # Assume the arguments are JSON string and parse them 262 | tool_arguments = json.loads(tool_function.arguments) 263 | 264 | print(f"run_thread() Calling {tool_name}({tool_arguments})") 265 | 266 | # Assuming arguments are passed as a dictionary 267 | function_output = self.map_function_tools[tool_name].function( 268 | **tool_arguments 269 | ) 270 | 271 | tool_outputs.append( 272 | ToolOutput(tool_call_id=tool_call.id, output=function_output) 273 | ) 274 | 275 | # Submit the tool outputs back to the API 276 | self.client.beta.threads.runs.submit_tool_outputs( 277 | thread_id=self.current_thread_id, 278 | run_id=self.run_id, 279 | tool_outputs=[to for to in tool_outputs], 280 | ) 281 | elif run_status.status == "completed": 282 | self.load_threads() 283 | return self 284 | 285 | time.sleep(self.polling_interval) # Wait a little before polling again 286 | 287 | def enable_retrieval(self): 288 | print(f"enable_retrieval()") 289 | if self.assistant_id is None: 290 | raise ValueError( 291 | "No assistant has been created or retrieved. Call get_or_create_assistant() first." 292 | ) 293 | 294 | # Update the assistant with the new list of tools, replacing any existing tools 295 | updated_assistant = self.client.beta.assistants.update( 296 | tools=[{"type": "retrieval"}], assistant_id=self.assistant_id 297 | ) 298 | 299 | return self 300 | 301 | def upsert_files(self, file_paths: List[str]) -> List[str]: 302 | """ 303 | 304 | Upserts files to the file api - does not attach to assistant at all 305 | 306 | 1. Get all current files from the files api 307 | 2. Check if file exists 308 | 3. If file exists, look for byte size differences if changed: update it 309 | 4. If file does not exist, create it 310 | """ 311 | print(f"upsert_file({file_paths})") 312 | 313 | # confirm all files exist 314 | for file_path in file_paths: 315 | if not os.path.exists(file_path): 316 | raise ValueError(f"File does not exist: {file_path}") 317 | 318 | if self.assistant_id is None: 319 | raise ValueError( 320 | "No assistant has been created or retrieved. Call get_or_create_assistant() first." 321 | ) 322 | 323 | existing_files: List[FileObject] = self.client.files.list().data 324 | 325 | print("existing_files", existing_files) 326 | 327 | matched_files: List[FileObject] = [] 328 | 329 | # create file objects 330 | for i in file_paths: 331 | file_name = os.path.basename(i) 332 | file_object = open(i, "rb") 333 | 334 | print("file_object", file_object) 335 | 336 | local_file_size = os.path.getsize(i) 337 | 338 | file_found = False # Flag to check if file was found 339 | 340 | # check if file exists 341 | for existing_file in existing_files: 342 | if existing_file.filename == file_name: 343 | print(f"Found existing file {file_name}") 344 | file_found = True # Set flag to True since file is found 345 | 346 | # check if file has changed - delete and reupload if so 347 | if existing_file.bytes != local_file_size: 348 | print(f"File {file_name} has changed - updating") 349 | self.client.files.delete(file_id=existing_file) 350 | updated_file_object: FileObject = self.client.files.create( 351 | file=file_object, 352 | purpose="assistants", 353 | ) 354 | matched_files.append(updated_file_object) 355 | break # Exit the loop after handling the existing file 356 | 357 | # If the file was not found, create it 358 | if not file_found: 359 | print(f"Creating file {file_name}") 360 | new_file_object: FileObject = self.client.files.create( 361 | file=file_object, 362 | purpose="assistants", 363 | ) 364 | matched_files.append(new_file_object) 365 | else: 366 | print(f"File {file_name} already exists - no updates needed") 367 | matched_files.append(existing_file) 368 | 369 | return [f.id for f in matched_files] 370 | 371 | def get_files(self, file_ids: Optional[List[str]] = None): 372 | print(f"list_files()") 373 | files = self.client.files.list().data 374 | if file_ids is not None: 375 | print(f"filtering files by {file_ids}") 376 | files = [file for file in files if file.id in file_ids] 377 | print("files", files) 378 | return files 379 | 380 | def get_files_by_name(self, file_names: List[str]): 381 | print(f"get_files_by_name({file_names})") 382 | files: List[FileObject] = self.client.files.list().data 383 | 384 | output_files = [] 385 | 386 | for file_name in file_names: 387 | for file in files: 388 | print("file.filename", file.filename) 389 | if file.filename == file_name: 390 | output_files.append(file) 391 | break 392 | 393 | print("files", files) 394 | 395 | return output_files 396 | 397 | def get_file_ids_by_name(self, file_paths: List[str]): 398 | print(f"get_file_ids_by_name({file_paths})") 399 | file_names = [os.path.basename(file_path) for file_path in file_paths] 400 | print("file_names", file_names) 401 | files = self.get_files_by_name(file_names) 402 | print("files", files) 403 | file_ids = [file.id for file in files] 404 | print("file_ids", file_ids) 405 | return file_ids 406 | -------------------------------------------------------------------------------- /api-server/requirements.txt: -------------------------------------------------------------------------------- 1 | Flask==3.0.0 2 | openai 3 | psycopg2-binary 4 | python-dotenv -------------------------------------------------------------------------------- /api-server/vercel.json: -------------------------------------------------------------------------------- 1 | { 2 | "rewrites": [{ "source": "/(.*)", "destination": "/api/index" }] 3 | } 4 | -------------------------------------------------------------------------------- /fe-clients/react-ts/.eslintrc.cjs: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | root: true, 3 | env: { browser: true, es2020: true }, 4 | extends: [ 5 | 'eslint:recommended', 6 | 'plugin:@typescript-eslint/recommended', 7 | 'plugin:react-hooks/recommended', 8 | ], 9 | ignorePatterns: ['dist', '.eslintrc.cjs'], 10 | parser: '@typescript-eslint/parser', 11 | plugins: ['react-refresh'], 12 | rules: { 13 | 'react-refresh/only-export-components': [ 14 | 'warn', 15 | { allowConstantExport: true }, 16 | ], 17 | }, 18 | } 19 | -------------------------------------------------------------------------------- /fe-clients/react-ts/.gitignore: -------------------------------------------------------------------------------- 1 | .aider.ident.cache.v1 2 | .aider.tags.cache.v1 3 | .aider.chat.history.md 4 | .aider.input.history 5 | 6 | 7 | # Logs 8 | logs 9 | *.log 10 | npm-debug.log* 11 | yarn-debug.log* 12 | yarn-error.log* 13 | pnpm-debug.log* 14 | lerna-debug.log* 15 | 16 | node_modules 17 | dist 18 | dist-ssr 19 | *.local 20 | 21 | # Editor directories and files 22 | .vscode/* 23 | !.vscode/extensions.json 24 | .idea 25 | .DS_Store 26 | *.suo 27 | *.ntvs* 28 | *.njsproj 29 | *.sln 30 | *.sw? 31 | -------------------------------------------------------------------------------- /fe-clients/react-ts/README.md: -------------------------------------------------------------------------------- 1 | # React + TypeScript + Vite 2 | 3 | This template provides a minimal setup to get React working in Vite with HMR and some ESLint rules. 4 | 5 | ## Setup 6 | - Install Dependencies 7 | - `bun install` 8 | - Start the dev server 9 | - `bun run dev` 10 | 11 | Currently, two official plugins are available: 12 | 13 | - [@vitejs/plugin-react](https://github.com/vitejs/vite-plugin-react/blob/main/packages/plugin-react/README.md) uses [Babel](https://babeljs.io/) for Fast Refresh 14 | - [@vitejs/plugin-react-swc](https://github.com/vitejs/vite-plugin-react-swc) uses [SWC](https://swc.rs/) for Fast Refresh 15 | 16 | ## Expanding the ESLint configuration 17 | 18 | If you are developing a production application, we recommend updating the configuration to enable type aware lint rules: 19 | 20 | - Configure the top-level `parserOptions` property like this: 21 | 22 | ```js 23 | export default { 24 | // other rules... 25 | parserOptions: { 26 | ecmaVersion: 'latest', 27 | sourceType: 'module', 28 | project: ['./tsconfig.json', './tsconfig.node.json'], 29 | tsconfigRootDir: __dirname, 30 | }, 31 | } 32 | ``` 33 | 34 | - Replace `plugin:@typescript-eslint/recommended` to `plugin:@typescript-eslint/recommended-type-checked` or `plugin:@typescript-eslint/strict-type-checked` 35 | - Optionally add `plugin:@typescript-eslint/stylistic-type-checked` 36 | - Install [eslint-plugin-react](https://github.com/jsx-eslint/eslint-plugin-react) and add `plugin:react/recommended` & `plugin:react/jsx-runtime` to the `extends` list 37 | -------------------------------------------------------------------------------- /fe-clients/react-ts/bun.lockb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/fe-clients/react-ts/bun.lockb -------------------------------------------------------------------------------- /fe-clients/react-ts/index.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | Vite + React + TS 8 | 9 | 10 |
11 | 12 | 13 | 14 | -------------------------------------------------------------------------------- /fe-clients/react-ts/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "react-ts", 3 | "private": true, 4 | "version": "0.0.0", 5 | "type": "module", 6 | "scripts": { 7 | "dev": "vite", 8 | "build": "tsc && vite build", 9 | "lint": "eslint . --ext ts,tsx --report-unused-disable-directives --max-warnings 0", 10 | "preview": "vite preview" 11 | }, 12 | "dependencies": { 13 | "react": "^18.2.0", 14 | "react-dom": "^18.2.0" 15 | }, 16 | "devDependencies": { 17 | "@types/react": "^18.2.37", 18 | "@types/react-dom": "^18.2.15", 19 | "@typescript-eslint/eslint-plugin": "^6.10.0", 20 | "@typescript-eslint/parser": "^6.10.0", 21 | "@vitejs/plugin-react": "^4.2.0", 22 | "eslint": "^8.53.0", 23 | "eslint-plugin-react-hooks": "^4.6.0", 24 | "eslint-plugin-react-refresh": "^0.4.4", 25 | "typescript": "^5.2.2", 26 | "vite": "^5.0.0" 27 | } 28 | } 29 | -------------------------------------------------------------------------------- /fe-clients/react-ts/public/vite.svg: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /fe-clients/react-ts/src/App.css: -------------------------------------------------------------------------------- 1 | #root { 2 | max-width: 1280px; 3 | margin: 0 auto; 4 | padding: 2rem; 5 | text-align: center; 6 | } 7 | 8 | .logo { 9 | height: 6em; 10 | padding: 1.5em; 11 | will-change: filter; 12 | transition: filter 300ms; 13 | } 14 | .logo:hover { 15 | filter: drop-shadow(0 0 2em #646cffaa); 16 | } 17 | .logo.react:hover { 18 | filter: drop-shadow(0 0 2em #61dafbaa); 19 | } 20 | 21 | @keyframes logo-spin { 22 | from { 23 | transform: rotate(0deg); 24 | } 25 | to { 26 | transform: rotate(360deg); 27 | } 28 | } 29 | 30 | @media (prefers-reduced-motion: no-preference) { 31 | a:nth-of-type(2) .logo { 32 | animation: logo-spin infinite 20s linear; 33 | } 34 | } 35 | 36 | .card { 37 | padding: 2em; 38 | } 39 | 40 | .read-the-docs { 41 | color: #888; 42 | } 43 | /* Dark theme styles */ 44 | input, button, pre, section, code { 45 | background-color: #333; 46 | color: white; 47 | } 48 | 49 | /* Scrollable pre tag */ 50 | pre { 51 | overflow: auto; 52 | max-height: 200px; 53 | } 54 | 55 | /* Flex column with gap */ 56 | .flex-col-gap { 57 | display: flex; 58 | flex-direction: column; 59 | gap: 16px; 60 | } 61 | -------------------------------------------------------------------------------- /fe-clients/react-ts/src/App.tsx: -------------------------------------------------------------------------------- 1 | // code: add a enter key event that takes the 'prompt' and makes a http post request to localhost:3000/prompt using fetch. Take the json response and push it into the 'promptResults' list. Make sure to json.parse the response.results' it will be a json string. 2 | 3 | import { useState } from 'react' 4 | import reactLogo from './assets/react.svg' 5 | import viteLogo from '/vite.svg' 6 | import './App.css' 7 | 8 | // Define the PromptResult type 9 | type PromptResult = { 10 | prompt: string; 11 | results: Record[]; 12 | sql: string; 13 | }; 14 | 15 | function App() { 16 | 17 | // State variables 18 | const [prompt, setPrompt] = useState(''); 19 | // Load promptResults from local storage or default to empty list 20 | const [promptResults, setPromptResults] = useState( 21 | () => JSON.parse(localStorage.getItem('promptResults') || '[]') 22 | ); 23 | 24 | // Function to handle the Enter key press 25 | function handleKeyPress(event: React.KeyboardEvent) { 26 | if (event.key === 'Enter') { 27 | handleSubmit(); 28 | } 29 | } 30 | 31 | // Function to handle the submit action 32 | function handleSubmit() { 33 | fetch('http://localhost:3000/prompt', { 34 | method: 'POST', 35 | headers: { 36 | 'Content-Type': 'application/json', 37 | }, 38 | body: JSON.stringify({ prompt }), 39 | }) 40 | .then(response => response.json()) 41 | .then(data => { 42 | const resultsParsed = { 43 | ...data, 44 | results: JSON.parse(data.results), 45 | }; 46 | const newPromptResults = [...promptResults, resultsParsed]; 47 | setPromptResults(newPromptResults); 48 | localStorage.setItem('promptResults', JSON.stringify(newPromptResults)); 49 | }) 50 | .catch(error => console.error('Error:', error)); 51 | } 52 | 53 | return ( 54 | <> 55 | 63 |

Framework: React-TS

64 |

TTYDB Prototype

65 | {/* Dark theme input and button */} 66 | setPrompt(e.target.value)} 70 | onKeyPress={handleKeyPress} 71 | className="dark-theme" 72 | /> 73 | 74 | 75 | 76 | {/* List of prompt results */} 77 |
78 | {promptResults.map((result, index) => { 79 | return ( 80 |
81 | {/* Display the prompt used */} 82 |

Prompt: {result.prompt}

83 | 
{JSON.stringify(result.results, null, 2)}
84 | {result.sql} 85 |
86 | ); 87 | })} 88 |
89 | 90 | ) 91 | } 92 | 93 | export default App 94 | -------------------------------------------------------------------------------- /fe-clients/react-ts/src/assets/react.svg: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /fe-clients/react-ts/src/index.css: -------------------------------------------------------------------------------- 1 | :root { 2 | font-family: Inter, system-ui, Avenir, Helvetica, Arial, sans-serif; 3 | line-height: 1.5; 4 | font-weight: 400; 5 | 6 | color-scheme: light dark; 7 | color: rgba(255, 255, 255, 0.87); 8 | background-color: #242424; 9 | 10 | font-synthesis: none; 11 | text-rendering: optimizeLegibility; 12 | -webkit-font-smoothing: antialiased; 13 | -moz-osx-font-smoothing: grayscale; 14 | } 15 | 16 | a { 17 | font-weight: 500; 18 | color: #646cff; 19 | text-decoration: inherit; 20 | } 21 | a:hover { 22 | color: #535bf2; 23 | } 24 | 25 | body { 26 | margin: 0; 27 | display: flex; 28 | place-items: center; 29 | min-width: 320px; 30 | min-height: 100vh; 31 | } 32 | 33 | h1 { 34 | font-size: 3.2em; 35 | line-height: 1.1; 36 | } 37 | 38 | button { 39 | border-radius: 8px; 40 | border: 1px solid transparent; 41 | padding: 0.6em 1.2em; 42 | font-size: 1em; 43 | font-weight: 500; 44 | font-family: inherit; 45 | background-color: #1a1a1a; 46 | cursor: pointer; 47 | transition: border-color 0.25s; 48 | } 49 | button:hover { 50 | border-color: #646cff; 51 | } 52 | button:focus, 53 | button:focus-visible { 54 | outline: 4px auto -webkit-focus-ring-color; 55 | } 56 | 57 | @media (prefers-color-scheme: light) { 58 | :root { 59 | color: #213547; 60 | background-color: #ffffff; 61 | } 62 | a:hover { 63 | color: #747bff; 64 | } 65 | button { 66 | background-color: #f9f9f9; 67 | } 68 | } 69 | -------------------------------------------------------------------------------- /fe-clients/react-ts/src/main.tsx: -------------------------------------------------------------------------------- 1 | import React from 'react' 2 | import ReactDOM from 'react-dom/client' 3 | import App from './App.tsx' 4 | import './index.css' 5 | 6 | ReactDOM.createRoot(document.getElementById('root')!).render( 7 | 8 | 9 | , 10 | ) 11 | -------------------------------------------------------------------------------- /fe-clients/react-ts/src/vite-env.d.ts: -------------------------------------------------------------------------------- 1 | /// 2 | -------------------------------------------------------------------------------- /fe-clients/react-ts/tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "ES2020", 4 | "useDefineForClassFields": true, 5 | "lib": ["ES2020", "DOM", "DOM.Iterable"], 6 | "module": "ESNext", 7 | "skipLibCheck": true, 8 | 9 | /* Bundler mode */ 10 | "moduleResolution": "bundler", 11 | "allowImportingTsExtensions": true, 12 | "resolveJsonModule": true, 13 | "isolatedModules": true, 14 | "noEmit": true, 15 | "jsx": "react-jsx", 16 | 17 | /* Linting */ 18 | "strict": true, 19 | "noUnusedLocals": true, 20 | "noUnusedParameters": true, 21 | "noFallthroughCasesInSwitch": true 22 | }, 23 | "include": ["src"], 24 | "references": [{ "path": "./tsconfig.node.json" }] 25 | } 26 | -------------------------------------------------------------------------------- /fe-clients/react-ts/tsconfig.node.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "composite": true, 4 | "skipLibCheck": true, 5 | "module": "ESNext", 6 | "moduleResolution": "bundler", 7 | "allowSyntheticDefaultImports": true 8 | }, 9 | "include": ["vite.config.ts"] 10 | } 11 | -------------------------------------------------------------------------------- /fe-clients/react-ts/vite.config.ts: -------------------------------------------------------------------------------- 1 | import { defineConfig } from 'vite' 2 | import react from '@vitejs/plugin-react' 3 | 4 | // https://vitejs.dev/config/ 5 | export default defineConfig({ 6 | plugins: [react()], 7 | }) 8 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/.gitignore: -------------------------------------------------------------------------------- 1 | .aider.ident.cache.v1 2 | .aider.tags.cache.v1 3 | .aider.chat.history.md 4 | .aider.input.history 5 | 6 | 7 | # Logs 8 | logs 9 | *.log 10 | npm-debug.log* 11 | yarn-debug.log* 12 | yarn-error.log* 13 | pnpm-debug.log* 14 | lerna-debug.log* 15 | 16 | node_modules 17 | dist 18 | dist-ssr 19 | *.local 20 | 21 | # Editor directories and files 22 | .vscode/* 23 | !.vscode/extensions.json 24 | .idea 25 | .DS_Store 26 | *.suo 27 | *.ntvs* 28 | *.njsproj 29 | *.sln 30 | *.sw? 31 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/.vscode/extensions.json: -------------------------------------------------------------------------------- 1 | { 2 | "recommendations": ["svelte.svelte-vscode"] 3 | } 4 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/README.md: -------------------------------------------------------------------------------- 1 | # Svelte + TS + Vite 2 | 3 | This template should help get you started developing with Svelte and TypeScript in Vite. 4 | 5 | ## Setup 6 | - Install Dependencies 7 | - `bun install` 8 | - Start the dev server 9 | - `bun run dev` 10 | 11 | ## Recommended IDE Setup 12 | 13 | [VS Code](https://code.visualstudio.com/) + [Svelte](https://marketplace.visualstudio.com/items?itemName=svelte.svelte-vscode). 14 | 15 | ## Need an official Svelte framework? 16 | 17 | Check out [SvelteKit](https://github.com/sveltejs/kit#readme), which is also powered by Vite. Deploy anywhere with its serverless-first approach and adapt to various platforms, with out of the box support for TypeScript, SCSS, and Less, and easily-added support for mdsvex, GraphQL, PostCSS, Tailwind CSS, and more. 18 | 19 | ## Technical considerations 20 | 21 | **Why use this over SvelteKit?** 22 | 23 | - It brings its own routing solution which might not be preferable for some users. 24 | - It is first and foremost a framework that just happens to use Vite under the hood, not a Vite app. 25 | 26 | This template contains as little as possible to get started with Vite + TypeScript + Svelte, while taking into account the developer experience with regards to HMR and intellisense. It demonstrates capabilities on par with the other `create-vite` templates and is a good starting point for beginners dipping their toes into a Vite + Svelte project. 27 | 28 | Should you later need the extended capabilities and extensibility provided by SvelteKit, the template has been structured similarly to SvelteKit so that it is easy to migrate. 29 | 30 | **Why `global.d.ts` instead of `compilerOptions.types` inside `jsconfig.json` or `tsconfig.json`?** 31 | 32 | Setting `compilerOptions.types` shuts out all other types not explicitly listed in the configuration. Using triple-slash references keeps the default TypeScript setting of accepting type information from the entire workspace, while also adding `svelte` and `vite/client` type information. 33 | 34 | **Why include `.vscode/extensions.json`?** 35 | 36 | Other templates indirectly recommend extensions via the README, but this file allows VS Code to prompt the user to install the recommended extension upon opening the project. 37 | 38 | **Why enable `allowJs` in the TS template?** 39 | 40 | While `allowJs: false` would indeed prevent the use of `.js` files in the project, it does not prevent the use of JavaScript syntax in `.svelte` files. In addition, it would force `checkJs: false`, bringing the worst of both worlds: not being able to guarantee the entire codebase is TypeScript, and also having worse typechecking for the existing JavaScript. In addition, there are valid use cases in which a mixed codebase may be relevant. 41 | 42 | **Why is HMR not preserving my local component state?** 43 | 44 | HMR state preservation comes with a number of gotchas! It has been disabled by default in both `svelte-hmr` and `@sveltejs/vite-plugin-svelte` due to its often surprising behavior. You can read the details [here](https://github.com/rixo/svelte-hmr#svelte-hmr). 45 | 46 | If you have state that's important to retain within a component, consider creating an external store which would not be replaced by HMR. 47 | 48 | ```ts 49 | // store.ts 50 | // An extremely simple external store 51 | import { writable } from 'svelte/store' 52 | export default writable(0) 53 | ``` 54 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/bun.lockb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/fe-clients/svelte-ts/bun.lockb -------------------------------------------------------------------------------- /fe-clients/svelte-ts/index.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | Vite + Svelte + TS 8 | 9 | 10 |
11 | 12 | 13 | 14 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "svelte-ts", 3 | "private": true, 4 | "version": "0.0.0", 5 | "type": "module", 6 | "scripts": { 7 | "dev": "vite", 8 | "build": "vite build", 9 | "preview": "vite preview", 10 | "check": "svelte-check --tsconfig ./tsconfig.json" 11 | }, 12 | "devDependencies": { 13 | "@sveltejs/vite-plugin-svelte": "^3.0.0", 14 | "@tsconfig/svelte": "^5.0.2", 15 | "svelte": "^4.2.3", 16 | "svelte-check": "^3.6.0", 17 | "tslib": "^2.6.2", 18 | "typescript": "^5.2.2", 19 | "vite": "^5.0.0" 20 | } 21 | } 22 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/public/vite.svg: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/src/App.svelte: -------------------------------------------------------------------------------- 1 | 2 | 3 | 45 | 46 |
47 | 55 |

Framework: Svelete-TS

56 |

TTYDB Prototype

57 | 63 | 64 | {#each promptResults as result (result.prompt)} 65 |
66 |

{result.prompt}

67 | 
{JSON.stringify(result.results, null, 2)}
68 | {result.sql} 69 |
70 | {/each} 71 |
72 | 73 | 103 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/src/app.css: -------------------------------------------------------------------------------- 1 | :root { 2 | font-family: Inter, system-ui, Avenir, Helvetica, Arial, sans-serif; 3 | line-height: 1.5; 4 | font-weight: 400; 5 | 6 | color-scheme: light dark; 7 | color: rgba(255, 255, 255, 0.87); 8 | background-color: #242424; 9 | 10 | font-synthesis: none; 11 | text-rendering: optimizeLegibility; 12 | -webkit-font-smoothing: antialiased; 13 | -moz-osx-font-smoothing: grayscale; 14 | } 15 | 16 | a { 17 | font-weight: 500; 18 | color: #646cff; 19 | text-decoration: inherit; 20 | } 21 | a:hover { 22 | color: #535bf2; 23 | } 24 | 25 | body { 26 | margin: 0; 27 | display: flex; 28 | place-items: center; 29 | min-width: 320px; 30 | min-height: 100vh; 31 | } 32 | 33 | h1 { 34 | font-size: 3.2em; 35 | line-height: 1.1; 36 | } 37 | 38 | .card { 39 | padding: 2em; 40 | } 41 | 42 | #app { 43 | max-width: 1280px; 44 | margin: 0 auto; 45 | padding: 2rem; 46 | text-align: center; 47 | } 48 | 49 | pre { 50 | text-align: left; 51 | } 52 | 53 | button { 54 | border-radius: 8px; 55 | border: 1px solid transparent; 56 | padding: 0.6em 1.2em; 57 | font-size: 1em; 58 | font-weight: 500; 59 | font-family: inherit; 60 | background-color: #1a1a1a; 61 | cursor: pointer; 62 | transition: border-color 0.25s; 63 | } 64 | button:hover { 65 | border-color: #646cff; 66 | } 67 | button:focus, 68 | button:focus-visible { 69 | outline: 4px auto -webkit-focus-ring-color; 70 | } 71 | 72 | @media (prefers-color-scheme: light) { 73 | :root { 74 | color: #213547; 75 | background-color: #ffffff; 76 | } 77 | a:hover { 78 | color: #747bff; 79 | } 80 | button { 81 | background-color: #f9f9f9; 82 | } 83 | } 84 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/src/assets/svelte.svg: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/src/lib/Counter.svelte: -------------------------------------------------------------------------------- 1 | 7 | 8 | 11 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/src/main.ts: -------------------------------------------------------------------------------- 1 | import './app.css' 2 | import App from './App.svelte' 3 | 4 | const app = new App({ 5 | target: document.getElementById('app'), 6 | }) 7 | 8 | export default app 9 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/src/vite-env.d.ts: -------------------------------------------------------------------------------- 1 | /// 2 | /// 3 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/svelte.config.js: -------------------------------------------------------------------------------- 1 | import { vitePreprocess } from '@sveltejs/vite-plugin-svelte' 2 | 3 | export default { 4 | // Consult https://svelte.dev/docs#compile-time-svelte-preprocess 5 | // for more information about preprocessors 6 | preprocess: vitePreprocess(), 7 | } 8 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "extends": "@tsconfig/svelte/tsconfig.json", 3 | "compilerOptions": { 4 | "target": "ESNext", 5 | "useDefineForClassFields": true, 6 | "module": "ESNext", 7 | "resolveJsonModule": true, 8 | /** 9 | * Typecheck JS in `.svelte` and `.js` files by default. 10 | * Disable checkJs if you'd like to use dynamic types in JS. 11 | * Note that setting allowJs false does not prevent the use 12 | * of JS in `.svelte` files. 13 | */ 14 | "allowJs": true, 15 | "checkJs": true, 16 | "isolatedModules": true 17 | }, 18 | "include": ["src/**/*.ts", "src/**/*.js", "src/**/*.svelte"], 19 | "references": [{ "path": "./tsconfig.node.json" }] 20 | } 21 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/tsconfig.node.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "composite": true, 4 | "skipLibCheck": true, 5 | "module": "ESNext", 6 | "moduleResolution": "bundler" 7 | }, 8 | "include": ["vite.config.ts"] 9 | } 10 | -------------------------------------------------------------------------------- /fe-clients/svelte-ts/vite.config.ts: -------------------------------------------------------------------------------- 1 | import { defineConfig } from 'vite' 2 | import { svelte } from '@sveltejs/vite-plugin-svelte' 3 | 4 | // https://vitejs.dev/config/ 5 | export default defineConfig({ 6 | plugins: [svelte()], 7 | }) 8 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/.gitignore: -------------------------------------------------------------------------------- 1 | .aider.ident.cache.v1 2 | .aider.tags.cache.v1 3 | .aider.chat.history.md 4 | .aider.input.history 5 | 6 | 7 | # Logs 8 | logs 9 | *.log 10 | npm-debug.log* 11 | yarn-debug.log* 12 | yarn-error.log* 13 | pnpm-debug.log* 14 | lerna-debug.log* 15 | 16 | node_modules 17 | dist 18 | dist-ssr 19 | *.local 20 | 21 | # Editor directories and files 22 | .vscode/* 23 | !.vscode/extensions.json 24 | .idea 25 | .DS_Store 26 | *.suo 27 | *.ntvs* 28 | *.njsproj 29 | *.sln 30 | *.sw? 31 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/.vscode/extensions.json: -------------------------------------------------------------------------------- 1 | { 2 | "recommendations": ["Vue.volar", "Vue.vscode-typescript-vue-plugin"] 3 | } 4 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/README.md: -------------------------------------------------------------------------------- 1 | # Vue 3 + TypeScript + Vite 2 | 3 | This template should help get you started developing with Vue 3 and TypeScript in Vite. The template uses Vue 3 ` 12 | 13 | 14 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "vue-ts", 3 | "private": true, 4 | "version": "0.0.0", 5 | "type": "module", 6 | "scripts": { 7 | "dev": "vite", 8 | "build": "vue-tsc && vite build", 9 | "preview": "vite preview" 10 | }, 11 | "dependencies": { 12 | "vue": "^3.3.8" 13 | }, 14 | "devDependencies": { 15 | "@vitejs/plugin-vue": "^4.5.0", 16 | "typescript": "^5.2.2", 17 | "vite": "^5.0.0", 18 | "vue-tsc": "^1.8.22" 19 | } 20 | } 21 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/public/vite.svg: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/src/App.vue: -------------------------------------------------------------------------------- 1 | 30 | 31 | 82 | 83 | 140 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/src/assets/ttydb.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/src/assets/vue.svg: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/src/components/HelloWorld.vue: -------------------------------------------------------------------------------- 1 | 8 | 9 | 33 | 34 | 39 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/src/main.ts: -------------------------------------------------------------------------------- 1 | import { createApp } from 'vue' 2 | import './style.css' 3 | import App from './App.vue' 4 | 5 | createApp(App).mount('#app') 6 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/src/style.css: -------------------------------------------------------------------------------- 1 | :root { 2 | font-family: Inter, system-ui, Avenir, Helvetica, Arial, sans-serif; 3 | line-height: 1.5; 4 | font-weight: 400; 5 | 6 | color-scheme: light dark; 7 | color: rgba(255, 255, 255, 0.87); 8 | background-color: #242424; 9 | 10 | font-synthesis: none; 11 | text-rendering: optimizeLegibility; 12 | -webkit-font-smoothing: antialiased; 13 | -moz-osx-font-smoothing: grayscale; 14 | } 15 | 16 | a { 17 | font-weight: 500; 18 | color: #646cff; 19 | text-decoration: inherit; 20 | } 21 | a:hover { 22 | color: #535bf2; 23 | } 24 | 25 | body { 26 | margin: 0; 27 | display: flex; 28 | place-items: center; 29 | min-width: 320px; 30 | min-height: 100vh; 31 | } 32 | 33 | h1 { 34 | font-size: 3.2em; 35 | line-height: 1.1; 36 | } 37 | 38 | button { 39 | border-radius: 8px; 40 | border: 1px solid transparent; 41 | padding: 0.6em 1.2em; 42 | font-size: 1em; 43 | font-weight: 500; 44 | font-family: inherit; 45 | background-color: #1a1a1a; 46 | cursor: pointer; 47 | transition: border-color 0.25s; 48 | } 49 | button:hover { 50 | border-color: #646cff; 51 | } 52 | button:focus, 53 | button:focus-visible { 54 | outline: 4px auto -webkit-focus-ring-color; 55 | } 56 | 57 | .card { 58 | padding: 2em; 59 | } 60 | 61 | #app { 62 | max-width: 1280px; 63 | margin: 0 auto; 64 | padding: 2rem; 65 | text-align: center; 66 | } 67 | 68 | @media (prefers-color-scheme: light) { 69 | :root { 70 | color: #213547; 71 | background-color: #ffffff; 72 | } 73 | a:hover { 74 | color: #747bff; 75 | } 76 | button { 77 | background-color: #f9f9f9; 78 | } 79 | } 80 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/src/vite-env.d.ts: -------------------------------------------------------------------------------- 1 | /// 2 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "ES2020", 4 | "useDefineForClassFields": true, 5 | "module": "ESNext", 6 | "lib": ["ES2020", "DOM", "DOM.Iterable"], 7 | "skipLibCheck": true, 8 | 9 | /* Bundler mode */ 10 | "moduleResolution": "bundler", 11 | "allowImportingTsExtensions": true, 12 | "resolveJsonModule": true, 13 | "isolatedModules": true, 14 | "noEmit": true, 15 | "jsx": "preserve", 16 | 17 | /* Linting */ 18 | "strict": true, 19 | "noUnusedLocals": true, 20 | "noUnusedParameters": true, 21 | "noFallthroughCasesInSwitch": true 22 | }, 23 | "include": ["src/**/*.ts", "src/**/*.tsx", "src/**/*.vue"], 24 | "references": [{ "path": "./tsconfig.node.json" }] 25 | } 26 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/tsconfig.node.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "composite": true, 4 | "skipLibCheck": true, 5 | "module": "ESNext", 6 | "moduleResolution": "bundler", 7 | "allowSyntheticDefaultImports": true 8 | }, 9 | "include": ["vite.config.ts"] 10 | } 11 | -------------------------------------------------------------------------------- /fe-clients/vue-ts/vite.config.ts: -------------------------------------------------------------------------------- 1 | import { defineConfig } from 'vite' 2 | import vue from '@vitejs/plugin-vue' 3 | 4 | // https://vitejs.dev/config/ 5 | export default defineConfig({ 6 | plugins: [vue()], 7 | }) 8 | -------------------------------------------------------------------------------- /imgs/1-prompt-engineering-postgres-ai-data-analytics-agent.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/1-prompt-engineering-postgres-ai-data-analytics-agent.png -------------------------------------------------------------------------------- /imgs/10-talk-to-your-database.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/10-talk-to-your-database.png -------------------------------------------------------------------------------- /imgs/2-using-autogen-to-build-our-multi-agent-postgres-data-analytics-tool.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/2-using-autogen-to-build-our-multi-agent-postgres-data-analytics-tool.png -------------------------------------------------------------------------------- /imgs/2024-predictions.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/2024-predictions.png -------------------------------------------------------------------------------- /imgs/3-make-autogen-consistent-to-build-our-multi-agent-postgres-data-analytics-tool.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/3-make-autogen-consistent-to-build-our-multi-agent-postgres-data-analytics-tool.png -------------------------------------------------------------------------------- /imgs/4-autogen-token-tactics-managing-llm-memory-and-costs-multi-agent-postgres-ai-data-analytics.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/4-autogen-token-tactics-managing-llm-memory-and-costs-multi-agent-postgres-ai-data-analytics.png -------------------------------------------------------------------------------- /imgs/5-autogen-spyware-for-ai-agents-postgres-data-analytics-tool-ai.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/5-autogen-spyware-for-ai-agents-postgres-data-analytics-tool-ai.png -------------------------------------------------------------------------------- /imgs/6-autogen-and-guidance-for-autonomous-control-flow.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/6-autogen-and-guidance-for-autonomous-control-flow.png -------------------------------------------------------------------------------- /imgs/7-turbo4-assistants-threads-messages.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/7-turbo4-assistants-threads-messages.png -------------------------------------------------------------------------------- /imgs/8-ccc-ai-engineering-with-aider.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/8-ccc-ai-engineering-with-aider.png -------------------------------------------------------------------------------- /imgs/9-self-correcting-openai-gpt4-turbo-assistant.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/9-self-correcting-openai-gpt4-turbo-assistant.png -------------------------------------------------------------------------------- /imgs/multi-agent-coding.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/disler/multi-agent-postgres-data-analytics/b1bfdbb4d798eb0cd8b39066d04717874ee68158/imgs/multi-agent-coding.png -------------------------------------------------------------------------------- /postgres_da_ai_agent/agents/agent_config.py: -------------------------------------------------------------------------------- 1 | import autogen 2 | 3 | 4 | # build the gpt_configuration object 5 | # Base Configuration 6 | base_config = { 7 | "use_cache": False, 8 | "temperature": 0, 9 | "config_list": autogen.config_list_from_models(["gpt-4"]), 10 | "request_timeout": 120, 11 | } 12 | 13 | # Configuration with "run_sql" 14 | run_sql_config = { 15 | **base_config, # Inherit base configuration 16 | "functions": [ 17 | { 18 | "name": "run_sql", 19 | "description": "Run a SQL query against the postgres database", 20 | "parameters": { 21 | "type": "object", 22 | "properties": { 23 | "sql": { 24 | "type": "string", 25 | "description": "The SQL query to run", 26 | } 27 | }, 28 | "required": ["sql"], 29 | }, 30 | } 31 | ], 32 | } 33 | 34 | # Configuration with "write_file" 35 | write_file_config = { 36 | **base_config, # Inherit base configuration 37 | "functions": [ 38 | { 39 | "name": "write_file", 40 | "description": "Write a file to the filesystem", 41 | "parameters": { 42 | "type": "object", 43 | "properties": { 44 | "fname": { 45 | "type": "string", 46 | "description": "The name of the file to write", 47 | }, 48 | "content": { 49 | "type": "string", 50 | "description": "The content of the file to write", 51 | }, 52 | }, 53 | "required": ["fname", "content"], 54 | }, 55 | } 56 | ], 57 | } 58 | 59 | # Configuration with "write_json_file" 60 | write_json_file_config = { 61 | **base_config, # Inherit base configuration 62 | "functions": [ 63 | { 64 | "name": "write_json_file", 65 | "description": "Write a json file to the filesystem", 66 | "parameters": { 67 | "type": "object", 68 | "properties": { 69 | "fname": { 70 | "type": "string", 71 | "description": "The name of the file to write", 72 | }, 73 | "json_str": { 74 | "type": "string", 75 | "description": "The content of the file to write", 76 | }, 77 | }, 78 | "required": ["fname", "json_str"], 79 | }, 80 | } 81 | ], 82 | } 83 | 84 | write_yaml_file_config = { 85 | **base_config, # Inherit base configuration 86 | "functions": [ 87 | { 88 | "name": "write_yml_file", 89 | "description": "Write a yml file to the filesystem", 90 | "parameters": { 91 | "type": "object", 92 | "properties": { 93 | "fname": { 94 | "type": "string", 95 | "description": "The name of the file to write", 96 | }, 97 | "json_str": { 98 | "type": "string", 99 | "description": "The json content of the file to write", 100 | }, 101 | }, 102 | "required": ["fname", "json_str"], 103 | }, 104 | } 105 | ], 106 | } 107 | 108 | 109 | write_innovation_file_config = { 110 | **base_config, # Inherit base configuration 111 | "functions": [ 112 | { 113 | "name": "write_innovation_file", 114 | "description": "Write a file to the filesystem", 115 | "parameters": { 116 | "type": "object", 117 | "properties": { 118 | "content": { 119 | "type": "string", 120 | "description": "The content of the file to write", 121 | }, 122 | }, 123 | "required": ["content"], 124 | }, 125 | } 126 | ], 127 | } 128 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/agents/agents.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List, Dict, Any 2 | from postgres_da_ai_agent.agents.instruments import PostgresAgentInstruments 3 | from postgres_da_ai_agent.modules import orchestrator 4 | from postgres_da_ai_agent.agents import agent_config 5 | import autogen 6 | import guidance 7 | 8 | # ------------------------ PROMPTS ------------------------ 9 | 10 | 11 | USER_PROXY_PROMPT = "A human admin. Interact with the Product Manager to discuss the plan. Plan execution needs to be approved by this admin." 12 | DATA_ENGINEER_PROMPT = "A Data Engineer. Generate the initial SQL based on the requirements provided. Send it to the Sr Data Analyst to be executed. " 13 | SR_DATA_ANALYST_PROMPT = "Sr Data Analyst. You run the SQL query using the run_sql function, send the raw response to the data viz team. You use the run_sql function exclusively." 14 | 15 | 16 | GUIDANCE_SCRUM_MASTER_SQL_NLQ_PROMPT = """ 17 | Is the following block of text a SQL Natural Language Query (NLQ)? Please rank from 1 to 5, where: 18 | 1: Definitely not NLQ 19 | 2: Likely not NLQ 20 | 3: Neutral / Unsure 21 | 4: Likely NLQ 22 | 5: Definitely NLQ 23 | 24 | Return the rank as a number exclusively using the rank variable to be casted as an integer. 25 | 26 | Block of Text: {{potential_nlq}} 27 | {{#select "rank" logprobs='logprobs'}} 1{{or}} 2{{or}} 3{{or}} 4{{or}} 5{{/select}} 28 | """ 29 | 30 | DATA_INSIGHTS_GUIDANCE_PROMPT = """ 31 | You're a data innovator. You analyze SQL databases table structure and generate 3 novel insights for your team to reflect on and query. 32 | Format your insights in JSON format. 33 | ```json 34 | [{{#geneach 'insight' num_iterations=3 join=','}} 35 | { 36 | "insight": "{{gen 'insight' temperature=0.7}}", 37 | "actionable_business_value": "{{gen 'actionable_value' temperature=0.7}}", 38 | "sql": "{{gen 'new_query' temperature=0.7}}" 39 | } 40 | {{/geneach}}] 41 | ```""" 42 | 43 | 44 | INSIGHTS_FILE_REPORTER_PROMPT = "You're a data reporter. You write json data you receive directly into a file using the write_innovation_file function." 45 | 46 | 47 | # unused prompts 48 | COMPLETION_PROMPT = "If everything looks good, respond with APPROVED" 49 | PRODUCT_MANAGER_PROMPT = ( 50 | "Product Manager. Validate the response to make sure it's correct" 51 | + COMPLETION_PROMPT 52 | ) 53 | TEXT_REPORT_ANALYST_PROMPT = "Text File Report Analyst. You exclusively use the write_file function on a summarized report." 54 | JSON_REPORT_ANALYST_PROMPT = "Json Report Analyst. You exclusively use the write_json_file function on the report." 55 | YML_REPORT_ANALYST_PROMPT = "Yaml Report Analyst. You exclusively use the write_yml_file function on the report." 56 | 57 | 58 | # ------------------------ BUILD AGENT TEAMS ------------------------ 59 | 60 | 61 | def build_data_eng_team(instruments: PostgresAgentInstruments): 62 | """ 63 | Build a team of agents that can generate, execute, and report an SQL query 64 | """ 65 | 66 | # create a set of agents with specific roles 67 | # admin user proxy agent - takes in the prompt and manages the group chat 68 | user_proxy = autogen.UserProxyAgent( 69 | name="Admin", 70 | system_message=USER_PROXY_PROMPT, 71 | code_execution_config=False, 72 | human_input_mode="NEVER", 73 | ) 74 | 75 | # data engineer agent - generates the sql query 76 | data_engineer = autogen.AssistantAgent( 77 | name="Engineer", 78 | llm_config=agent_config.base_config, 79 | system_message=DATA_ENGINEER_PROMPT, 80 | code_execution_config=False, 81 | human_input_mode="NEVER", 82 | ) 83 | 84 | sr_data_analyst = autogen.AssistantAgent( 85 | name="Sr_Data_Analyst", 86 | llm_config=agent_config.run_sql_config, 87 | system_message=SR_DATA_ANALYST_PROMPT, 88 | code_execution_config=False, 89 | human_input_mode="NEVER", 90 | function_map={ 91 | "run_sql": instruments.run_sql, 92 | }, 93 | ) 94 | 95 | return [ 96 | user_proxy, 97 | data_engineer, 98 | sr_data_analyst, 99 | ] 100 | 101 | 102 | def build_data_viz_team(instruments: PostgresAgentInstruments): 103 | # admin user proxy agent - takes in the prompt and manages the group chat 104 | user_proxy = autogen.UserProxyAgent( 105 | name="Admin", 106 | system_message=USER_PROXY_PROMPT, 107 | code_execution_config=False, 108 | human_input_mode="NEVER", 109 | ) 110 | 111 | # text report analyst - writes a summary report of the results and saves them to a local text file 112 | text_report_analyst = autogen.AssistantAgent( 113 | name="Text_Report_Analyst", 114 | llm_config=agent_config.write_file_config, 115 | system_message=TEXT_REPORT_ANALYST_PROMPT, 116 | human_input_mode="NEVER", 117 | function_map={ 118 | "write_file": instruments.write_file, 119 | }, 120 | ) 121 | 122 | # json report analyst - writes a summary report of the results and saves them to a local json file 123 | json_report_analyst = autogen.AssistantAgent( 124 | name="Json_Report_Analyst", 125 | llm_config=agent_config.write_json_file_config, 126 | system_message=JSON_REPORT_ANALYST_PROMPT, 127 | human_input_mode="NEVER", 128 | function_map={ 129 | "write_json_file": instruments.write_json_file, 130 | }, 131 | ) 132 | 133 | yaml_report_analyst = autogen.AssistantAgent( 134 | name="Yml_Report_Analyst", 135 | llm_config=agent_config.write_yaml_file_config, 136 | system_message=YML_REPORT_ANALYST_PROMPT, 137 | human_input_mode="NEVER", 138 | function_map={ 139 | "write_yml_file": instruments.write_yml_file, 140 | }, 141 | ) 142 | 143 | return [ 144 | user_proxy, 145 | text_report_analyst, 146 | json_report_analyst, 147 | yaml_report_analyst, 148 | ] 149 | 150 | 151 | def build_scrum_master_team(instruments: PostgresAgentInstruments): 152 | user_proxy = autogen.UserProxyAgent( 153 | name="Admin", 154 | system_message=USER_PROXY_PROMPT, 155 | code_execution_config=False, 156 | human_input_mode="NEVER", 157 | ) 158 | 159 | scrum_agent = DefensiveScrumMasterAgent( 160 | name="Scrum_Master", 161 | llm_config=agent_config.base_config, 162 | system_message=GUIDANCE_SCRUM_MASTER_SQL_NLQ_PROMPT, 163 | human_input_mode="NEVER", 164 | ) 165 | 166 | return [user_proxy, scrum_agent] 167 | 168 | 169 | def build_insights_team(instruments: PostgresAgentInstruments): 170 | user_proxy = autogen.UserProxyAgent( 171 | name="Admin", 172 | system_message=USER_PROXY_PROMPT, 173 | code_execution_config=False, 174 | human_input_mode="NEVER", 175 | ) 176 | 177 | insights_agent = InsightsAgent( 178 | name="Insights", 179 | llm_config=agent_config.base_config, 180 | system_message=DATA_INSIGHTS_GUIDANCE_PROMPT, 181 | human_input_mode="NEVER", 182 | ) 183 | 184 | insights_data_reporter = autogen.AssistantAgent( 185 | name="Insights_Data_Reporter", 186 | llm_config=agent_config.write_innovation_file_config, 187 | system_message=INSIGHTS_FILE_REPORTER_PROMPT, 188 | human_input_mode="NEVER", 189 | function_map={ 190 | "write_innovation_file": instruments.write_innovation_file, 191 | }, 192 | ) 193 | 194 | return [user_proxy, insights_agent, insights_data_reporter] 195 | 196 | 197 | # ------------------------ ORCHESTRATION ------------------------ 198 | 199 | 200 | def build_team_orchestrator( 201 | team: str, 202 | agent_instruments: PostgresAgentInstruments, 203 | validate_results: callable = None, 204 | ) -> orchestrator.Orchestrator: 205 | """ 206 | Based on a team name, build a team of agents and return an orchestrator 207 | """ 208 | if team == "data_eng": 209 | return orchestrator.Orchestrator( 210 | name="data_eng_team", 211 | agents=build_data_eng_team(agent_instruments), 212 | instruments=agent_instruments, 213 | validate_results_func=validate_results, 214 | ) 215 | elif team == "data_viz": 216 | return orchestrator.Orchestrator( 217 | name="data_viz_team", 218 | agents=build_data_viz_team(agent_instruments), 219 | validate_results_func=validate_results, 220 | ) 221 | elif team == "scrum_master": 222 | return orchestrator.Orchestrator( 223 | name="scrum_master_team", 224 | agents=build_scrum_master_team(agent_instruments), 225 | instruments=agent_instruments, 226 | validate_results_func=validate_results, 227 | ) 228 | elif team == "data_insights": 229 | return orchestrator.Orchestrator( 230 | name="data_insights_team", 231 | agents=build_insights_team(agent_instruments), 232 | instruments=agent_instruments, 233 | validate_results_func=validate_results, 234 | ) 235 | 236 | raise Exception("Unknown team: " + team) 237 | 238 | 239 | # ------------------------ CUSTOM AGENTS ------------------------ 240 | 241 | 242 | class DefensiveScrumMasterAgent(autogen.ConversableAgent): 243 | """ 244 | Custom agent that uses the guidance function to determine if a message is a SQL NLQ 245 | """ 246 | 247 | def __init__(self, *args, **kwargs): 248 | super().__init__(*args, **kwargs) 249 | # Register the new reply function for this specific agent 250 | self.register_reply(self, self.check_sql_nlq, position=0) 251 | 252 | def check_sql_nlq( 253 | self, 254 | messages: Optional[List[Dict]] = None, 255 | sender: Optional[autogen.Agent] = None, 256 | config: Optional[Any] = None, # Persistent state. 257 | ): 258 | # Check the last received message 259 | last_message = messages[-1]["content"] 260 | 261 | # Use the guidance string to determine if the message is a SQL NLQ 262 | response = guidance( 263 | GUIDANCE_SCRUM_MASTER_SQL_NLQ_PROMPT, potential_nlq=last_message 264 | ) 265 | 266 | # You can return the exact response or just a simplified version, 267 | # here we are just returning the rank for simplicity 268 | rank = response.get("choices", [{}])[0].get("rank", "3") 269 | 270 | return True, rank 271 | 272 | 273 | class InsightsAgent(autogen.ConversableAgent): 274 | """ 275 | Custom agent that uses the guidance function to generate insights in JSON format 276 | """ 277 | 278 | def __init__(self, *args, **kwargs): 279 | super().__init__(*args, **kwargs) 280 | self.register_reply(self, self.generate_insights, position=0) 281 | 282 | def generate_insights( 283 | self, 284 | messages: Optional[List[Dict]] = None, 285 | sender: Optional[autogen.Agent] = None, 286 | config: Optional[Any] = None, 287 | ): 288 | insights = guidance(DATA_INSIGHTS_GUIDANCE_PROMPT) 289 | return True, insights 290 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/agents/instruments.py: -------------------------------------------------------------------------------- 1 | from postgres_da_ai_agent.modules.db import PostgresManager 2 | from postgres_da_ai_agent.modules import file 3 | import os 4 | 5 | BASE_DIR = os.environ.get("BASE_DIR", "./agent_results") 6 | 7 | 8 | class AgentInstruments: 9 | """ 10 | Base class for multli-agent instruments that are tools, state, and functions that an agent can use across the lifecycle of conversations 11 | """ 12 | 13 | def __init__(self) -> None: 14 | self.session_id = None 15 | self.messages = [] 16 | 17 | def __enter__(self): 18 | return self 19 | 20 | def __exit__(self, exc_type, exc_value, traceback): 21 | pass 22 | 23 | def sync_messages(self, messages: list): 24 | """ 25 | Syncs messages with the orchestrator 26 | """ 27 | raise NotImplementedError 28 | 29 | def make_agent_chat_file(self, team_name: str): 30 | return os.path.join(self.root_dir, f"agent_chats_{team_name}.json") 31 | 32 | def make_agent_cost_file(self, team_name: str): 33 | return os.path.join(self.root_dir, f"agent_cost_{team_name}.json") 34 | 35 | @property 36 | def root_dir(self): 37 | return os.path.join(BASE_DIR, self.session_id) 38 | 39 | 40 | class PostgresAgentInstruments(AgentInstruments): 41 | """ 42 | Unified Toolset for the Postgres Data Analytics Multi-Agent System 43 | 44 | Advantages: 45 | - All agents have access to the same state and functions 46 | - Gives agent functions awareness of changing context 47 | - Clear and concise capabilities for agents 48 | - Clean database connection management 49 | 50 | Guidelines: 51 | - Agent Functions should not call other agent functions directly 52 | - Instead Agent Functions should call external lower level modules 53 | - Prefer 1 to 1 mapping of agents and their functions 54 | - The state lifecycle lives between all agent orchestrations 55 | """ 56 | 57 | def __init__(self, db_url: str, session_id: str) -> None: 58 | super().__init__() 59 | 60 | self.db_url = db_url 61 | self.db = None 62 | self.session_id = session_id 63 | self.messages = [] 64 | self.innovation_index = 0 65 | 66 | def __enter__(self): 67 | """ 68 | Support entering the 'with' statement 69 | """ 70 | self.reset_files() 71 | self.db = PostgresManager() 72 | self.db.connect_with_url(self.db_url) 73 | return self, self.db 74 | 75 | def __exit__(self, exc_type, exc_val, exc_tb): 76 | """ 77 | Support exiting the 'with' statement 78 | """ 79 | self.db.close() 80 | 81 | def sync_messages(self, messages: list): 82 | """ 83 | Syncs messages with the orchestrator 84 | """ 85 | self.messages = messages 86 | 87 | def reset_files(self): 88 | """ 89 | Clear everything in the root_dir 90 | """ 91 | 92 | # if it does not exist create it 93 | if not os.path.exists(self.root_dir): 94 | os.makedirs(self.root_dir) 95 | 96 | for fname in os.listdir(self.root_dir): 97 | os.remove(os.path.join(self.root_dir, fname)) 98 | 99 | def get_file_path(self, fname: str): 100 | """ 101 | Get the full path to a file in the root_dir 102 | """ 103 | return os.path.join(self.root_dir, fname) 104 | 105 | # -------------------------- Agent Properties -------------------------- # 106 | 107 | @property 108 | def run_sql_results_file(self): 109 | return self.get_file_path("run_sql_results.json") 110 | 111 | @property 112 | def sql_query_file(self): 113 | return self.get_file_path("sql_query.sql") 114 | 115 | # -------------------------- Agent Functions -------------------------- # 116 | 117 | def run_sql(self, sql: str) -> str: 118 | """ 119 | Run a SQL query against the postgres database 120 | """ 121 | results_as_json = self.db.run_sql(sql) 122 | 123 | fname = self.run_sql_results_file 124 | 125 | # dump these results to a file 126 | with open(fname, "w") as f: 127 | f.write(results_as_json) 128 | 129 | with open(self.sql_query_file, "w") as f: 130 | f.write(sql) 131 | 132 | return "Successfully delivered results to json file" 133 | 134 | def validate_run_sql(self): 135 | """ 136 | validate that the run_sql results file exists and has content 137 | """ 138 | fname = self.run_sql_results_file 139 | 140 | with open(fname, "r") as f: 141 | content = f.read() 142 | 143 | if not content: 144 | return False, f"File {fname} is empty" 145 | 146 | return True, "" 147 | 148 | def write_file(self, content: str): 149 | fname = self.get_file_path(f"write_file.txt") 150 | return file.write_file(fname, content) 151 | 152 | def write_json_file(self, json_str: str): 153 | fname = self.get_file_path(f"write_json_file.json") 154 | return file.write_json_file(fname, json_str) 155 | 156 | def write_yml_file(self, json_str: str): 157 | fname = self.get_file_path(f"write_yml_file.yml") 158 | return file.write_yml_file(fname, json_str) 159 | 160 | def write_innovation_file(self, content: str): 161 | fname = self.get_file_path(f"{self.innovation_index}_innovation_file.json") 162 | file.write_file(fname, content) 163 | self.innovation_index += 1 164 | return f"Successfully wrote innovation file. You can check my work." 165 | 166 | def validate_innovation_files(self): 167 | """ 168 | loop from 0 to innovation_index and verify file exists with content 169 | """ 170 | for i in range(self.innovation_index): 171 | fname = self.get_file_path(f"{i}_innovation_file.json") 172 | with open(fname, "r") as f: 173 | content = f.read() 174 | if not content: 175 | return False, f"File {fname} is empty" 176 | 177 | return True, "" 178 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/agents/turbo4.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | import openai 4 | import time 5 | from typing import Callable, Dict, Any, List, Optional, Union, Tuple 6 | import dotenv 7 | from dataclasses import dataclass, asdict 8 | from openai.types.beta import Thread, Assistant 9 | from openai.types import FileObject 10 | from openai.types.beta.threads.thread_message import ThreadMessage 11 | from openai.types.beta.threads.run_submit_tool_outputs_params import ToolOutput 12 | from postgres_da_ai_agent.modules import llm 13 | from postgres_da_ai_agent.types import Chat, TurboTool 14 | 15 | dotenv.load_dotenv() 16 | 17 | 18 | class Turbo4: 19 | """ 20 | Simple, chainable class for the OpenAI's GPT-4 Assistant APIs. 21 | """ 22 | 23 | def __init__(self): 24 | openai.api_key = os.environ.get("OPENAI_API_KEY") 25 | self.client = openai.OpenAI() 26 | self.map_function_tools: Dict[str, TurboTool] = {} 27 | self.current_thread_id = None 28 | self.thread_messages: List[ThreadMessage] = [] 29 | self.local_messages = [] 30 | self.assistant_id = None 31 | self.polling_interval = ( 32 | 0.5 # Interval in seconds to poll the API for thread run completion 33 | ) 34 | self.model = "gpt-4-1106-preview" 35 | 36 | @property 37 | def chat_messages(self) -> List[Chat]: 38 | return [ 39 | Chat( 40 | from_name=msg.role, 41 | to_name="assistant" if msg.role == "user" else "user", 42 | message=llm.safe_get(msg.model_dump(), "content.0.text.value"), 43 | created=msg.created_at, 44 | ) 45 | for msg in self.thread_messages 46 | ] 47 | 48 | @property 49 | def tool_config(self): 50 | return [tool.config for tool in self.map_function_tools.values()] 51 | 52 | # ------------- Additional Utility Functions ----------------- 53 | 54 | def run_validation(self, validation_func: Callable): 55 | print(f"run_validation({validation_func.__name__})") 56 | validation_func() 57 | return self 58 | 59 | def spy_on_assistant(self, output_file: str): 60 | sorted_messages = sorted( 61 | self.chat_messages, key=lambda msg: msg.created, reverse=False 62 | ) 63 | messages_as_json = [asdict(msg) for msg in sorted_messages] 64 | with open(output_file, "w") as f: 65 | json.dump(messages_as_json, f, indent=2) 66 | 67 | return self 68 | 69 | def get_costs_and_tokens(self, output_file: str) -> Tuple[float, float]: 70 | """ 71 | Get the estimated cost and token usage for the current thread. 72 | 73 | https://openai.com/pricing 74 | 75 | Open questions - how to calculate retrieval and code interpreter costs? 76 | """ 77 | 78 | retrival_costs = 0 79 | code_interpreter_costs = 0 80 | 81 | msgs = [ 82 | llm.safe_get(msg.model_dump(), "content.0.text.value") 83 | for msg in self.thread_messages 84 | ] 85 | joined_msgs = " ".join(msgs) 86 | 87 | msg_cost, tokens = llm.estimate_price_and_tokens(joined_msgs) 88 | 89 | with open(output_file, "w") as f: 90 | json.dump( 91 | { 92 | "cost": msg_cost, 93 | "tokens": tokens, 94 | }, 95 | f, 96 | indent=2, 97 | ) 98 | 99 | return self 100 | 101 | # ------------- CORE ASSISTANTS API FUNCTIONS ----------------- 102 | 103 | def get_or_create_assistant(self, name: str, model: str = "gpt-4-1106-preview"): 104 | print(f"get_or_create_assistant({name}, {model})") 105 | # Retrieve the list of existing assistants 106 | assistants: List[Assistant] = self.client.beta.assistants.list().data 107 | 108 | # Check if an assistant with the given name already exists 109 | for assistant in assistants: 110 | if assistant.name == name: 111 | self.assistant_id = assistant.id 112 | 113 | # update model if different 114 | if assistant.model != model: 115 | print(f"Updating assistant model from {assistant.model} to {model}") 116 | self.client.beta.assistants.update( 117 | assistant_id=self.assistant_id, model=model 118 | ) 119 | break 120 | else: # If no assistant was found with the name, create a new one 121 | assistant = self.client.beta.assistants.create(model=model, name=name) 122 | self.assistant_id = assistant.id 123 | 124 | self.model = model 125 | 126 | return self 127 | 128 | def set_instructions(self, instructions: str): 129 | print(f"set_instructions()") 130 | if self.assistant_id is None: 131 | raise ValueError( 132 | "No assistant has been created or retrieved. Call get_or_create_assistant() first." 133 | ) 134 | # Update the assistant with the new instructions 135 | updated_assistant = self.client.beta.assistants.update( 136 | assistant_id=self.assistant_id, instructions=instructions 137 | ) 138 | return self 139 | 140 | def equip_tools( 141 | self, turbo_tools: List[TurboTool], equip_on_assistant: bool = False 142 | ): 143 | print(f"equip_tools({turbo_tools}, {equip_on_assistant})") 144 | if self.assistant_id is None: 145 | raise ValueError( 146 | "No assistant has been created or retrieved. Call get_or_create_assistant() first." 147 | ) 148 | 149 | # Update the functions dictionary with the new tools 150 | self.map_function_tools = {tool.name: tool for tool in turbo_tools} 151 | 152 | if equip_on_assistant: 153 | # Update the assistant with the new list of tools, replacing any existing tools 154 | updated_assistant = self.client.beta.assistants.update( 155 | tools=self.tool_config, assistant_id=self.assistant_id 156 | ) 157 | 158 | return self 159 | 160 | def make_thread(self): 161 | print(f"make_thread()") 162 | 163 | if self.assistant_id is None: 164 | raise ValueError( 165 | "No assistant has been created. Call create_assistant() first." 166 | ) 167 | 168 | response = self.client.beta.threads.create() 169 | self.current_thread_id = response.id 170 | self.thread_messages = [] 171 | return self 172 | 173 | def add_message(self, message: str, refresh_threads: bool = False): 174 | print(f"add_message({message})") 175 | self.local_messages.append(message) 176 | self.client.beta.threads.messages.create( 177 | thread_id=self.current_thread_id, content=message, role="user" 178 | ) 179 | if refresh_threads: 180 | self.load_threads() 181 | return self 182 | 183 | def load_threads(self): 184 | self.thread_messages = self.client.beta.threads.messages.list( 185 | thread_id=self.current_thread_id 186 | ).data 187 | 188 | def list_steps(self): 189 | print(f"list_steps()") 190 | steps = self.client.beta.threads.runs.steps.list( 191 | thread_id=self.current_thread_id, 192 | run_id=self.run_id, 193 | ) 194 | print("steps", steps) 195 | return steps 196 | 197 | def run_thread(self, toolbox: Optional[List[str]] = None): 198 | print(f"run_thread({toolbox})") 199 | if self.current_thread_id is None: 200 | raise ValueError("No thread has been created. Call make_thread() first.") 201 | if self.local_messages == []: 202 | raise ValueError("No messages have been added to the thread.") 203 | 204 | if toolbox is None: 205 | tools = None 206 | else: 207 | # get tools from toolbox 208 | tools = [self.map_function_tools[tool_name].config for tool_name in toolbox] 209 | 210 | # throw if tool not found 211 | if len(tools) != len(toolbox): 212 | raise ValueError( 213 | f"Tool not found in toolbox. toolbox={toolbox}, tools={tools}. Make sure all tools are equipped on the assistant." 214 | ) 215 | 216 | # refresh current thread 217 | self.load_threads() 218 | 219 | # Start the thread running 220 | run = self.client.beta.threads.runs.create( 221 | thread_id=self.current_thread_id, 222 | assistant_id=self.assistant_id, 223 | tools=tools, 224 | ) 225 | self.run_id = run.id 226 | 227 | # Polling mechanism to wait for thread's run completion or required actions 228 | while True: 229 | # self.list_steps() 230 | 231 | run_status = self.client.beta.threads.runs.retrieve( 232 | thread_id=self.current_thread_id, run_id=self.run_id 233 | ) 234 | if run_status.status == "requires_action": 235 | tool_outputs: List[ToolOutput] = [] 236 | for ( 237 | tool_call 238 | ) in run_status.required_action.submit_tool_outputs.tool_calls: 239 | tool_function = tool_call.function 240 | tool_name = tool_function.name 241 | 242 | # Check if tool_arguments is already a dictionary, if so, proceed directly 243 | if isinstance(tool_function.arguments, dict): 244 | tool_arguments = tool_function.arguments 245 | else: 246 | # Assume the arguments are JSON string and parse them 247 | tool_arguments = json.loads(tool_function.arguments) 248 | 249 | print(f"run_thread() Calling {tool_name}({tool_arguments})") 250 | 251 | # Assuming arguments are passed as a dictionary 252 | function_output = self.map_function_tools[tool_name].function( 253 | **tool_arguments 254 | ) 255 | 256 | tool_outputs.append( 257 | ToolOutput(tool_call_id=tool_call.id, output=function_output) 258 | ) 259 | 260 | # Submit the tool outputs back to the API 261 | self.client.beta.threads.runs.submit_tool_outputs( 262 | thread_id=self.current_thread_id, 263 | run_id=self.run_id, 264 | tool_outputs=[to for to in tool_outputs], 265 | ) 266 | elif run_status.status == "completed": 267 | self.load_threads() 268 | return self 269 | 270 | time.sleep(self.polling_interval) # Wait a little before polling again 271 | 272 | def enable_retrieval(self): 273 | print(f"enable_retrieval()") 274 | if self.assistant_id is None: 275 | raise ValueError( 276 | "No assistant has been created or retrieved. Call get_or_create_assistant() first." 277 | ) 278 | 279 | # Update the assistant with the new list of tools, replacing any existing tools 280 | updated_assistant = self.client.beta.assistants.update( 281 | tools=[{"type": "retrieval"}], assistant_id=self.assistant_id 282 | ) 283 | 284 | return self 285 | 286 | # Future versions: 287 | 288 | # enable code interpreter 289 | # crud files 290 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/main.py: -------------------------------------------------------------------------------- 1 | """ 2 | Heads up: in v7 pyautogen doesn't work with the latest openai version so this file has been commented out via pyproject.toml 3 | """ 4 | 5 | import os 6 | from postgres_da_ai_agent.agents.instruments import PostgresAgentInstruments 7 | from postgres_da_ai_agent.modules.db import PostgresManager 8 | from postgres_da_ai_agent.modules import llm 9 | from postgres_da_ai_agent.modules import orchestrator 10 | from postgres_da_ai_agent.modules import rand 11 | from postgres_da_ai_agent.modules import file 12 | from postgres_da_ai_agent.modules import embeddings 13 | from postgres_da_ai_agent.agents import agents 14 | import dotenv 15 | import argparse 16 | import autogen 17 | 18 | from postgres_da_ai_agent.types import ConversationResult 19 | 20 | 21 | # ---------------- Your Environment Variables ---------------- 22 | 23 | dotenv.load_dotenv() 24 | 25 | assert os.environ.get("DATABASE_URL"), "POSTGRES_CONNECTION_URL not found in .env file" 26 | assert os.environ.get( 27 | "OPENAI_API_KEY" 28 | ), "POSTGRES_CONNECTION_URL not found in .env file" 29 | 30 | 31 | # ---------------- Constants ---------------- 32 | 33 | 34 | DB_URL = os.environ.get("DATABASE_URL") 35 | OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") 36 | 37 | POSTGRES_TABLE_DEFINITIONS_CAP_REF = "TABLE_DEFINITIONS" 38 | 39 | 40 | def main(): 41 | # ---------------- Parse '--prompt' CLI Parameter ---------------- 42 | 43 | parser = argparse.ArgumentParser() 44 | parser.add_argument("--prompt", help="The prompt for the AI") 45 | args = parser.parse_args() 46 | 47 | if not args.prompt: 48 | print("Please provide a prompt") 49 | return 50 | 51 | raw_prompt = args.prompt 52 | 53 | prompt = f"Fulfill this database query: {raw_prompt}. " 54 | 55 | session_id = rand.generate_session_id(raw_prompt) 56 | 57 | # ---------------- Create Agent Instruments And Build Database Connection ---------------- 58 | 59 | with PostgresAgentInstruments(DB_URL, session_id) as (agent_instruments, db): 60 | # ----------- Gate Team: Prevent bad prompts from running and burning your $$$ ------------- 61 | 62 | gate_orchestrator = agents.build_team_orchestrator( 63 | "scrum_master", 64 | agent_instruments, 65 | validate_results=lambda: (True, ""), 66 | ) 67 | 68 | gate_orchestrator: ConversationResult = ( 69 | gate_orchestrator.sequential_conversation(prompt) 70 | ) 71 | 72 | print("gate_orchestrator.last_message_str", gate_orchestrator.last_message_str) 73 | 74 | nlq_confidence = int(gate_orchestrator.last_message_str) 75 | 76 | match nlq_confidence: 77 | case (1 | 2): 78 | print(f"❌ Gate Team Rejected - Confidence too low: {nlq_confidence}") 79 | return 80 | case (3 | 4 | 5): 81 | print(f"✅ Gate Team Approved - Valid confidence: {nlq_confidence}") 82 | case _: 83 | print("❌ Gate Team Rejected - Invalid response") 84 | return 85 | 86 | # -------- BUILD TABLE DEFINITIONS ----------- 87 | 88 | map_table_name_to_table_def = db.get_table_definition_map_for_embeddings() 89 | 90 | database_embedder = embeddings.DatabaseEmbedder() 91 | 92 | for name, table_def in map_table_name_to_table_def.items(): 93 | database_embedder.add_table(name, table_def) 94 | 95 | similar_tables = database_embedder.get_similar_tables(raw_prompt, n=5) 96 | 97 | table_definitions = database_embedder.get_table_definitions_from_names( 98 | similar_tables 99 | ) 100 | 101 | related_table_names = db.get_related_tables(similar_tables, n=3) 102 | 103 | core_and_related_table_definitions = ( 104 | database_embedder.get_table_definitions_from_names( 105 | related_table_names + similar_tables 106 | ) 107 | ) 108 | 109 | prompt = llm.add_cap_ref( 110 | prompt, 111 | f"Use these {POSTGRES_TABLE_DEFINITIONS_CAP_REF} to satisfy the database query.", 112 | POSTGRES_TABLE_DEFINITIONS_CAP_REF, 113 | table_definitions, 114 | ) 115 | 116 | # ----------- Data Eng Team: Based on a sql table definitions and a prompt create an sql statement and execute it ------------- 117 | 118 | data_eng_orchestrator = agents.build_team_orchestrator( 119 | "data_eng", 120 | agent_instruments, 121 | validate_results=agent_instruments.validate_run_sql, 122 | ) 123 | 124 | data_eng_conversation_result: ConversationResult = ( 125 | data_eng_orchestrator.sequential_conversation(prompt) 126 | ) 127 | 128 | match data_eng_conversation_result: 129 | case ConversationResult( 130 | success=True, cost=data_eng_cost, tokens=data_eng_tokens 131 | ): 132 | print( 133 | f"✅ Orchestrator was successful. Team: {data_eng_orchestrator.name}" 134 | ) 135 | print( 136 | f"💰📊🤖 {data_eng_orchestrator.name} Cost: {data_eng_cost}, tokens: {data_eng_tokens}" 137 | ) 138 | case _: 139 | print( 140 | f"❌ Orchestrator failed. Team: {data_eng_orchestrator.name} Failed" 141 | ) 142 | 143 | # ----------- Data Insights Team: Based on sql table definitions and a prompt generate novel insights ------------- 144 | 145 | innovation_prompt = f"Given this database query: '{raw_prompt}'. Generate novel insights and new database queries to give business insights." 146 | 147 | insights_prompt = llm.add_cap_ref( 148 | innovation_prompt, 149 | f"Use these {POSTGRES_TABLE_DEFINITIONS_CAP_REF} to satisfy the database query.", 150 | POSTGRES_TABLE_DEFINITIONS_CAP_REF, 151 | core_and_related_table_definitions, 152 | ) 153 | 154 | data_insights_orchestrator = agents.build_team_orchestrator( 155 | "data_insights", 156 | agent_instruments, 157 | validate_results=agent_instruments.validate_innovation_files, 158 | ) 159 | 160 | data_insights_conversation_result: ConversationResult = ( 161 | data_insights_orchestrator.round_robin_conversation( 162 | insights_prompt, loops=1 163 | ) 164 | ) 165 | 166 | match data_insights_conversation_result: 167 | case ConversationResult( 168 | success=True, cost=data_insights_cost, tokens=data_insights_tokens 169 | ): 170 | print( 171 | f"✅ Orchestrator was successful. Team: {data_insights_orchestrator.name}" 172 | ) 173 | print( 174 | f"💰📊🤖 {data_insights_orchestrator.name} Cost: {data_insights_cost}, tokens: {data_insights_tokens}" 175 | ) 176 | case _: 177 | print( 178 | f"❌ Orchestrator failed. Team: {data_insights_orchestrator.name} Failed" 179 | ) 180 | 181 | 182 | if __name__ == "__main__": 183 | main() 184 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/modules/db.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | import json 3 | import psycopg2 4 | from psycopg2.sql import SQL, Identifier 5 | 6 | 7 | class PostgresManager: 8 | """ 9 | A class to manage postgres connections and queries 10 | """ 11 | 12 | def __init__(self): 13 | self.conn = None 14 | self.cur = None 15 | 16 | def __enter__(self): 17 | return self 18 | 19 | def __exit__(self, exc_type, exc_val, exc_tb): 20 | if self.cur: 21 | self.cur.close() 22 | if self.conn: 23 | self.conn.close() 24 | 25 | def connect_with_url(self, url): 26 | self.conn = psycopg2.connect(url) 27 | self.cur = self.conn.cursor() 28 | 29 | def close(self): 30 | if self.cur: 31 | self.cur.close() 32 | if self.conn: 33 | self.conn.close() 34 | 35 | def run_sql(self, sql) -> str: 36 | """ 37 | Run a SQL query against the postgres database 38 | """ 39 | self.cur.execute(sql) 40 | columns = [desc[0] for desc in self.cur.description] 41 | res = self.cur.fetchall() 42 | 43 | list_of_dicts = [dict(zip(columns, row)) for row in res] 44 | 45 | json_result = json.dumps(list_of_dicts, indent=4, default=self.datetime_handler) 46 | 47 | return json_result 48 | 49 | def datetime_handler(self, obj): 50 | """ 51 | Handle datetime objects when serializing to JSON. 52 | """ 53 | if isinstance(obj, datetime): 54 | return obj.isoformat() 55 | return str(obj) # or just return the object unchanged, or another default value 56 | 57 | def get_table_definition(self, table_name): 58 | """ 59 | Generate the 'create' definition for a table 60 | """ 61 | 62 | get_def_stmt = """ 63 | SELECT pg_class.relname as tablename, 64 | pg_attribute.attnum, 65 | pg_attribute.attname, 66 | format_type(atttypid, atttypmod) 67 | FROM pg_class 68 | JOIN pg_namespace ON pg_namespace.oid = pg_class.relnamespace 69 | JOIN pg_attribute ON pg_attribute.attrelid = pg_class.oid 70 | WHERE pg_attribute.attnum > 0 71 | AND pg_class.relname = %s 72 | AND pg_namespace.nspname = 'public' -- Assuming you're interested in public schema 73 | """ 74 | self.cur.execute(get_def_stmt, (table_name,)) 75 | rows = self.cur.fetchall() 76 | create_table_stmt = "CREATE TABLE {} (\n".format(table_name) 77 | for row in rows: 78 | create_table_stmt += "{} {},\n".format(row[2], row[3]) 79 | create_table_stmt = create_table_stmt.rstrip(",\n") + "\n);" 80 | return create_table_stmt 81 | 82 | def get_all_table_names(self): 83 | """ 84 | Get all table names in the database 85 | """ 86 | get_all_tables_stmt = ( 87 | "SELECT tablename FROM pg_tables WHERE schemaname = 'public';" 88 | ) 89 | self.cur.execute(get_all_tables_stmt) 90 | return [row[0] for row in self.cur.fetchall()] 91 | 92 | def get_table_definitions_for_prompt(self): 93 | """ 94 | Get all table 'create' definitions in the database 95 | """ 96 | table_names = self.get_all_table_names() 97 | definitions = [] 98 | for table_name in table_names: 99 | definitions.append(self.get_table_definition(table_name)) 100 | return "\n\n".join(definitions) 101 | 102 | def get_table_definition_map_for_embeddings(self): 103 | """ 104 | Creates a map of table names to table definitions 105 | """ 106 | table_names = self.get_all_table_names() 107 | definitions = {} 108 | for table_name in table_names: 109 | definitions[table_name] = self.get_table_definition(table_name) 110 | return definitions 111 | 112 | def get_related_tables(self, table_list, n=2): 113 | """ 114 | Get tables that have foreign keys referencing the given table 115 | """ 116 | 117 | related_tables_dict = {} 118 | 119 | for table in table_list: 120 | # Query to fetch tables that have foreign keys referencing the given table 121 | self.cur.execute( 122 | """ 123 | SELECT 124 | a.relname AS table_name 125 | FROM 126 | pg_constraint con 127 | JOIN pg_class a ON a.oid = con.conrelid 128 | WHERE 129 | confrelid = (SELECT oid FROM pg_class WHERE relname = %s) 130 | LIMIT %s; 131 | """, 132 | (table, n), 133 | ) 134 | 135 | related_tables = [row[0] for row in self.cur.fetchall()] 136 | 137 | # Query to fetch tables that the given table references 138 | self.cur.execute( 139 | """ 140 | SELECT 141 | a.relname AS referenced_table_name 142 | FROM 143 | pg_constraint con 144 | JOIN pg_class a ON a.oid = con.confrelid 145 | WHERE 146 | conrelid = (SELECT oid FROM pg_class WHERE relname = %s) 147 | LIMIT %s; 148 | """, 149 | (table, n), 150 | ) 151 | 152 | related_tables += [row[0] for row in self.cur.fetchall()] 153 | 154 | related_tables_dict[table] = related_tables 155 | 156 | # convert dict to list and remove dups 157 | related_tables_list = [] 158 | for table, related_tables in related_tables_dict.items(): 159 | related_tables_list += related_tables 160 | 161 | related_tables_list = list(set(related_tables_list)) 162 | 163 | return related_tables_list 164 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/modules/embeddings.py: -------------------------------------------------------------------------------- 1 | from sklearn.metrics.pairwise import cosine_similarity 2 | from transformers import BertTokenizer, BertModel 3 | 4 | from postgres_da_ai_agent.modules.db import PostgresManager 5 | 6 | 7 | class DatabaseEmbedder: 8 | """ 9 | This class is responsible for embedding database table definitions and 10 | computing similarity between user queries and table definitions. 11 | """ 12 | 13 | def __init__(self, db: PostgresManager): 14 | self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") 15 | self.model = BertModel.from_pretrained("bert-base-uncased") 16 | self.map_name_to_embeddings = {} 17 | self.map_name_to_table_def = {} 18 | self.db = db 19 | 20 | def get_similar_table_defs_for_prompt(self, prompt: str, n_similar=5, n_foreign=0): 21 | map_table_name_to_table_def = self.db.get_table_definition_map_for_embeddings() 22 | for name, table_def in map_table_name_to_table_def.items(): 23 | self.add_table(name, table_def) 24 | 25 | similar_tables = self.get_similar_tables(prompt, n=n_similar) 26 | 27 | table_definitions = self.get_table_definitions_from_names(similar_tables) 28 | 29 | if n_foreign > 0: 30 | foreign_table_names = self.db.get_foreign_tables(similar_tables, n=3) 31 | 32 | table_definitions = self.get_table_definitions_from_names( 33 | foreign_table_names + similar_tables 34 | ) 35 | 36 | return table_definitions 37 | 38 | def add_table(self, table_name: str, text_representation: str): 39 | """ 40 | Add a table to the database embedder. 41 | Map the table name to its embedding and text representation. 42 | """ 43 | self.map_name_to_embeddings[table_name] = self.compute_embeddings( 44 | text_representation 45 | ) 46 | 47 | self.map_name_to_table_def[table_name] = text_representation 48 | 49 | def compute_embeddings(self, text): 50 | """ 51 | Compute embeddings for a given text using the BERT model. 52 | """ 53 | inputs = self.tokenizer( 54 | text, return_tensors="pt", truncation=True, padding=True, max_length=512 55 | ) 56 | outputs = self.model(**inputs) 57 | return outputs["pooler_output"].detach().numpy() 58 | 59 | def get_similar_tables_via_embeddings(self, query, n=3): 60 | """ 61 | Given a query, find the top 'n' tables that are most similar to it. 62 | 63 | Args: 64 | - query (str): The user's natural language query. 65 | - n (int, optional): Number of top tables to return. Defaults to 3. 66 | 67 | Returns: 68 | - list: Top 'n' table names ranked by their similarity to the query. 69 | """ 70 | # Compute the embedding for the user's query 71 | query_embedding = self.compute_embeddings(query) 72 | # Calculate cosine similarity between the query and all tables 73 | similarities = { 74 | table: cosine_similarity(query_embedding, emb)[0][0] 75 | for table, emb in self.map_name_to_embeddings.items() 76 | } 77 | # Rank tables based on their similarity scores and return top 'n' 78 | return sorted(similarities, key=similarities.get, reverse=True)[:n] 79 | 80 | def get_similar_table_names_via_word_match(self, query: str): 81 | """ 82 | if any word in our query is a table name, add the table to a list 83 | """ 84 | 85 | tables = [] 86 | 87 | for table_name in self.map_name_to_table_def.keys(): 88 | if table_name.lower() in query.lower(): 89 | tables.append(table_name) 90 | 91 | return tables 92 | 93 | def get_similar_tables(self, query: str, n=3): 94 | """ 95 | combines results from get_similar_tables_via_embeddings and get_similar_table_names_via_word_match 96 | """ 97 | 98 | similar_tables_via_embeddings = self.get_similar_tables_via_embeddings(query, n) 99 | similar_tables_via_word_match = self.get_similar_table_names_via_word_match( 100 | query 101 | ) 102 | 103 | return similar_tables_via_embeddings + similar_tables_via_word_match 104 | 105 | def get_table_definitions_from_names(self, table_names: list) -> str: 106 | """ 107 | Given a list of table names, return their table definitions. 108 | """ 109 | table_defs = [ 110 | self.map_name_to_table_def[table_name] for table_name in table_names 111 | ] 112 | return "\n\n".join(table_defs) 113 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/modules/file.py: -------------------------------------------------------------------------------- 1 | import json 2 | import yaml 3 | 4 | 5 | def write_file(fname, content): 6 | with open(fname, "w") as f: 7 | f.write(content) 8 | 9 | 10 | def write_json_file(fname, json_str: str): 11 | # convert ' to " 12 | json_str = json_str.replace("'", '"') 13 | 14 | # Convert the string to a Python object 15 | data = json.loads(json_str) 16 | 17 | # Write the Python object to the file as JSON 18 | with open(fname, "w") as f: 19 | json.dump(data, f, indent=4) 20 | 21 | 22 | def write_yml_file(fname, json_str: str): 23 | # Try to replace single quotes with double quotes for JSON 24 | cleaned_json_str = json_str.replace("'", '"') 25 | 26 | # Safely convert the JSON string to a Python object 27 | try: 28 | data = json.loads(cleaned_json_str) 29 | except json.JSONDecodeError as e: 30 | print(f"Error decoding JSON: {e}") 31 | return 32 | 33 | # Write the Python object to the file as YAML 34 | with open(fname, "w") as f: 35 | yaml.dump(data, f) 36 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/modules/llm.py: -------------------------------------------------------------------------------- 1 | """ 2 | Purpose: 3 | Interact with the OpenAI API. 4 | Provide supporting prompt engineering functions. 5 | """ 6 | 7 | import json 8 | import sys 9 | from dotenv import load_dotenv 10 | import os 11 | from typing import Any, Dict, List 12 | import openai 13 | import tiktoken 14 | 15 | from postgres_da_ai_agent.types import TurboTool 16 | 17 | # load .env file 18 | load_dotenv() 19 | 20 | assert os.environ.get("OPENAI_API_KEY") 21 | 22 | # get openai api key 23 | openai.api_key = os.environ.get("OPENAI_API_KEY") 24 | 25 | # ------------------ helpers ------------------ 26 | 27 | 28 | def safe_get(data, dot_chained_keys): 29 | """ 30 | {'a': {'b': [{'c': 1}]}} 31 | safe_get(data, 'a.b.0.c') -> 1 32 | """ 33 | keys = dot_chained_keys.split(".") 34 | for key in keys: 35 | try: 36 | if isinstance(data, list): 37 | data = data[int(key)] 38 | else: 39 | data = data[key] 40 | except (KeyError, TypeError, IndexError): 41 | return None 42 | return data 43 | 44 | 45 | def response_parser(response: Dict[str, Any]): 46 | return safe_get(response, "choices.0.message.content") 47 | 48 | 49 | # ------------------ content generators ------------------ 50 | 51 | 52 | def prompt( 53 | prompt: str, 54 | model: str = "gpt-4-1106-preview", 55 | instructions: str = "You are a helpful assistant.", 56 | ) -> str: 57 | """ 58 | Generate a response from a prompt using the OpenAI API. 59 | """ 60 | 61 | if not openai.api_key: 62 | sys.exit( 63 | """ 64 | ERORR: OpenAI API key not found. Please export your key to OPENAI_API_KEY 65 | Example bash command: 66 | export OPENAI_API_KEY= 67 | """ 68 | ) 69 | 70 | response = openai.chat.completions.create( 71 | model=model, 72 | messages=[ 73 | { 74 | "role": "system", 75 | "content": instructions, # Added instructions as a system message 76 | }, 77 | { 78 | "role": "user", 79 | "content": prompt, 80 | }, 81 | ], 82 | ) 83 | 84 | return response_parser(response.model_dump()) 85 | 86 | 87 | def prompt_func( 88 | prompt: str, 89 | turbo_tools: List[TurboTool], 90 | model: str = "gpt-4-1106-preview", 91 | instructions: str = "You are a helpful assistant.", 92 | ) -> str: 93 | """ 94 | Generate a response from a prompt using the OpenAI API. 95 | Force function calls to the provided turbo tools. 96 | 97 | :param prompt: The prompt to send to the model. 98 | :param turbo_tools: List of TurboTool objects each containing the tool's name, configuration, and function. 99 | :param model: The model version to use, default is 'gpt-4-1106-preview'. 100 | :return: The response generated by the model. 101 | """ 102 | 103 | messages = [{"role": "user", "content": prompt}] 104 | tools = [turbo_tool.config for turbo_tool in turbo_tools] 105 | 106 | tool_choice = ( 107 | "auto" 108 | if len(turbo_tools) > 1 109 | else {"type": "function", "function": {"name": turbo_tools[0].name}} 110 | ) 111 | 112 | messages.insert( 113 | 0, {"role": "system", "content": instructions} 114 | ) # Insert instructions as the first system message 115 | response = openai.chat.completions.create( 116 | model=model, messages=messages, tools=tools, tool_choice=tool_choice 117 | ) 118 | 119 | response_message = response.choices[0].message 120 | tool_calls = response_message.tool_calls 121 | 122 | func_responses = [] 123 | 124 | if tool_calls: 125 | messages.append(response_message) 126 | 127 | for tool_call in tool_calls: 128 | for turbo_tool in turbo_tools: 129 | if tool_call.function.name == turbo_tool.name: 130 | function_response = turbo_tool.function( 131 | **json.loads(tool_call.function.arguments) 132 | ) 133 | 134 | func_responses.append(function_response) 135 | 136 | message_to_append = { 137 | "tool_call_id": tool_call.id, 138 | "role": "tool", 139 | "name": turbo_tool.name, 140 | "content": function_response, 141 | } 142 | messages.append(message_to_append) 143 | break 144 | 145 | return func_responses 146 | 147 | 148 | def prompt_json_response( 149 | prompt: str, 150 | model: str = "gpt-4-1106-preview", 151 | instructions: str = "You are a helpful assistant.", 152 | ) -> str: 153 | """ 154 | Generate a response from a prompt using the OpenAI API. 155 | 156 | Example: 157 | res = llm.prompt_json_response( 158 | f"You're a data innovator. You analyze SQL databases table structure and generate 3 novel insights for your team to reflect on and query. 159 | Generate insights for this this prompt: {prompt}. 160 | Format your insights in JSON format. Respond in this json format [{{insight, sql, actionable_business_value}}, ...]", 161 | ) 162 | """ 163 | 164 | if not openai.api_key: 165 | sys.exit( 166 | """ 167 | ERORR: OpenAI API key not found. Please export your key to OPENAI_API_KEY 168 | Example bash command: 169 | export OPENAI_API_KEY= 170 | """ 171 | ) 172 | 173 | response = openai.chat.completions.create( 174 | model=model, 175 | messages=[ 176 | { 177 | "role": "system", 178 | "content": instructions, # Added instructions as a system message 179 | }, 180 | { 181 | "role": "user", 182 | "content": prompt, 183 | }, 184 | ], 185 | response_format={"type": "json_object"}, 186 | ) 187 | 188 | return response_parser(response.model_dump()) 189 | 190 | 191 | def add_cap_ref( 192 | prompt: str, prompt_suffix: str, cap_ref: str, cap_ref_content: str 193 | ) -> str: 194 | """ 195 | Attaches a capitalized reference to the prompt. 196 | Example 197 | prompt = 'Refactor this code.' 198 | prompt_suffix = 'Make it more readable using this EXAMPLE.' 199 | cap_ref = 'EXAMPLE' 200 | cap_ref_content = 'def foo():\n return True' 201 | returns 'Refactor this code. Make it more readable using this EXAMPLE.\n\nEXAMPLE\n\ndef foo():\n return True' 202 | """ 203 | 204 | new_prompt = f"""{prompt} {prompt_suffix}\n\n{cap_ref}\n\n{cap_ref_content}""" 205 | 206 | return new_prompt 207 | 208 | 209 | def count_tokens(text: str): 210 | """ 211 | Count the number of tokens in a string. 212 | """ 213 | enc = tiktoken.get_encoding("cl100k_base") 214 | return len(enc.encode(text)) 215 | 216 | 217 | map_model_to_cost_per_1k_tokens = { 218 | "gpt-4": 0.075, # ($0.03 Input Tokens + $0.06 Output Tokens) / 2 219 | "gpt-4-1106-preview": 0.02, # ($0.01 Input Tokens + $0.03 Output Tokens) / 2 220 | "gpt-4-1106-vision-preview": 0.02, # ($0.01 Input Tokens + $0.03 Output Tokens) / 2 221 | "gpt-3.5-turbo-1106": 0.0015, # ($0.001 Input Tokens + $0.002 Output Tokens) / 2 222 | } 223 | 224 | 225 | def estimate_price_and_tokens(text, model="gpt-4"): 226 | """ 227 | Conservative estimate the price and tokens for a given text. 228 | """ 229 | # round up to the output tokens 230 | COST_PER_1k_TOKENS = map_model_to_cost_per_1k_tokens[model] 231 | 232 | tokens = count_tokens(text) 233 | 234 | estimated_cost = (tokens / 1000) * COST_PER_1k_TOKENS 235 | 236 | # round 237 | estimated_cost = round(estimated_cost, 2) 238 | 239 | return estimated_cost, tokens 240 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/modules/orchestrator.py: -------------------------------------------------------------------------------- 1 | import dataclasses 2 | import json 3 | from typing import List, Optional, Tuple 4 | import autogen 5 | from postgres_da_ai_agent.agents.instruments import AgentInstruments 6 | from postgres_da_ai_agent.modules import llm 7 | from postgres_da_ai_agent.types import Chat, ConversationResult 8 | 9 | 10 | class Orchestrator: 11 | """ 12 | Orchestrators manage conversations between multi-agent teams. 13 | """ 14 | 15 | def __init__( 16 | self, 17 | name: str, 18 | agents: List[autogen.ConversableAgent], 19 | instruments: AgentInstruments, 20 | validate_results_func: callable = None, 21 | ): 22 | # Name of agent team 23 | self.name = name 24 | 25 | # List of agents 26 | self.agents = agents 27 | 28 | # List of raw messages - partially redundant due to self.chats 29 | self.messages = [] 30 | 31 | # Agent instruments - state and functions that agents can use 32 | self.instruments = instruments 33 | 34 | # List of chats - {from, to, message} 35 | self.chats: List[Chat] = [] 36 | 37 | # Function to validate results at the end of every conversation 38 | self.validate_results_func: callable = validate_results_func 39 | 40 | if len(self.agents) < 2: 41 | raise Exception("Orchestrator needs at least two agents") 42 | 43 | @property 44 | def total_agents(self): 45 | return len(self.agents) 46 | 47 | @property 48 | def last_message_is_dict(self): 49 | return isinstance(self.messages[-1], dict) 50 | 51 | @property 52 | def last_message_is_string(self): 53 | return isinstance(self.messages[-1], str) 54 | 55 | @property 56 | def last_message_is_func_call(self): 57 | return self.last_message_is_dict and self.latest_message.get( 58 | "function_call", None 59 | ) 60 | 61 | @property 62 | def last_message_is_content(self): 63 | return self.last_message_is_dict and self.latest_message.get("content", None) 64 | 65 | @property 66 | def latest_message(self) -> Optional[str]: 67 | if not self.messages: 68 | return None 69 | return self.messages[-1] 70 | 71 | @property 72 | def last_message_always_string(self): 73 | if not self.messages: 74 | return "" 75 | if self.last_message_is_content: 76 | return self.latest_message.get("content", "") 77 | return str(self.messages[-1]) 78 | 79 | def handle_validate_func(self) -> Tuple[bool, str]: 80 | """ 81 | Run the validate_results_func if it exists 82 | """ 83 | if self.validate_results_func: 84 | return self.validate_results_func() 85 | return True, "" 86 | 87 | def send_message( 88 | self, 89 | from_agent: autogen.ConversableAgent, 90 | to_agent: autogen.ConversableAgent, 91 | message: str, 92 | ): 93 | """ 94 | Send a message from one agent to another. 95 | Record the message in chat log in the orchestrator 96 | """ 97 | 98 | from_agent.send(message, to_agent) 99 | 100 | self.chats.append( 101 | Chat( 102 | from_name=from_agent.name, 103 | to_name=to_agent.name, 104 | message=str(message), 105 | ) 106 | ) 107 | 108 | def add_message(self, message: str): 109 | """ 110 | Add a message to the orchestrator 111 | """ 112 | self.messages.append(message) 113 | 114 | def get_message_as_str(self): 115 | """ 116 | Get all messages as a string 117 | """ 118 | 119 | messages_as_str = "" 120 | 121 | for message in self.messages: 122 | if message is None: 123 | continue 124 | 125 | if isinstance(message, dict): 126 | content_from_dict = message.get("content", None) 127 | func_call_from_dict = message.get("function_call", None) 128 | content = content_from_dict or func_call_from_dict 129 | if not content: 130 | continue 131 | messages_as_str += str(content) 132 | else: 133 | messages_as_str += str(message) 134 | 135 | return messages_as_str 136 | 137 | def get_cost_and_tokens(self): 138 | return llm.estimate_price_and_tokens(self.get_message_as_str()) 139 | 140 | def has_functions(self, agent: autogen.ConversableAgent): 141 | return len(agent._function_map) > 0 142 | 143 | def basic_chat( 144 | self, 145 | agent_a: autogen.ConversableAgent, 146 | agent_b: autogen.ConversableAgent, 147 | message: str, 148 | ): 149 | print(f"basic_chat(): {agent_a.name} -> {agent_b.name}") 150 | 151 | self.send_message(agent_a, agent_b, message) 152 | 153 | reply = agent_b.generate_reply(sender=agent_a) 154 | 155 | self.add_message(reply) 156 | 157 | print(f"basic_chat(): replied with:", reply) 158 | 159 | def memory_chat( 160 | self, 161 | agent_a: autogen.ConversableAgent, 162 | agent_b: autogen.ConversableAgent, 163 | message: str, 164 | ): 165 | print(f"memory_chat() '{agent_a.name}' --> '{agent_b.name}'") 166 | 167 | self.send_message(agent_a, agent_b, message) 168 | 169 | reply = agent_b.generate_reply(sender=agent_a) 170 | 171 | self.send_message(agent_b, agent_b, message) 172 | 173 | self.add_message(reply) 174 | 175 | def function_chat( 176 | self, 177 | agent_a: autogen.ConversableAgent, 178 | agent_b: autogen.ConversableAgent, 179 | message: str, 180 | ): 181 | print(f"function_call(): {agent_a.name} -> {agent_b.name}") 182 | 183 | self.basic_chat(agent_a, agent_a, message) 184 | 185 | assert self.last_message_is_content 186 | 187 | self.basic_chat(agent_a, agent_b, self.latest_message) 188 | 189 | def self_function_chat(self, agent: autogen.ConversableAgent, message: str): 190 | print(f"self_function_chat(): {agent.name} -> {agent.name}") 191 | 192 | self.send_message(agent, agent, message) 193 | 194 | reply = agent.generate_reply(sender=agent) 195 | 196 | self.send_message(agent, agent, message) 197 | 198 | self.add_message(reply) 199 | 200 | print(f"self_function_chat(): replied with:", reply) 201 | 202 | def spy_on_agents(self, append_to_file: bool = True): 203 | conversations = [] 204 | 205 | for chat in self.chats: 206 | conversations.append(dataclasses.asdict(chat)) 207 | 208 | if append_to_file: 209 | file_name = self.instruments.make_agent_chat_file(self.name) 210 | with open(file_name, "w") as f: 211 | f.write(json.dumps(conversations, indent=4)) 212 | 213 | def sequential_conversation(self, prompt: str) -> ConversationResult: 214 | """ 215 | Runs a sequential conversation between agents. 216 | 217 | The most common type of conversation. 218 | 219 | For example 220 | "Agent A" -> "Agent B" -> "Agent C" -> "Agent D" -> "Agent E" 221 | """ 222 | 223 | print(f"\n\n--------- {self.name} Orchestrator Starting ---------\n\n") 224 | 225 | self.add_message(prompt) 226 | 227 | for idx, agent in enumerate(self.agents): 228 | agent_a = self.agents[idx] 229 | agent_b = self.agents[idx + 1] 230 | 231 | print( 232 | f"\n\n--------- Running iteration {idx} with (agent_a: {agent_a.name}, agent_b: {agent_b.name}) ---------\n\n" 233 | ) 234 | 235 | # agent_a -> chat -> agent_b 236 | if self.last_message_is_string: 237 | self.basic_chat(agent_a, agent_b, self.latest_message) 238 | 239 | # agent_a -> func() -> agent_b 240 | if self.last_message_is_func_call and self.has_functions(agent_a): 241 | self.function_chat(agent_a, agent_b, self.latest_message) 242 | 243 | self.spy_on_agents() 244 | 245 | if idx == self.total_agents - 2: 246 | if self.has_functions(agent_b): 247 | # agent_b -> func() -> agent_b 248 | self.self_function_chat(agent_b, self.latest_message) 249 | 250 | print(f"-------- Orchestrator Complete --------\n\n") 251 | 252 | was_successful, error_message = self.handle_validate_func() 253 | 254 | self.spy_on_agents() 255 | 256 | cost, tokens = self.get_cost_and_tokens() 257 | 258 | return ConversationResult( 259 | success=was_successful, 260 | messages=self.messages, 261 | cost=cost, 262 | tokens=tokens, 263 | last_message_str=self.last_message_always_string, 264 | error_message=error_message, 265 | ) 266 | 267 | def broadcast_conversation(self, prompt: str) -> ConversationResult: 268 | """ 269 | Broadcast a message from agent_a to all agents. 270 | 271 | For example 272 | "Agent A" -> "Agent B" 273 | "Agent A" -> "Agent C" 274 | "Agent A" -> "Agent D" 275 | "Agent A" -> "Agent E" 276 | """ 277 | 278 | print(f"\n\n--------- {self.name} Orchestrator Starting ---------\n\n") 279 | 280 | self.add_message(prompt) 281 | 282 | broadcast_agent = self.agents[0] 283 | 284 | for idx, agent_iterate in enumerate(self.agents[1:]): 285 | print( 286 | f"\n\n--------- Running iteration {idx} with (agent_broadcast: {broadcast_agent.name}, agent_iteration: {agent_iterate.name}) ---------\n\n" 287 | ) 288 | 289 | # agent_a -> chat -> agent_b 290 | if self.last_message_is_string: 291 | self.memory_chat(broadcast_agent, agent_iterate, prompt) 292 | 293 | # agent_b -> func() -> agent_b 294 | if self.last_message_is_func_call and self.has_functions(agent_iterate): 295 | self.function_chat(agent_iterate, agent_iterate, self.latest_message) 296 | 297 | self.spy_on_agents() 298 | 299 | print(f"-------- Orchestrator Complete --------\n\n") 300 | 301 | was_successful, error_message = self.handle_validate_func() 302 | 303 | if was_successful: 304 | print(f"✅ Orchestrator was successful") 305 | else: 306 | print(f"❌ Orchestrator failed") 307 | 308 | cost, tokens = self.get_cost_and_tokens() 309 | 310 | return ConversationResult( 311 | success=was_successful, 312 | messages=self.messages, 313 | cost=cost, 314 | tokens=tokens, 315 | last_message_str=self.last_message_always_string, 316 | error_message=error_message, 317 | ) 318 | 319 | def round_robin_conversation( 320 | self, prompt: str, loops: int = 1 321 | ) -> ConversationResult: 322 | """ 323 | Runs a basic round robin conversation between agents: 324 | 325 | Example for a setup with agents A, B, and C: 326 | (1) 327 | A -> B 328 | B -> C 329 | C -> A 330 | 331 | (2) 332 | A -> B 333 | B -> C 334 | C -> A 335 | 336 | ... 337 | 338 | `loops` determines the number of times the sequence is repeated. 339 | """ 340 | 341 | print( 342 | f"\n\n🚀 --------- {self.name} ::: Orchestrator Starting ::: Round Robin Conversation ---------\n\n" 343 | ) 344 | 345 | self.add_message(prompt) 346 | 347 | total_iterations = loops * len(self.agents) 348 | for iteration in range(total_iterations): 349 | idx = iteration % len(self.agents) 350 | agent_a = self.agents[idx] 351 | agent_b = self.agents[(idx + 1) % len(self.agents)] 352 | 353 | print( 354 | f"\n\n💬 --------- Running iteration {iteration} with conversation ({agent_a.name} -> {agent_b.name}) ---------\n\n", 355 | ) 356 | 357 | # if we're back at the first agent, we need to reset the last message to the prompt 358 | if iteration % (len(self.agents)) == 0: 359 | self.add_message(prompt) 360 | 361 | # agent_a -> chat -> agent_b 362 | if self.last_message_is_string: 363 | self.basic_chat(agent_a, agent_b, self.latest_message) 364 | 365 | # agent_a -> func() -> agent_b 366 | if self.last_message_is_func_call and self.has_functions(agent_a): 367 | self.function_chat(agent_a, agent_b, self.latest_message) 368 | 369 | self.spy_on_agents() 370 | 371 | print(f"-------- Orchestrator Complete --------\n\n") 372 | 373 | self.spy_on_agents() 374 | 375 | agents_were_successful, error_message = self.handle_validate_func() 376 | 377 | cost, tokens = self.get_cost_and_tokens() 378 | 379 | conversation_result: ConversationResult = ConversationResult( 380 | success=agents_were_successful, 381 | messages=self.messages, 382 | cost=cost, 383 | tokens=tokens, 384 | last_message_str=self.last_message_always_string, 385 | error_message=error_message, 386 | ) 387 | 388 | return conversation_result 389 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/modules/rand.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | import uuid 3 | 4 | 5 | def generate_session_id(raw_prompt: str): 6 | """ 7 | "get jobs with 'Completed' or 'Started' status" 8 | 9 | -> 10 | 11 | "get_jobs_with_Completed_or_Started_status__12_22_22" 12 | """ 13 | 14 | now = datetime.now() 15 | hours = now.hour 16 | minutes = now.minute 17 | seconds = now.second 18 | 19 | short_time_mm_ss = f"{hours:02}_{minutes:02}_{seconds:02}" 20 | 21 | lower_case = raw_prompt.lower() 22 | no_spaces = lower_case.replace(" ", "_") 23 | no_quotes = no_spaces.replace("'", "") 24 | shorter = no_quotes[:30] 25 | with_uuid = shorter + "__" + short_time_mm_ss 26 | return with_uuid 27 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/turbo_main.py: -------------------------------------------------------------------------------- 1 | from postgres_da_ai_agent.agents.turbo4 import Turbo4 2 | from postgres_da_ai_agent.types import Chat, TurboTool 3 | from typing import List, Callable 4 | import os 5 | from postgres_da_ai_agent.agents.instruments import PostgresAgentInstruments 6 | from postgres_da_ai_agent.modules import llm 7 | from postgres_da_ai_agent.modules import rand 8 | from postgres_da_ai_agent.modules import embeddings 9 | import argparse 10 | 11 | DB_URL = os.environ.get("DATABASE_URL") 12 | POSTGRES_TABLE_DEFINITIONS_CAP_REF = "TABLE_DEFINITIONS" 13 | 14 | 15 | custom_function_tool_config = { 16 | "type": "function", 17 | "function": { 18 | "name": "store_fact", 19 | "description": "A function that stores a fact.", 20 | "parameters": { 21 | "type": "object", 22 | "properties": {"fact": {"type": "string"}}, 23 | }, 24 | }, 25 | } 26 | 27 | run_sql_tool_config = { 28 | "type": "function", 29 | "function": { 30 | "name": "run_sql", 31 | "description": "Run a SQL query against the postgres database", 32 | "parameters": { 33 | "type": "object", 34 | "properties": { 35 | "sql": { 36 | "type": "string", 37 | "description": "The SQL query to run", 38 | } 39 | }, 40 | "required": ["sql"], 41 | }, 42 | }, 43 | } 44 | 45 | 46 | def store_fact(fact: str): 47 | print(f"------store_fact({fact})------") 48 | return "Fact stored." 49 | 50 | 51 | def main(): 52 | parser = argparse.ArgumentParser() 53 | parser.add_argument("--prompt", help="The prompt for the AI") 54 | args = parser.parse_args() 55 | 56 | if not args.prompt: 57 | print("Please provide a prompt") 58 | return 59 | 60 | raw_prompt = args.prompt 61 | 62 | prompt = f"Fulfill this database query: {raw_prompt}. " 63 | 64 | assistant_name = "Turbo4" 65 | 66 | assistant = Turbo4() 67 | 68 | session_id = rand.generate_session_id(assistant_name + raw_prompt) 69 | 70 | with PostgresAgentInstruments(DB_URL, session_id) as (agent_instruments, db): 71 | database_embedder = embeddings.DatabaseEmbedder(db) 72 | 73 | table_definitions = database_embedder.get_similar_table_defs_for_prompt( 74 | raw_prompt 75 | ) 76 | 77 | prompt = llm.add_cap_ref( 78 | prompt, 79 | f"Use these {POSTGRES_TABLE_DEFINITIONS_CAP_REF} to satisfy the database query.", 80 | POSTGRES_TABLE_DEFINITIONS_CAP_REF, 81 | table_definitions, 82 | ) 83 | 84 | tools = [ 85 | TurboTool("run_sql", run_sql_tool_config, agent_instruments.run_sql), 86 | ] 87 | 88 | ( 89 | assistant.get_or_create_assistant(assistant_name) 90 | .set_instructions( 91 | "You're an elite SQL developer. You generate the most concise and performant SQL queries." 92 | ) 93 | .equip_tools(tools) 94 | .make_thread() 95 | .add_message(prompt) 96 | .run_thread() 97 | .add_message( 98 | "Use the run_sql function to run the SQL you've just generated.", 99 | ) 100 | .run_thread(toolbox=[tools[0].name]) 101 | .run_validation(agent_instruments.validate_run_sql) 102 | .spy_on_assistant(agent_instruments.make_agent_chat_file(assistant_name)) 103 | .get_costs_and_tokens( 104 | agent_instruments.make_agent_cost_file(assistant_name) 105 | ) 106 | ) 107 | 108 | print(f"✅ Turbo4 Assistant finished.") 109 | 110 | # ---------- Simple Prompt Solution - Same thing, only 2 api calls instead of 8+ ------------ 111 | # sql_response = llm.prompt( 112 | # prompt, 113 | # model="gpt-4-1106-preview", 114 | # instructions="You're an elite SQL developer. You generate the most concise and performant SQL queries.", 115 | # ) 116 | # llm.prompt_func( 117 | # "Use the run_sql function to run the SQL you've just generated: " 118 | # + sql_response, 119 | # model="gpt-4-1106-preview", 120 | # instructions="You're an elite SQL developer. You generate the most concise and performant SQL queries.", 121 | # turbo_tools=tools, 122 | # ) 123 | # agent_instruments.validate_run_sql() 124 | 125 | # ----------- Example use case of Turbo4 and the Assistants API ------------ 126 | 127 | # ( 128 | # assistant.get_or_create_assistant(assistant_name) 129 | # .make_thread() 130 | # .equip_tools(tools) 131 | # .add_message("Generate 10 random facts about LLM technology.") 132 | # .run_thread() 133 | # .add_message("Use the store_fact function to 1 fact.") 134 | # .run_thread(toolbox=["store_fact"]) 135 | # ) 136 | 137 | 138 | if __name__ == "__main__": 139 | main() 140 | -------------------------------------------------------------------------------- /postgres_da_ai_agent/types.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from typing import Callable, List 3 | from dataclasses import dataclass, field 4 | import time 5 | 6 | 7 | @dataclass 8 | class Chat: 9 | from_name: str 10 | to_name: str 11 | message: str 12 | created: int = field(default_factory=time.time) 13 | 14 | 15 | @dataclass 16 | class ConversationResult: 17 | success: bool 18 | messages: List[Chat] 19 | cost: float 20 | tokens: int 21 | last_message_str: str 22 | error_message: str 23 | 24 | 25 | @dataclass 26 | class TurboTool: 27 | name: str 28 | config: dict 29 | function: Callable 30 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "postgres-da-ai-agent" 3 | version = "0.1.0" 4 | description = "" 5 | authors = ["indydevdan "] 6 | readme = "README.md" 7 | packages = [{include = "postgres_da_ai_agent"}] 8 | 9 | [tool.poetry.dependencies] 10 | python = "^3.10" 11 | openai = "^1.2.3" 12 | psycopg2-binary = "^2.9.8" 13 | argparse = "^1.4.0" 14 | python-dotenv = "^1.0.0" 15 | pyautogen = "^0.1.7" 16 | transformers = "^4.34.1" 17 | torch = "^2.1.0" 18 | scikit-learn = "^1.3.1" 19 | tiktoken = "^0.5.1" 20 | guidance = "^0.0.64" 21 | 22 | 23 | [build-system] 24 | requires = ["poetry-core"] 25 | build-backend = "poetry.core.masonry.api" 26 | 27 | [tool.poetry.scripts] 28 | # Not available in current version due to pyautogen not supporting openai ^1.2.3 29 | # old_start = "postgres_da_ai_agent.main:main" 30 | start = "postgres_da_ai_agent.turbo_main:main" 31 | turbo = "postgres_da_ai_agent.turbo_main:main" 32 | --------------------------------------------------------------------------------