pgvector-java/examples/cohere/src/main/java/com/example/Example.java at master · pgvector/pgvector-java · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
packagecom.example;

importjava.io.IOException;
importjava.net.URI;
importjava.net.http.HttpClient;
importjava.net.http.HttpRequest;
importjava.net.http.HttpRequest.BodyPublishers;
importjava.net.http.HttpResponse;
importjava.net.http.HttpResponse.BodyHandlers;
importjava.sql.Connection;
importjava.sql.DriverManager;
importjava.sql.PreparedStatement;
importjava.sql.ResultSet;
importjava.sql.SQLException;
importjava.sql.Statement;
importjava.util.ArrayList;
importjava.util.List;
importcom.fasterxml.jackson.databind.ObjectMapper;
importcom.fasterxml.jackson.databind.JsonNode;
importcom.fasterxml.jackson.databind.node.ObjectNode;
importcom.pgvector.PGbit;

publicclassExample {
publicstaticvoidmain(String[] args) throwsIOException, InterruptedException, SQLException {
StringapiKey = System.getenv("CO_API_KEY");
if (apiKey == null) {
System.out.println("Set CO_API_KEY");
System.exit(1);
 }

Connectionconn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_example");

StatementsetupStmt = conn.createStatement();
setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
setupStmt.executeUpdate("DROP TABLE IF EXISTS documents");

StatementcreateStmt = conn.createStatement();
createStmt.executeUpdate("CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1536))");

String[] input = {
"The dog is barking",
"The cat is purring",
"The bear is growling"
 };
List<byte[]> embeddings = embed(input, "search_document", apiKey);

for (inti = 0; i < input.length; i++) {
PreparedStatementinsertStmt = conn.prepareStatement("INSERT INTO documents (content, embedding) VALUES (?, ?)");
insertStmt.setString(1, input[i]);
insertStmt.setObject(2, newPGbit(embeddings.get(i)));
insertStmt.executeUpdate();
 }

Stringquery = "forest";
byte[] queryEmbedding = embed(newString[] {query}, "search_query", apiKey).get(0);
PreparedStatementneighborStmt = conn.prepareStatement("SELECT * FROM documents ORDER BY embedding <~> ? LIMIT 5");
neighborStmt.setObject(1, newPGbit(queryEmbedding));
ResultSetrs = neighborStmt.executeQuery();
while (rs.next()) {
System.out.println(rs.getString("content"));
 }

conn.close();
 }

// https://docs.cohere.com/reference/embed
privatestaticList<byte[]> embed(String[] texts, StringinputType, StringapiKey) throwsIOException, InterruptedException {
ObjectMappermapper = newObjectMapper();
ObjectNoderoot = mapper.createObjectNode();
for (Stringv : texts) {
root.withArray("texts").add(v);
 }
root.put("model", "embed-v4.0");
root.put("input_type", inputType);
root.withArray("embedding_types").add("binary");
Stringjson = mapper.writeValueAsString(root);

HttpClientclient = HttpClient.newHttpClient();
HttpRequestrequest = HttpRequest.newBuilder()
 .uri(URI.create("https://api.cohere.com/v2/embed"))
 .header("Authorization", "Bearer " + apiKey)
 .header("Content-Type", "application/json")
 .POST(BodyPublishers.ofString(json))
 .build();
HttpResponse<String> response = client.send(request, BodyHandlers.ofString());

List<byte[]> embeddings = newArrayList<>();
for (JsonNoden : mapper.readTree(response.body()).get("embeddings").get("binary")) {
byte[] embedding = newbyte[n.size()];
inti = 0;
for (JsonNodev : n) {
embedding[i++] = (byte) v.asInt();
 }
embeddings.add(embedding);
 }
returnembeddings;
 }
}