Datenbankabfragen in natürlicher Sprache: Text-to-SQL mit BigQuery

Andreas Nigg
Andreas Nigg · Mitgründer, Pondhouse Data
Veröffentlicht Aktualisiert 10 Min. Lesezeit
Datenbankabfragen in natürlicher Sprache: Text-to-SQL mit BigQuery
Inhalt

Mit Ihren Daten in BigQuery chatten können Sie, wenn Sie einen kleinen KI-Agenten zwischen Nutzer und Datenbank setzen. Der Agent liest das Tabellenschema, lässt ein Large Language Model (hier GPT-4o von OpenAI) die Frage in BigQuery-SQL übersetzen, führt die Abfrage aus und schickt einen etwaigen Fehler von BigQuery an das Modell zurück, damit es die Abfrage repariert. Auf dem Demo-Datensatz zur Kundenabwanderung eines Telekomanbieters, den wir unten verwenden, beantwortet diese einfache Schleife Fragen wie "How many churns are in the churn_table?" (1869), ohne dass jemand SQL schreibt.

Large Language Models schreiben sehr gutes SQL für die verschiedensten Datenbanken, auch für Google BigQuery. Das ist nützlich, denn BigQuery-Datensätze sind oft komplex und enthalten einige der wichtigsten Geschäftsdaten eines Unternehmens. Statt zuerst das richtige Datenmodell zu suchen und dann eine Abfrage zu schreiben, fragen Sie in normaler Sprache "Wie viele Nutzer haben sich letzten Monat registriert?" und haben in Sekunden eine Antwort.

In diesem Beitrag zeigen wir, wie Sie eine solche Chat-Oberfläche für Google BigQuery mit GPT-4o bauen. Wir richten die Entwicklungsumgebung ein, verbinden uns mit BigQuery, erzeugen die SQL-Abfragen über die OpenAI-API und führen sie auf Ihrer BigQuery-Datenbank aus.

Als Zugabe zeigen wir, wie Sie mit RAG (Retrieval Augmented Generation) (auf Englisch) die Qualität der erzeugten SQL-Abfragen verbessern. Wenn Sie BigQuery lieber an bestehende KI-Agenten und Chat-Clients anbinden möchten, lesen Sie unsere Anleitung zum MCP-Server für BigQuery.

Update (November 2026): OpenAI hat GPT-4o als veraltet markiert. Laut OpenAIs Deprecations-Seite werden die unten verwendeten gpt-4o-Snapshots am 14. Januar 2027 abgeschaltet. Der Agent funktioniert genauso mit jedem neueren Chat-Modell, das Function Calling unterstützt. Sie müssen nur den Parameter model ändern. Der Embedding-Code verwendet text-embedding-ada-002, das weiterhin funktioniert; OpenAIs aktuelle Embedding-Modelle sind text-embedding-3-small und text-embedding-3-large. Wenn Sie wechseln, prüfen Sie den Schwellenwert in query_chromadb erneut, denn die Abstände ändern sich mit dem Modell. Der ChromaDB-Code verwendet get_or_create_collection, das in allen ChromaDB-Versionen funktioniert (der Rückgabetyp von list_collections() hat sich zwischen den Versionen geändert).

Architektur einer Chat-Oberfläche für BigQuery

Eine Lösung, mit der Sie "mit BigQuery chatten", braucht drei Bausteine:

  1. Eine Möglichkeit, eine Frage in natürlicher Sprache einzugeben.
  2. Eine Schicht, die die Frage in SQL übersetzt.
  3. Eine Laufzeitumgebung, die die SQL-Abfrage auf der BigQuery-Datenbank ausführt.

Baustein 1 und 3 sind unkompliziert. Die eigentliche Logik und Komplexität steckt in Baustein 2, der Übersetzung von natürlicher Sprache in SQL.

