From f60a6e76d9554eb0600d54f52aac1418fb476913 Mon Sep 17 00:00:00 2001
From: Josh Salomon <jsalomon@redhat.com>
Date: Wed, 9 Apr 2025 21:46:29 +0300
Subject: [PATCH] Add documentation for postgresql uri

and the Generated documentation fixes

Signed-off-by: Josh Salomon <jsalomon@redhat.com>

Signed-off-by: Josh Salomon <jsalomon@redhat.com>
---
 docs/_static/llama-stack-spec.html    | 2 +-
 docs/_static/llama-stack-spec.yaml    | 1 +
 llama_stack/apis/datasets/datasets.py | 5 ++++-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html
index 4990d845e..3b144487b 100644
--- a/docs/_static/llama-stack-spec.html
+++ b/docs/_static/llama-stack-spec.html
@@ -9399,7 +9399,7 @@
                     },
                     "source": {
                         "$ref": "#/components/schemas/DataSource",
-                        "description": "The data source of the dataset. Ensure that the data source schema is compatible with the purpose of the dataset. Examples: - { \"type\": \"uri\", \"uri\": \"https://mywebsite.com/mydata.jsonl\" } - { \"type\": \"uri\", \"uri\": \"lsfs://mydata.jsonl\" } - { \"type\": \"uri\", \"uri\": \"data:csv;base64,{base64_content}\" } - { \"type\": \"uri\", \"uri\": \"huggingface://llamastack/simpleqa?split=train\" } - { \"type\": \"rows\", \"rows\": [ { \"messages\": [ {\"role\": \"user\", \"content\": \"Hello, world!\"}, {\"role\": \"assistant\", \"content\": \"Hello, world!\"}, ] } ] }"
+                        "description": "The data source of the dataset. Ensure that the data source schema is compatible with the purpose of the dataset. Examples: - { \"type\": \"uri\", \"uri\": \"https://mywebsite.com/mydata.jsonl\" } - { \"type\": \"uri\", \"uri\": \"lsfs://mydata.jsonl\" } - { \"type\": \"uri\", \"uri\": \"data:csv;base64,{base64_content}\" } - { \"type\": \"uri\", \"uri\": \"huggingface://llamastack/simpleqa?split=train\" } - { \"type\": \"uri\", \"uri\": \"postgresql://<user>:<password>@<host>:<port>/<database>?table=<tablename>\" } - { \"type\": \"rows\", \"rows\": [ { \"messages\": [ {\"role\": \"user\", \"content\": \"Hello, world!\"}, {\"role\": \"assistant\", \"content\": \"Hello, world!\"}, ] } ] }"
                     },
                     "metadata": {
                         "type": "object",
diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml
index ba3868560..e08ba8af9 100644
--- a/docs/_static/llama-stack-spec.yaml
+++ b/docs/_static/llama-stack-spec.yaml
@@ -6420,6 +6420,7 @@ components:
             "uri": "https://mywebsite.com/mydata.jsonl" } - { "type": "uri", "uri":
             "lsfs://mydata.jsonl" } - { "type": "uri", "uri": "data:csv;base64,{base64_content}"
             } - { "type": "uri", "uri": "huggingface://llamastack/simpleqa?split=train"
+            } - { "type": "uri", "uri": "postgresql://<user>:<password>@<host>:<port>/<database>?table=<tablename>"
             } - { "type": "rows", "rows": [ { "messages": [ {"role": "user", "content":
             "Hello, world!"}, {"role": "assistant", "content": "Hello, world!"}, ]
             } ] }
diff --git a/llama_stack/apis/datasets/datasets.py b/llama_stack/apis/datasets/datasets.py
index 8ff3e7d15..80b3762b1 100644
--- a/llama_stack/apis/datasets/datasets.py
+++ b/llama_stack/apis/datasets/datasets.py
@@ -176,7 +176,10 @@ class Datasets(Protocol):
                "type": "uri",
                "uri": "huggingface://llamastack/simpleqa?split=train"
            }
-           TODO: Add postgresql example here
+           - {
+               "type": "uri",
+               "uri": "postgresql://<user>:<password>@<host>:<port>/<database>?table=<tablename>"
+           }
            - {
                "type": "rows",
                "rows": [