Add docs for LLM Extension

astro.config.mjs

@@ -287,6 +287,7 @@ export default defineConfig({
                         { label: 'JSON', link: '/extensions/json' },
                         { label: 'Neo4j', link: '/extensions/neo4j', badge: { text: 'New' }},
                         { label: 'Vector search', link: '/extensions/vector'},
+                        { label: 'LLM', link: '/extensions/llm', badge: { text: 'New' }},
                     ],
                 },
             ],

src/content/docs/extensions/llm.mdx

@@ -0,0 +1,239 @@
+---
+title: "LLM Extension"
+---
+
+The **LLM extension** provides an interface to fetch text embeddings from supported providers and their models.
+
+It can be installed and loaded using the following commands:
+
+```sql
+INSTALL llm;
+LOAD llm;
+```
+
+The list of supported providers and tested models is as follows.
+
+| Provider | Models |
+| --- | --- |
+| `amazon-bedrock` | amazon.titan-embed-text-v1 |
+| `google-vertex` | gemini-embedding-001, text-embedding-005, text-multilingual-embedding-002 |
+| `google-gemini` | gemini-embedding-exp-03-07, text-embedding-004, embedding-001 |
+| `ollama` | nomic-embed-text, all-minilm:l6-v2 |
+| `open-ai` | text-embedding-3-large, text-embedding-3-small, text-embedding-ada-002 |
+| `voyage-ai` | voyage-3-large, voyage-3.5, voyage-3.5-lite, voyage-code-3, voyage-finance-2, voyage-law-2, voyage-code-2 |
+
+Note: Model names are dynamically validated at runtime. If an unsupported model is specified, the provider's API will return an appropriate error.
+
+---
+
+## Usage
+
+### Configuration
+
+The LLM extension supports configurable options such as credentials, dimensions, and region where applicable.
+
+- All credentials are provided through environment variables.
+- Dimensions and region are passed to the `CREATE_EMBEDDING` function call, when supported.
+
+The following environment variables must be set for the corresponding providers:
+
+| Provider | Required Environment Variables |
+| --- | --- |
+| `amazon-bedrock` | `AWS_ACCESS_KEY`, `AWS_SECRET_ACCESS_KEY` |
+| `google-vertex` | `GOOGLE_CLOUD_PROJECT_ID`, `GOOGLE_VERTEX_ACCESS_KEY` |
+| `google-gemini` | `GOOGLE_GEMINI_API_KEY` |
+| `open-ai` | `OPENAI_API_KEY` |
+| `voyage-ai` | `VOYAGE_API_KEY` |
+
+Note that the `ollama` provider requires no environment variables but does require the embedding model to be available at `http://localhost:11434`.
+
+---
+
+## Embedding Function
+
+The core feature of the LLM extension is a scalar function that generates text embeddings.
+
+### Syntax
+
+```sql
+CREATE_EMBEDDING(prompt: STRING, provider: STRING, model: STRING) -> LIST<FLOAT>
+```
+
+### Overloads
+
+```sql
+CREATE_EMBEDDING(prompt: STRING, provider: STRING, model: STRING, dimensions: INT64) -> LIST<FLOAT>
+```
+
+```sql
+CREATE_EMBEDDING(prompt: STRING, provider: STRING, model: STRING, region: STRING) -> LIST<FLOAT>
+```
+
+```sql
+CREATE_EMBEDDING(prompt: STRING, provider: STRING, model: STRING, dimensions: INT64, region: STRING) -> LIST<FLOAT>
+```
+
+### Parameters
+
+- **prompt**: The input text to embed.
+- **provider**: The embedding provider, e.g., `'open-ai'`, `'ollama'`, `'voyage-ai'`, etc.
+- **model**: The identifier of the embedding model. This may change overtime; reference the provider's API reference before use.
+- **dimensions** *(optional)*: The size of the embedding vector to be returned, if supported by the provider. Must be a positive integer.
+- **region** *(optional)*: A region or endpoint hint, if supported by the provider.
+
+### Examples
+
+```sql
+CREATE_EMBEDDING("Hello world", "openai", "text-embedding-3-small")
+```
+
+With optional parameters:
+
+```sql
+CREATE_EMBEDDING("Hello world", "openai", "text-embedding-3-small", 512)
+```
+
+```sql
+CREATE_EMBEDDING("Hello world", "amazon-bedrock", "amazon.titan-embed-text-v1", "us-east-1");
+```
+
+```sql
+CREATE_EMBEDDING("Hello world", "google-vertex", "gemini-embedding-001", 256, "us-east1");
+```
+
+---
+
+## Providers and Configuration
+
+### Region Support
+
+The following providers support region configuration. The list of regions may not be exhaustive. Please reference the provider's API reference for more information about what regions are supported.
+There is no default region; a region argument must be passed to the function call to use these providers.
+
+| Provider | Regions |
+| --- | --- |
+| `amazon-bedrock` | us-east-1, us-west-2, ap-southeast-1, ap-northeast-1, eu-central-1 |
+| `google-vertex` | us-central1, us-west1, us-west2, us-west3, us-west4, us-east1, us-east4, us-south1, northamerica-northeast1, northamerica-northeast2, southamerica-east1, southamerica-west1, europe-west2, europe-west1, europe-west4, europe-west6, europe-west3, europe-north1, europe-central2, europe-west8, europe-west9, europe-southwest1, asia-south1, asia-southeast1, asia-southeast2, asia-east2, asia-east1, asia-northeast1, asia-northeast2, australia-southeast1, australia-southeast2, asia-northeast3, me-west1 |
+
+Note: Regions are dynamically validated at runtime. If an unsupported region is specified, the provider's API will return an error.
+
+### Dimensions Support
+
+The following providers support dimension configuration. 
+
+| Provider | Models |
+| --- | --- |
+| `google-vertex` | gemini-embedding-001, text-embedding-005, text-multilingual-embedding-002 |
+| `open-ai` | text-embedding-3-large, text-embedding-3-small |
+| `voyage-ai` | voyage-3-large, voyage-3.5, voyage-3.5-lite, voyage-code-3 |
+
+Note: The list of models that support dimension configuration may change
+overtime. The set of valid dimensions is determined by the provider and may also
+change over time. It is highly recommended to check the provider's API reference
+before use. If an unsupported dimension is specified, the provider's API will return an error.
+
+---
+
+## Synergy with Vector Extension
+The following example uses the same dataset as the example in the Vector
+documentation. Here we use the embedding function to fetch embeddings.
+
+```python
+# create_embeddings.py
+import pyarrow
+import polars
+import kuzu
+
+DB_NAME = ""
+
+# Initialize the database
+db = kuzu.Database(DB_NAME)
+conn = kuzu.Connection(db)
+
+# Install and load vector extension
+conn.execute("INSTALL vector; LOAD vector;")
+conn.execute("INSTALL llm; LOAD llm;")
+
+
+# Create tables
+conn.execute("CREATE NODE TABLE Book(id SERIAL PRIMARY KEY, title STRING, title_embedding FLOAT[1536], published_year INT64);")
+conn.execute("CREATE NODE TABLE Publisher(name STRING PRIMARY KEY);")
+conn.execute("CREATE REL TABLE PublishedBy(FROM Book TO Publisher);")
+
+# Sample data
+titles = [
+    "The Quantum World",
+    "Chronicles of the Universe",
+    "Learning Machines",
+    "Echoes of the Past",
+    "The Dragon's Call"
+]
+publishers = ["Harvard University Press", "Independent Publisher", "Pearson", "McGraw-Hill Ryerson", "O'Reilly"]
+published_years = [2004, 2022, 2019, 2010, 2015]
+
+# Insert sample data - Books with embeddings
+for title, published_year in zip(titles, published_years):
+    conn.execute(
+        """CREATE (b:Book {title: $title, title_embedding: create_embedding($title, 'open-ai', 'text-embedding-3-small'), published_year: $year});""",
+        {"title": title, "year": published_year}
+    )
+    print(f"Inserted book: {title}")
+
+# Insert sample data - Publishers
+for publisher in publishers:
+    conn.execute(
+        """CREATE (p:Publisher {name: $publisher});""",
+        {"publisher": publisher}
+    )
+    print(f"Inserted publisher: {publisher}")
+
+# Create relationships between Books and Publishers
+for title, publisher in zip(titles, publishers):
+    conn.execute("""
+        MATCH (b:Book {title: $title})
+        MATCH (p:Publisher {name: $publisher})
+        CREATE (b)-[:PublishedBy]->(p);
+        """,
+            {"title": title, "publisher": publisher}
+    )
+    print(f"Created relationship between {title} and {publisher}")
+
+
+
+conn.execute(
+            """
+            CALL CREATE_VECTOR_INDEX(
+                'Book',
+                'title_vec_index',
+                'title_embedding'
+                )
+            """
+            )
+
+result = conn.execute(
+    """
+    CALL QUERY_VECTOR_INDEX(
+        'Book',
+        'title_vec_index',
+        create_embedding('quantum machine learning', 'open-ai', 'text-embedding-3-small'),
+        2
+    )
+    RETURN node.title ORDER BY distance;
+    """
+)
+print(result.get_as_pl())
+```
+
+In the above query, we asked for the 2 nearest neighbors of the query vector "quantum machine learning".
+The result is a list of book titles that are most similar to this concept.
+
+```
+┌───────────────────┐
+│ node.title        │
+│ ---               │
+│ str               │
+╞═══════════════════╡
+│ Learning Machines │
+│ The Quantum World │
+└───────────────────┘
+```