Skip to content

Commit

Permalink
Merge pull request #220 from 0xThresh/text_to_sql_1.1
Browse files Browse the repository at this point in the history
Made env vars more generic, added Ollama port to dev-docker.sh
  • Loading branch information
0xThresh authored Aug 12, 2024
2 parents 45e3a21 + f45d666 commit 231e032
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 13 deletions.
2 changes: 1 addition & 1 deletion dev-docker.sh
Original file line number Diff line number Diff line change
Expand Up @@ -6,4 +6,4 @@

# Runs the containers with Ollama image for Open WebUI and the Pipelines endpoint in place
docker run -d -p 9099:9099 --add-host=host.docker.internal:host-gateway -v pipelines:/app/pipelines --name pipelines --restart always --env-file .env ghcr.io/open-webui/pipelines:latest
docker run -d -p 3000:8080 -v ~/.ollama:/root/.ollama -v open-webui:/app/backend/data --name open-webui --restart always -e OPENAI_API_BASE_URL=http://host.docker.internal:9099 -e OPENAI_API_KEY=0p3n-w3bu! ghcr.io/open-webui/open-webui:ollama
docker run -d -p 3000:8080 -p 11434:11434 --add-host=host.docker.internal:host-gateway -v ~/.ollama:/root/.ollama -v open-webui:/app/backend/data --name open-webui --restart always -e OPENAI_API_BASE_URL=http://host.docker.internal:9099 -e OPENAI_API_KEY=0p3n-w3bu! -e OLLAMA_HOST=0.0.0.0 ghcr.io/open-webui/open-webui:ollama
24 changes: 12 additions & 12 deletions examples/pipelines/rag/text_to_sql_pipeline.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
"""
title: Llama Index DB Pipeline
author: 0xThresh
date: 2024-07-01
version: 1.0
date: 2024-08-11
version: 1.1
license: MIT
description: A pipeline for using text-to-SQL for retrieving relevant information from a database using the Llama Index library.
requirements: llama_index, sqlalchemy, psycopg2-binary
Expand All @@ -24,7 +24,7 @@ class Valves(BaseModel):
DB_USER: str
DB_PASSWORD: str
DB_DATABASE: str
DB_TABLES: list[str]
DB_TABLE: str
OLLAMA_HOST: str
TEXT_TO_SQL_MODEL: str

Expand All @@ -39,14 +39,14 @@ def __init__(self):
self.valves = self.Valves(
**{
"pipelines": ["*"], # Connect to all pipelines
"DB_HOST": os.getenv("PG_HOST", "http://localhost:5432"), # Database hostname
"DB_PORT": os.getenv("PG_PORT", 5432), # Database port
"DB_USER": os.getenv("PG_USER", "postgres"), # User to connect to the database with
"DB_PASSWORD": os.getenv("PG_PASSWORD", "password"), # Password to connect to the database with
"DB_DATABASE": os.getenv("PG_DB", "postgres"), # Database to select on the DB instance
"DB_TABLES": ["albums"], # Table(s) to run queries against
"DB_HOST": os.getenv("DB_HOST", "http://localhost"), # Database hostname
"DB_PORT": os.getenv("DB_PORT", 5432), # Database port
"DB_USER": os.getenv("DB_USER", "postgres"), # User to connect to the database with
"DB_PASSWORD": os.getenv("DB_PASSWORD", "password"), # Password to connect to the database with
"DB_DATABASE": os.getenv("DB_DATABASE", "postgres"), # Database to select on the DB instance
"DB_TABLE": os.getenv("DB_TABLE", "table_name"), # Table(s) to run queries against
"OLLAMA_HOST": os.getenv("OLLAMA_HOST", "http://host.docker.internal:11434"), # Make sure to update with the URL of your Ollama host, such as http://localhost:11434 or remote server address
"TEXT_TO_SQL_MODEL": "phi3:latest" # Model to use for text-to-SQL generation
"TEXT_TO_SQL_MODEL": os.getenv("TEXT_TO_SQL_MODEL", "llama3.1:latest") # Model to use for text-to-SQL generation
}
)

Expand All @@ -69,7 +69,7 @@ def pipe(
# Debug logging is required to see what SQL query is generated by the LlamaIndex library; enable on Pipelines server if needed

# Create database reader for Postgres
sql_database = SQLDatabase(self.engine, include_tables=self.valves.DB_TABLES)
sql_database = SQLDatabase(self.engine, include_tables=[self.valves.DB_TABLE])

# Set up LLM connection; uses phi3 model with 128k context limit since some queries have returned 20k+ tokens
llm = Ollama(model=self.valves.TEXT_TO_SQL_MODEL, base_url=self.valves.OLLAMA_HOST, request_timeout=180.0, context_window=30000)
Expand Down Expand Up @@ -99,7 +99,7 @@ def pipe(

query_engine = NLSQLTableQueryEngine(
sql_database=sql_database,
tables=self.valves.DB_TABLES,
tables=[self.valves.DB_TABLE],
llm=llm,
embed_model="local",
text_to_sql_prompt=text_to_sql_template,
Expand Down

0 comments on commit 231e032

Please sign in to comment.