Skip to content

Commit

Permalink
docs(examples): add ollama example (#289)
Browse files Browse the repository at this point in the history
* docs(examples): add ollama example

* remove unncessary pkg details

* return score

* docs: update readme
  • Loading branch information
notrab authored Jan 7, 2025
1 parent b41fa4b commit fbf8f63
Show file tree
Hide file tree
Showing 6 changed files with 515 additions and 0 deletions.
1 change: 1 addition & 0 deletions examples/ollama/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
local.db
29 changes: 29 additions & 0 deletions examples/ollama/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
# Ollama + Vector Search Example

This example demonstrates how to use libSQL vector search with a local database and Ollama.

## Install Dependencies

```bash
npm i
```

## Install Ollama

[Download Ollama](https://ollama.com/download) and install it.

## Running

Make sure Ollama is running with the model `mistral`:

```bash
ollama run mistral
```

Execute the example:

```bash
node index.mjs
```

This will setup a local SQLite database, generate embeddings using Ollama, and insert the data with embeddings, and then query the results using the vector similarity search function.
110 changes: 110 additions & 0 deletions examples/ollama/index.mjs
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
import { createClient } from "@libsql/client";
import ollama from "ollama";

const client = createClient({
url: "file:local.db",
});

await client.batch(
[
"CREATE TABLE IF NOT EXISTS movies (id INTEGER PRIMARY KEY, title TEXT NOT NULL, description TEXT NOT NULL, embedding F32_BLOB(4096))",
"CREATE INDEX IF NOT EXISTS movies_embedding_idx ON movies(libsql_vector_idx(embedding))",
],
"write",
);

async function getEmbedding(prompt) {
const response = await ollama.embeddings({
model: "mistral",
prompt,
});

return response.embedding;
}

async function insertMovie(id, title, description) {
const embedding = await getEmbedding(description);

await client.execute({
sql: `INSERT OR REPLACE INTO movies (id, title, description, embedding) VALUES (?, ?, ?, vector(?))`,
args: [id, title, description, JSON.stringify(embedding)],
});
}

async function insertMovieIfNotExists(id, title, description) {
const existing = await client.execute({
sql: "SELECT id FROM movies WHERE id = ?",
args: [id],
});

if (existing.rows.length === 0) {
await insertMovie(id, title, description);
console.log(`Inserted: ${title} (ID: ${id})`);
} else {
console.log(`Movie already exists: ${title} (ID: ${id})`);
}
}

async function findSimilarMovies(description, limit = 3) {
const queryEmbedding = await getEmbedding(description);

const results = await client.execute({
sql: `
WITH vector_scores AS (
SELECT DISTINCT
id,
title,
description,
1 - vector_distance_cos(embedding, vector32(?)) AS similarity
FROM movies
ORDER BY similarity DESC
LIMIT ?
)
SELECT id, title, description, similarity FROM vector_scores
`,
args: [JSON.stringify(queryEmbedding), limit],
});

return results.rows;
}

try {
const sampleMovies = [
{
id: 1,
title: "Inception",
description:
"A thief who enters the dreams of others to steal secrets from their subconscious.",
},
{
id: 2,
title: "The Matrix",
description:
"A computer programmer discovers that reality as he knows it is a simulation created by machines.",
},
{
id: 3,
title: "Interstellar",
description:
"Astronauts travel through a wormhole in search of a new habitable planet for humanity.",
},
];

for (const movie of sampleMovies) {
await insertMovieIfNotExists(movie.id, movie.title, movie.description);
}

const query =
"A sci-fi movie about virtual reality and artificial intelligence";
console.log("\nSearching for movies similar to:", query);

const similarMovies = await findSimilarMovies(query);
console.log("\nSimilar movies found:");
similarMovies.forEach((movie) => {
console.log(`\nTitle: ${movie.title}`);
console.log(`Description: ${movie.description}`);
console.log(`Similarity: ${movie.similarity.toFixed(4)}`);
});
} catch (error) {
console.error("Error:", error);
}
Loading

0 comments on commit fbf8f63

Please sign in to comment.