Für gute SQL-Generierung braucht es diese Teilschritte:

  1. Die Struktur der Datenbank auslesen, also das Datenbankschema.
  2. Eine Beschreibung der einzelnen Spalten beschaffen, damit die KI den Kontext der Daten versteht.
  3. Mit Retrieval Augmented Generation die BigQuery-Tabellen und -Spalten finden, die für die Frage am relevantesten sind.
  4. Die SQL-Abfrage mit einem Large Language Model erzeugen.
  5. Fehler in der erzeugten Abfrage automatisch korrigieren, denn LLMs sind nicht perfekt und liefern ab und zu fehlerhafte Abfragen.

Zusammengesetzt hat eine robuste Chat-Oberfläche für BigQuery diese Architektur:

Grobarchitektur einer Chat-Oberfläche für BigQuery

Bei genauerem Hinsehen wird klar, dass wir BigQuery nicht einfach mit einem KI-Modell "verbinden". Wir brauchen einen KI-Agenten.

Kurz gesagt sind KI-Agenten KI-Modelle, die eine Eingabe entgegennehmen und dann selbstständig ihr Wissen, externe Werkzeuge (etwa einen Datenbankzugriff) und ihr Gedächtnis nutzen, um die Anfrage zu lösen.

Unser Agent ist klein. Seine Aufgabe ist es, die Frage in eine funktionierende SQL-Abfrage zu übersetzen, die sie beantwortet. Als Werkzeuge bekommt er Zugriff auf die BigQuery-Datenbank und eine Python-Laufzeit. Das nötige Wissen, nämlich SQL zu schreiben, bringt das KI-Modell selbst mit.

KI-Agent, der Nutzereingabe, Datenbankzugriff und eine Python-Laufzeit kombiniert

Hinweis: Da die KI hier nur wenige Aufgaben selbstständig lösen muss, bleibt unser Agent recht klein. KI-Agenten können sehr unterschiedliche Aufgaben übernehmen, hier reicht aber eine sehr begrenzte Umsetzung. Unser Agent muss nur SQL-Abfragen erzeugen, sie ausführen, bei einem Fehler entscheiden, ob er es noch einmal versucht, und das Ergebnis zurückgeben.

Schritt für Schritt: einen Text-to-SQL-Agenten für BigQuery bauen

Mit der Architektur im Kopf gehen wir die Umsetzung durch. Wir bauen die gesamte Agentenschleife von Grund auf, ohne eine der vorhandenen LLM-Bibliotheken. (Einen breiteren Überblick über diesen Ansatz gibt unser Beitrag, wie Sie KI-Agenten ohne Frameworks bauen (auf Englisch).)

Warum nicht LangChain oder LlamaIndex?

Sowohl LangChain als auch LlamaIndex bieten SQL-Agenten als Teil ihres Funktionsumfangs an.

Für den Einstieg in Text-to-SQL-Agenten eignen sich beide gut, ihre Abstraktionen machen aber schwer nachvollziehbar, was darunter passiert. Vor allem sind ihre Umsetzungen recht allgemein gehalten, damit sie viele SQL-Datenbanken abdecken. Nach vielen Stunden und Tagen mit SQL-Agenten haben wir festgestellt, dass jeder Anwendungsfall Anpassungen braucht, sowohl bei der SQL-Generierung als auch bei der Fehlerbehandlung. Diese Flexibilität bieten die High-Level-Schnittstellen von LangChain und LlamaIndex nicht.

Vorbereitung

Stellen Sie zuerst sicher, dass Sie ein aktives Google-Cloud-Projekt haben und für dieses Projekt die Abrechnung aktiviert ist.

Aktivieren Sie dann die folgenden Google-Cloud-APIs:

Richten Sie anschließend die Authentifizierung für BigQuery so ein:

  1. Legen Sie für Ihr Projekt ein Dienstkonto an und geben Sie ihm die folgenden BigQuery-Rollen:

    • BigQuery User
    • BigQuery Data Viewer
    • BigQuery Job User
  2. Laden Sie den Dienstkontoschlüssel des neuen Dienstkontos herunter.

