Added Cohere example [skip ci]

ankane · ankane · commit b49c0ec35405 · 2025-01-13T01:44:22.000-08:00
diff --git a/README.md b/README.md
@@ -25,6 +25,7 @@ And follow the instructions for your database library:
 Or check out some examples:
 
 - [Embeddings](examples/openai/example.cpp) with OpenAI
+- [Binary embeddings](examples/cohere/example.cpp) with Cohere
 - [Recommendations](examples/disco/example.cpp) with Disco
 
 ## libpqxx
diff --git a/examples/cohere/CMakeLists.txt b/examples/cohere/CMakeLists.txt
@@ -0,0 +1,17 @@
+cmake_minimum_required(VERSION 3.18)
+
+project(example)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_FLAGS "-Wno-unknown-attributes")
+
+include(FetchContent)
+
+FetchContent_Declare(cpr GIT_REPOSITORY https://github.com/libcpr/cpr.git GIT_TAG 1.11.1)
+FetchContent_Declare(json GIT_REPOSITORY https://github.com/nlohmann/json.git GIT_TAG v3.11.3)
+FetchContent_Declare(libpqxx GIT_REPOSITORY https://github.com/jtv/libpqxx.git GIT_TAG 7.10.0)
+FetchContent_MakeAvailable(cpr json libpqxx)
+
+add_executable(example example.cpp)
+target_include_directories(example PRIVATE ${CMAKE_SOURCE_DIR}/../../include)
+target_link_libraries(example PRIVATE cpr::cpr nlohmann_json::nlohmann_json pqxx)
diff --git a/examples/cohere/example.cpp b/examples/cohere/example.cpp
@@ -0,0 +1,77 @@
+#include <bitset>
+#include <cstdint>
+#include <iostream>
+
+#include <cpr/cpr.h>
+#include <nlohmann/json.hpp>
+#include <pgvector/pqxx.hpp>
+#include <pqxx/pqxx>
+
+using json = nlohmann::json;
+
+// https://docs.cohere.com/reference/embed
+std::vector<std::string> fetch_embeddings(const std::vector<std::string>& texts, const std::string& input_type, char *api_key) {
+    std::string url = "https://api.cohere.com/v1/embed";
+    json data = {
+        {"texts", texts},
+        {"model", "embed-english-v3.0"},
+        {"input_type", input_type},
+        {"embedding_types", {"ubinary"}}
+    };
+
+    cpr::Response r = cpr::Post(
+        cpr::Url{url},
+        cpr::Body{data.dump()},
+        cpr::Bearer{api_key},
+        cpr::Header{{"Content-Type", "application/json"}}
+    );
+    json response = json::parse(r.text);
+
+    std::vector<std::string> embeddings;
+    for (auto& v: response["embeddings"]["ubinary"]) {
+        std::stringstream buf;
+        for (uint8_t c : v) {
+            std::bitset<8> b{c};
+            buf << b.to_string();
+        }
+        embeddings.emplace_back(buf.str());
+    }
+    return embeddings;
+}
+
+int main() {
+    char *api_key = std::getenv("CO_API_KEY");
+    if (!api_key) {
+        std::cout << "Set CO_API_KEY" << std::endl;
+        return 1;
+    }
+
+    pqxx::connection conn("dbname=pgvector_example");
+
+    pqxx::work tx(conn);
+    tx.exec("CREATE EXTENSION IF NOT EXISTS vector");
+    tx.exec("DROP TABLE IF EXISTS documents");
+    tx.exec("CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))");
+    tx.commit();
+
+    std::vector<std::string> input = {
+        "The dog is barking",
+        "The cat is purring",
+        "The bear is growling"
+    };
+    auto embeddings = fetch_embeddings(input, "search_document", api_key);
+
+    for (size_t i = 0; i < input.size(); i++) {
+        tx.exec("INSERT INTO documents (content, embedding) VALUES ($1, $2)", {input[i], embeddings[i]});
+    }
+    tx.commit();
+
+    std::string query = "forest";
+    auto query_embedding = fetch_embeddings({query}, "search_query", api_key)[0];
+    pqxx::result result = tx.exec("SELECT content FROM documents ORDER BY embedding <~> $1 LIMIT 5", pqxx::params{query_embedding});
+    for (auto const& row : result) {
+        std::cout << row[0].c_str() << std::endl;
+    }
+
+    return 0;
+}