Zum Schluss importieren wir Demodaten in BigQuery. Sie können auch Ihre eigenen Daten verwenden.

  1. Laden Sie den Demo-Datensatz herunter. Er enthält Daten aus einer bekannten Kaggle-Challenge zur Kundenabwanderung (Churn) bei einem Telekomanbieter.

  2. Importieren Sie die Datei in BigQuery, in ein neues Dataset aiagent_test und eine Tabelle churn_table.

Den KI-Agenten umsetzen

  1. Legen Sie eine neue venv oder conda-Umgebung an und aktivieren Sie sie.

  2. Installieren Sie die benötigten Python-Pakete:

    1pip install pandas openai sqlalchemy sqlalchemy-bigquery google-cloud-bigquery chromadb
  3. Importieren Sie die Module, legen Sie BigQuery-Dataset und -Tabelle fest und setzen Sie den OpenAI-API-Schlüssel.

    1from openai.types.chat import ChatCompletionToolParam
    2from sqlalchemy.sql import text
    3from sqlalchemy import create_engine, MetaData
    4import openai
    5from openai import AsyncOpenAI
    6import json
    7import pandas as pd
    8
    9project = "pondhouse-data"
    10dataset = "aiagent_test"
    11table = "churn_table"
    12openai.api_key = "your-openai-api"
  4. Erstellen Sie eine SQLAlchemy-Engine für die BigQuery-Datenbank.

    1engine = create_engine(f'bigquery://{project}', credentials_path='/path/to/keyfile.json')
  5. Testen wir die Datenbankverbindung:

    1statement = text(f"""SELECT count(*) from {dataset}.{table};""")
    2
    3with engine.connect() as connection:
    4 print(connection.execute(statement).fetchone())

    Das Ergebnis sollte so aussehen:

    1(7043,)
  6. Wenn die Verbindung steht, lesen Sie das Datenbankschema aus.

    1metadata = MetaData()
    2metadata.reflect(bind=engine)
    3
    4# Das ausgelesene Schema umformatieren
    5schema_info = []
    6
    7for table in metadata.tables.values():
    8 table_info = {
    9 "table_name": table.name,
    10 "columns": [{column.name: column.type} for column in table.columns],
    11 }
    12 schema_info.append(table_info)

    Damit erhalten Sie eine Liste aller Tabellen und Spalten im angegebenen BigQuery-Dataset.

    1 [{'table_name': 'aiagent_test.churn_table',
    2 'columns': [{'customerID': String()},
    3 {'gender': String()},
    4 {'SeniorCitizen': Integer()},
    5 {'Partner': Boolean()},
    6 {'Dependents': Boolean()},
    7 ...], ...]

Damit ist die Vorbereitung abgeschlossen, und wir können den eigentlichen KI-Agenten umsetzen. Wir erzeugen die SQL-Abfragen über die OpenAI-API und führen sie auf der BigQuery-Datenbank aus.

Die Agentenschleife umsetzen

Im Kern ist ein KI-Agent nichts anderes als eine While-Schleife, die eine Aufgabe oder Frage entgegennimmt und ein KI-Modell aufruft, um sie zu lösen. Die Schleife (der Agent) läuft, bis eine Abbruchbedingung erfüllt ist.

Unser Agent ist einfach gehalten. In Pseudocode sieht er so aus:

while:

  • Mit KI aus der Nutzereingabe und dem Datenbankschema eine SQL-Abfrage erzeugen
  • Die erzeugte SQL-Abfrage auf der BigQuery-Datenbank ausführen
  • Prüfen, ob es Fehler gab. Falls ja, die Abfrage per KI korrigieren
  • Optional: Falls keine Fehler, per KI prüfen, ob das Ergebnis die Frage beantwortet. Falls nicht, wiederholen
  • Aufhören, sobald die Frage beantwortet ist

Setzen wir das in Python um. Zuerst schreiben wir eine Funktion, die die Abfrage erzeugt. Dafür nutzen wir OpenAI Function Calling, weil es verlässliche JSON-Ausgaben liefert, die sich leicht weiterverarbeiten lassen.

Hinweis: In vielen Experimenten wirkten KI-Modelle mit Function Calling etwas weniger intelligent. Bei unserer SQL-Generierung war das bisher kein Problem, Sie sollten es aber im Hinterkopf behalten. Wenn Ihr Modell ständig schlechte Abfragen liefert, probieren Sie den "normalen" Chat-Endpunkt und werten Sie die Ausgabe selbst aus.

1query_database_function: list[ChatCompletionToolParam] = [
2 {
3 "type": "function",
4 "function": {
5 "name": "query_database",
6 "description": "Create a BigQuery SQL query which answers the question as accurately as possible. Make sure to return a complete query, without the user needing to add additional information. \n Also return the type of the query, like 'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'ALTER', 'DROP', etc.",
7 "parameters": {
8 "type": "object",
9 "properties": {
10 "query": {
11 "type": "string",
12 "description": "BigQuery dialect SQL query",
13 },
14 "type": {
15 "type": "string",
16 "description": "Type of query, like 'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'ALTER', 'DROP', etc.",
17 },
18 },
19 "required": ["query", "type"],
20 },
21 },
22 }
23]
24
25async def create_query(query: str, schema):
26 system_prompt = """Given the following database schema, create a BigQuery SQL query which answers the users question as accurately as possible. Make sure to return a complete query, without the user needing to add additional information.
27Also return the type of the query, like 'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'ALTER', 'DROP', etc."""
28
29 client = AsyncOpenAI(api_key=openai.api_key)
30 response = await client.chat.completions.create(
31 model="gpt-4o",
32 temperature=0.5,
33 messages=[
34 {
35 "role": "system",
36 "content": system_prompt + "\n" + "Schema: " + str(schema),
37 },
38 {"role": "user", "content": query},
39 ],
40 tools=query_database_function,
41 )
42
43 query = json.loads(response.choices[0].message.tool_calls[0].function.arguments)[
44 "query"
45 ]
46 query_type = json.loads(
47 response.choices[0].message.tool_calls[0].function.arguments
48 )["type"]
49
50 if query_type != "SELECT":
51 raise Exception("Only read statements are allowed!")
52
53 return query

Diese Funktion übernimmt den ersten Versuch, eine Abfrage zu erzeugen. Die resultierende Abfrage führen wir auf BigQuery aus und prüfen auf Fehler. Gibt es welche, korrigiert die KI die Abfrage.

Dafür schreiben wir als Nächstes die Funktion zur Fehlerbehandlung. Sie ist nichts anderes als eine weitere Abfragegenerierung per LLM, nur dass wir dem Modell diesmal auch die Fehlermeldung mitschicken, damit es die Abfrage korrigieren kann.

1async def error_handling(question: str, query: str, schema, error):
2 system_prompt = "Given the following database schema, you suggested a BigQuery SQL query to answer the following question which resulted in the following error. Please correct the error."
3 client = AsyncOpenAI(api_key=openai.api_key)
4 response = await client.chat.completions.create(
5 model="gpt-4o",
6 temperature=0.5,
7 messages=[
8 {
9 "role": "system",
10 "content": system_prompt
11 + "\n"
12 + "Schema: "
13 + str(schema)
14 + "\n Original question which failed: "
15 + question
16 + "\n Your answer, which failed: "
17 + query
18 + "\n Error: "
19 + str(error),
20 },
21 {
22 "role": "user",
23 "content": "Correct the error and answer the question.",
24 },
25 ],
26 tools=query_database_function,
27 )
28 query = json.loads(response.choices[0].message.tool_calls[0].function.arguments)[
29 "query"
30 ]
31
32 return query

Wir ändern nur den System-Prompt, sodass er den Fehler enthält und das Modell um eine Korrektur bittet. Ansonsten sind Anfrage und Funktion dieselben wie vorher.

Zum Schluss setzen wir alles in der Agentenschleife zusammen:

1async def execute_agent(question: str):
2 query = await create_query(question, schema_info)
3 print(query)
4
5 with engine.connect() as connection:
6 runs = 0
7 while runs < 10:
8 try:
9 # Abfrage ausführen
10 result = connection.execute(text(query))
11 break
12 except Exception as e:
13 print("Oh no, an error. Never mind... We'll figure this out.", e)
14 runs += 1
15 query = await error_handling(question, query, schema_info, e)
16 print("New query: ", query)
17 if runs >= 10:
18 print("Stopping agent after 10 failed attempts.")
19 return None
20 columns = [col for col in result.keys()]
21 data = result.fetchall()
22 return pd.DataFrame(data, columns=columns)

Jetzt starten wir den Agenten mit einer Frage:

1await execute_agent("How many churns are in the churn_table?")
2# Ausgabe: 1869

Der Agent erzeugt die SQL-Abfrage, führt sie aus, behebt Fehler selbst und liefert das richtige Ergebnis.

Jetzt etwas Komplizierteres: der Anteil männlicher und weiblicher Kunden unter denen, die abgewandert sind.

1await execute_agent(
2 "What's the percentage of male and female customers of users who churned?"
3)

Ausgabe:

1SELECT gender, COUNT(*) * 100.0 / (SELECT COUNT(*) FROM aiagent_test.churn_table WHERE Churn = TRUE) AS percentage FROM aiagent_test.churn_table WHERE Churn = TRUE GROUP BY gender
2
3 gender percentage
40 Male 49.75923
51 Female 50.24077

Nicht schlecht für ein paar Dutzend Zeilen Python.

Den BigQuery-SQL-Agenten verbessern

Mit dem Agenten lässt sich schon gut herumspielen, er erzeugt aber auch eine ganze Menge unbrauchbarer Abfragen. Diese Fehler behebt er erstaunlich gut selbst, zwei Maßnahmen verringern aber die Zahl der Fehlversuche.

  1. Beschreibungen zu den Datenbankspalten hinzufügen. Bisher geben wir dem KI-Modell nur den Tabellennamen und die Spaltennamen als Kontext. Das Modell muss dann "raten", welche Spalten für die Frage relevant sind. Das ist fehleranfällig, und Spaltenbeschreibungen verbessern es.

    Dafür gibt es zwei Wege: a. Jede Spalte von Hand beschreiben. Das ist etwas mühsam, aber der genaueste und verlässlichste Weg, dem Modell Kontext zu geben. b. Ein head(10) jeder Spalte an GPT-4o schicken und um eine Beschreibung bitten. Das Modell findet sehr wahrscheinlich eine brauchbare Beschreibung.

    Beide Methoden sind einfach, die Umsetzung überlassen wir deshalb Ihnen. Für diesen Beitrag habe ich unser Schema von GPT-4o so umschreiben lassen:

    1table_with_descriptions = [
    2 {'table_name': 'aiagent_test.churn_table', 'column_name': 'customerID', 'column_description': 'A unique identifier assigned to each customer.', 'column_type': 'String'},
    3 {'table_name': 'aiagent_test.churn_table', 'column_name': 'gender', 'column_description': 'The gender of the customer (Male/Female).', 'column_type': 'String'},
    4 {'table_name': 'aiagent_test.churn_table', 'column_name': 'SeniorCitizen', 'column_description': 'Indicates if the customer is a senior citizen (1 for Yes, 0 for No).', 'column_type': 'Integer'},
    5 {'table_name': 'aiagent_test.churn_table', 'column_name': 'Partner', 'column_description': 'Indicates if the customer has a partner (True for Yes, False for No).', 'column_type': 'Boolean'},
    6 {'table_name': 'aiagent_test.churn_table', 'column_name': 'Dependents', 'column_description': 'Indicates if the customer has dependents (True for Yes, False for No).', 'column_type': 'Boolean'},
    7 {'table_name': 'aiagent_test.churn_table', 'column_name': 'tenure', 'column_description': 'Number of months the customer has stayed with the company.', 'column_type': 'Integer'},
    8 {'table_name': 'aiagent_test.churn_table', 'column_name': 'PhoneService', 'column_description': 'Indicates if the customer has phone service (True for Yes, False for No).', 'column_type': 'Boolean'},
    9 {'table_name': 'aiagent_test.churn_table', 'column_name': 'MultipleLines', 'column_description': 'Indicates if the customer has multiple lines (Yes, No, No phone service).', 'column_type': 'String'},
    10 {'table_name': 'aiagent_test.churn_table', 'column_name': 'InternetService', 'column_description': 'Type of internet service provided to the customer (DSL, Fiber optic, No).', 'column_type': 'String'},
    11 {'table_name': 'aiagent_test.churn_table', 'column_name': 'OnlineSecurity', 'column_description': 'Indicates if the customer has online security service (Yes, No, No internet service).', 'column_type': 'String'},
    12 {'table_name': 'aiagent_test.churn_table', 'column_name': 'OnlineBackup', 'column_description': 'Indicates if the customer has online backup service (Yes, No, No internet service).', 'column_type': 'String'},
    13 {'table_name': 'aiagent_test.churn_table', 'column_name': 'DeviceProtection', 'column_description': 'Indicates if the customer has device protection plan (Yes, No, No internet service).', 'column_type': 'String'},
    14 {'table_name': 'aiagent_test.churn_table', 'column_name': 'TechSupport', 'column_description': 'Indicates if the customer has tech support service (Yes, No, No internet service).', 'column_type': 'String'},
    15 {'table_name': 'aiagent_test.churn_table', 'column_name': 'StreamingTV', 'column_description': 'Indicates if the customer has streaming TV service (Yes, No, No internet service).', 'column_type': 'String'},
    16 {'table_name': 'aiagent_test.churn_table', 'column_name': 'StreamingMovies', 'column_description': 'Indicates if the customer has streaming movies service (Yes, No, No internet service).', 'column_type': 'String'},
    17 {'table_name': 'aiagent_test.churn_table', 'column_name': 'Contract', 'column_description': 'The type of contract the customer has (Month-to-month, One year, Two year).', 'column_type': 'String'},
    18 {'table_name': 'aiagent_test.churn_table', 'column_name': 'PaperlessBilling', 'column_description': 'Indicates if the customer uses paperless billing (True for Yes, False for No).', 'column_type': 'Boolean'},
    19 {'table_name': 'aiagent_test.churn_table', 'column_name': 'PaymentMethod', 'column_description': 'The payment method used by the customer (Electronic check, Mailed check, Bank transfer, Credit card).', 'column_type': 'String'},
    20 {'table_name': 'aiagent_test.churn_table', 'column_name': 'MonthlyCharges', 'column_description': 'The amount charged to the customer on a monthly basis.', 'column_type': 'Float'},
    21 {'table_name': 'aiagent_test.churn_table', 'column_name': 'TotalCharges', 'column_description': 'The total amount charged to the customer.', 'column_type': 'String'},
    22 {'table_name': 'aiagent_test.churn_table', 'column_name': 'Churn', 'column_description': 'Indicates if the customer has churned (True for Yes, False for No).', 'column_type': 'Boolean'}
    23]
  2. Mit RAG (auf Englisch) zuerst die möglicherweise relevanten Tabellen und Spalten finden. Das verkleinert den Suchraum für das KI-Modell.

    Das funktioniert, weil sich die meisten Fragen nur auf einen Teil der vorhandenen Spalten beziehen. Wenn Sie relevante Spalten und Tabellen suchen, bevor die KI eine Abfrage schreibt, kann das die Qualität des erzeugten SQL deutlich verbessern.

Mit RAG die relevanten Spalten für Text-to-SQL finden

Um zu zeigen, wie ein RAG-gestütztes Text-to-SQL-System aussieht, erweitern wir den Agenten, den wir gerade gebaut haben.

Das RAG-Prinzip ist immer gleich: Sie nehmen eine Menge Texte (hier die Spaltenbeschreibungen), erstellen dafür einen Vektorindex und finden über diesen Index die Texte, die der Frage inhaltlich am nächsten sind.

Konkret erstellen wir einen Index über unsere Spaltenbeschreibungen und speichern die übrigen Schlüssel (table_name, column_name und column_type) als Metadaten.

Als Vektordatenbank verwenden wir in diesem Beispiel ChromaDB, jede andere funktioniert aber genauso. (Auch BigQuery selbst kann Embeddings speichern und eine Vektorsuche ausführen, siehe RAG, Embeddings und Vektorsuche mit BigQuery (auf Englisch).)

1
2import chromadb
3from chromadb.config import Settings
4
5async def get_embedding(description: str):
6 client = AsyncOpenAI(api_key=openai.api_key)
7 response = await client.embeddings.create(
8 input=[description], model="text-embedding-ada-002"
9 )
10
11 return response.data[0].embedding
12
13
14# Vektordarstellungen der Beschreibungen erzeugen
15for item in table_with_descriptions:
16 description = item["column_description"]
17 embedding = await get_embedding(description)
18 item["vector"] = embedding
19
20# ChromaDB initialisieren
21client = chromadb.Client(Settings())
22collection_name = "telecom_churn"
23collection = client.get_or_create_collection(collection_name)
24
25# Daten in ChromaDB einfügen
26documents = [item["column_description"] for item in table_with_descriptions]
27embeddings = [item["vector"] for item in table_with_descriptions]
28metadatas = [
29 {k: v for k, v in item.items() if k not in ["vector", "column_description"]}
30 for item in table_with_descriptions
31]
32ids = [f"doc{idx+1}" for idx in range(len(table_with_descriptions))]
33
34collection.add(documents=documents, embeddings=embeddings, metadatas=metadatas, ids=ids)

Hinweis: ChromaDB hat eine praktische Funktion: Sie können collection.add() nur das "Dokument" (den Text, den Sie indexieren wollen) ohne Embeddings übergeben, und ChromaDB erzeugt die Embeddings selbst. In diesem Beispiel erzeugen wir sie von Hand, damit sich der Code leichter auf andere Vektordatenbanken übertragen lässt. Wenn Sie ohnehin ChromaDB verwenden, nutzen Sie die Funktion gern.

Mit den Embeddings schreiben wir eine Funktion, die die relevantesten Spalten heraussucht.

1async def query_chromadb(question: str, threshold=0.5):
2 # Embedding für die Frage erzeugen
3 query_vector = await get_embedding(question)
4
5 # ChromaDB nach den relevantesten Dokumenten abfragen
6 results = collection.query(
7 query_embeddings=[query_vector],
8 n_results=10, # Anzahl der Ergebnisse nach Bedarf anpassen
9 )
10
11 # Relevante Metadaten aus den Ergebnissen holen
12 relevant_results = []
13
14 # Ergebnisse nach Ähnlichkeitsschwelle filtern
15 for i, score in enumerate(results["distances"][0]):
16 if score < threshold:
17 document = results["documents"][0][i]
18 metadata = results["metadatas"][0][i]
19 relevant_results.append(
20 {
21 "table_name": metadata["table_name"],
22 "column_name": metadata["column_name"],
23 "column_type": metadata["column_type"],
24 "column_description": document,
25 "score": score,
26 }
27 )
28
29 relevant_results = sorted(relevant_results, key=lambda x: x["score"], reverse=False)
30 relevant_results
31
32 return relevant_results if relevant_results else []

query_chromadb findet die 10 Spalten, die für die Frage am relevantesten sind. Passen Sie diese Zahl nach Bedarf an.

ChromaDB hat keinen Parameter für einen Schwellenwert, deshalb filtern wir die Ergebnisse selbst und behalten nur Spalten, die einigermaßen relevant sind.

Hinweis: Wenn Ihre Vektordatenbank einen Schwellenwert unterstützt, können Sie das manuelle Filtern weglassen.

Jetzt verwenden wir die Funktion in unserem Agenten:

1async def execute_agent_with_rag(question: str):
2 rag_schema = await query_chromadb(question) # <-- Das ist die einzige nötige Änderung
3 query = await create_query(question, rag_schema)
4 print(query)
5
6 with engine.connect() as connection:
7 runs = 0
8 while runs < 10:
9 try:
10 # Abfrage ausführen
11 result = connection.execute(text(query))
12 break
13 except Exception as e:
14 print("Oh no, an error. Never mind... We'll figure this out.", e)
15 runs += 1
16 query = await error_handling(question, query, schema_info, e)
17 print("New query: ", query)
18 if runs >= 10:
19 print("Stopping agent after 10 failed attempts.")
20 return None
21 columns = [col for col in result.keys()]
22 data = result.fetchall()
23 return pd.DataFrame(data, columns=columns)

Die einzige Änderung ist der Aufruf von query_chromadb vor create_query. Statt des ganzen Tabellenschemas bekommt das Modell nur die Spalten, die für die Frage relevant sein könnten.

Probieren wir es noch einmal:

1await execute_agent_with_rag(
2 "What's the percentage of male and female customers of users who churned?"
3)

Auch mit RAG liefert der Agent das richtige Ergebnis.

Fazit

Wir haben einen einfachen KI-Agenten gebaut, der Fragen in natürlicher Sprache in SQL übersetzt, die Abfragen auf einer BigQuery-Datenbank ausführt und Fehler automatisch behandelt. Er nutzt GPT-4o von OpenAI für die SQL-Generierung, Google BigQuery für Speicherung und Abfragen und Retrieval Augmented Generation (RAG), um die relevanten Spalten auszuwählen.

Zwei Maßnahmen verbessern die Genauigkeit des erzeugten SQL:

  • Beschreibungen zu den Datenbankspalten hinzufügen
  • Vor der Abfragegenerierung mit RAG die relevanten Tabellen und Spalten finden

Ein weiterer Schritt ist eine semantische Schicht, die Geschäftsbegriffe auf Tabellen, Join-Pfade und SQL-Ausdrücke abbildet. Ein Werkzeug, das eine solche Schicht automatisch erzeugt, haben wir in Text-to-SQL mit semantischer Schicht: NVIDIA Auto Ontology im Praxistest ausprobiert.

Weil die Umsetzung selbst gebaut ist, können Sie Prompts und Fehlerbehandlung auf BigQuery und Ihren konkreten Anwendungsfall abstimmen. Das ist der Hauptgrund, warum wir weder LangChain noch LlamaIndex verwendet haben.

Wir glauben, dass solche Schnittstellen in natürlicher Sprache verändern können, wie Unternehmen mit ihren Daten arbeiten. Wer kein SQL kann, bekommt Antworten aus komplexen Datenbanken, ohne auf das Datenteam angewiesen zu sein, und Datenfachleute verbringen mehr Zeit mit Analyse und weniger mit dem Heraussuchen von Daten.

Unfehlbar ist das System allerdings nicht. Wichtige Ergebnisse sollten Sie immer prüfen, besonders bei komplexen Abfragen oder wenn Sie wichtige Geschäftsentscheidungen darauf stützen.

Wenn Sie ein solches System in Ihrem Unternehmen umsetzen wollen, unterstützen wir Sie gern. Melden Sie sich bei uns.

Weiterlesen

Sie möchten zuverlässige KI-Agenten-Systeme bauen?

Wir haben auf Basis aktueller Forschung einen ausführlichen Leitfaden (auf Englisch) erstellt, wie Sie robuste KI-Agenten-Systeme bauen, die auch im produktiven Betrieb funktionieren. Der Leitfaden behandelt:

  • die 14 systematischen Fehlerquellen in Multi-Agenten-Systemen
  • belegte Best Practices für das Design von Agenten
  • strukturierte Kommunikationsprotokolle und Prüfmechanismen
Andreas Nigg

Über den Autor

Andreas Nigg

Mitgründer, Pondhouse Data

Andreas entwickelt seit mehr als acht Jahren Datenanwendungen. Bei Pondhouse Data konzipiert und baut er KI-Systeme und integriert sie in komplexe, gewachsene Unternehmensumgebungen.