Skip to content

Commit d826861

Browse files
michelerivamicheleriva
authored
feat: adds support for OramaCore to Orama Switch (#914)
1 parent b9cdce3 commit d826861

File tree

5 files changed

+308
-125
lines changed

5 files changed

+308
-125
lines changed

‎packages/switch/package.json

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
{
22
"name": "@orama/switch",
33
"version": "3.1.3",
4-
"description": "Orama Switch allows you to run queries on Orama Cloud and OSS with a single interface",
4+
"description": "Orama Switch allows you to run queries on Orama Cloud, OramaCore and Orama JS with a single interface",
55
"main": "./dist/index.js",
66
"types": "./dist/index.d.ts",
77
"exports": {
@@ -32,7 +32,8 @@
3232
},
3333
"peerDependencies": {
3434
"@orama/orama": "workspace:*",
35-
"@oramacloud/client": "^2.1.1"
35+
"@oramacloud/client": "^2.1.1",
36+
"@orama/core": "^0.0.10"
3637
},
3738
"license": "Apache-2.0",
3839
"devDependencies": {

‎packages/switch/src/index.ts

Lines changed: 128 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,70 +1,155 @@
1-
import type { AnyOrama, Results, SearchParams, Nullable, IAnswerSessionConfig as OSSAnswerSessionConfig } from '@orama/orama'
2-
import { search, AnswerSession as OSSAnswerSession } from '@orama/orama'
3-
import { OramaClient, ClientSearchParams, AnswerSessionParams as CloudAnswerSessionConfig, AnswerSession as CloudAnswerSession } from '@oramacloud/client'
1+
import type {
2+
AnyOrama,
3+
IAnswerSessionConfig as OSSAnswerSessionConfig,
4+
Nullable,
5+
Results,
6+
SearchParams,
7+
} from "@orama/orama";
8+
import { AnswerSession as OSSAnswerSession, search } from "@orama/orama";
9+
import {
10+
AnswerSession as CloudAnswerSession,
11+
AnswerSessionParams as CloudAnswerSessionConfig,
12+
ClientSearchParams,
13+
OramaClient,
14+
} from "@oramacloud/client";
15+
import type {
16+
AnswerSession as OramaCoreAnswerSession,
17+
AnswerSessionConfig as OramaCoreAnswerSessionConfig,
18+
SearchParams as OramaCoreSearchParams,
19+
} from "@orama/core";
20+
import { CollectionManager } from "@orama/core";
421

5-
export type OramaSwitchClient = AnyOrama | OramaClient
22+
export type OramaSwitchClient = AnyOrama | OramaClient | CollectionManager;
623

7-
export type ClientType = 'oss' | 'cloud'
24+
export type ClientType = "oss" | "cloud" | "core";
825

926
export type SearchConfig = {
10-
abortController?: AbortController
11-
fresh?: boolean
12-
debounce?: number
13-
}
27+
abortController?: AbortController;
28+
fresh?: boolean;
29+
debounce?: number;
30+
};
1431

1532
function isOramaClient(client: any): client is OramaClient {
16-
return client && typeof client === 'object' && 'api_key' in client && 'endpoint' in client;
33+
return client && typeof client === "object" && "api_key" in client &&
34+
"endpoint" in client;
35+
}
36+
37+
function isOramaCoreClient(client: any): client is CollectionManager {
38+
return client && client instanceof CollectionManager;
39+
}
40+
41+
function isOramaJSClient(client: any): client is AnyOrama {
42+
return client && typeof client === "object" && "id" in client &&
43+
"tokenizer" in client;
1744
}
1845

1946
export class Switch<T = OramaSwitchClient> {
20-
private invalidClientError = 'Invalid client. Expected either an OramaClient or an Orama OSS database.'
21-
client: OramaSwitchClient
22-
clientType: ClientType
23-
isCloud: boolean = false
24-
isOSS: boolean = false
47+
private invalidClientError =
48+
"Invalid client. Expected either an OramaClient, CollectionManager, or an Orama JS database.";
49+
client: OramaSwitchClient;
50+
clientType: ClientType;
51+
isCloud: boolean = false;
52+
isJS: boolean = false;
53+
isCore: boolean = false;
2554

2655
constructor(client: OramaSwitchClient) {
27-
this.client = client
28-
29-
if (isOramaClient(client)) {
30-
this.clientType = 'cloud'
31-
this.isCloud = true
32-
} else if (typeof client === 'object' && 'id' in client && 'tokenizer' in client) {
33-
this.clientType = 'oss'
34-
this.isOSS = true
35-
} else {
36-
throw new Error(this.invalidClientError)
56+
this.client = client;
57+
58+
switch (true) {
59+
case isOramaClient(client):
60+
this.clientType = "cloud";
61+
this.isCloud = true;
62+
break;
63+
case isOramaCoreClient(client):
64+
this.clientType = "core";
65+
this.isCore = true;
66+
break;
67+
case isOramaJSClient(client):
68+
this.clientType = "oss";
69+
this.isJS = true;
70+
break;
71+
default:
72+
throw new Error(this.invalidClientError);
3773
}
3874
}
3975

4076
async search<R = unknown>(
41-
params: T extends OramaClient ? ClientSearchParams : SearchParams<AnyOrama>,
42-
config?: SearchConfig
77+
params: T extends OramaClient ? ClientSearchParams
78+
: T extends CollectionManager ? SearchParams<AnyOrama>
79+
: never,
80+
config?: SearchConfig,
4381
): Promise<Nullable<Results<R>>> {
44-
if (this.isCloud) {
45-
return (this.client as OramaClient).search(params as T extends OramaClient ? ClientSearchParams : never, config)
46-
} else {
47-
return search(this.client as AnyOrama, params as SearchParams<AnyOrama>) as Promise<Nullable<Results<R>>>
82+
switch (true) {
83+
// OramaCloud - Old client
84+
case this.isCloud:
85+
return (this.client as OramaClient).search(
86+
params as T extends OramaClient ? ClientSearchParams : never,
87+
config,
88+
);
89+
90+
// OramaCore - New client
91+
case this.isCore: {
92+
const results = await (this.client as CollectionManager).search(
93+
params as T extends CollectionManager ? OramaCoreSearchParams : never,
94+
);
95+
96+
return results as unknown as Nullable<Results<R>>;
97+
}
98+
99+
// OramaJS - JavaScript client
100+
case this.isJS:
101+
return search(
102+
this.client as AnyOrama,
103+
params as SearchParams<AnyOrama>,
104+
) as Promise<Nullable<Results<R>>>;
105+
default:
106+
throw new Error(this.invalidClientError);
48107
}
49108
}
50109

51-
createAnswerSession(params: T extends OramaClient ? CloudAnswerSessionConfig : OSSAnswerSessionConfig): T extends OramaClient ? CloudAnswerSession<true> : OSSAnswerSession {
52-
if (this.isCloud) {
53-
const p = params as CloudAnswerSessionConfig
54-
return (this.client as OramaClient).createAnswerSession(p) as unknown as T extends OramaClient ? CloudAnswerSession<true> : OSSAnswerSession
55-
}
110+
createAnswerSession(
111+
params: T extends OramaClient ? CloudAnswerSessionConfig
112+
: OSSAnswerSessionConfig,
113+
): T extends OramaClient ? CloudAnswerSession<true>
114+
: T extends CollectionManager ? OramaCoreAnswerSession
115+
: OSSAnswerSession {
116+
switch (true) {
117+
// OramaCloud - Old client
118+
case this.isCloud: {
119+
const p = params as CloudAnswerSessionConfig;
120+
return (this.client as OramaClient).createAnswerSession(
121+
p,
122+
) as unknown as T extends OramaClient ? CloudAnswerSession<true>
123+
: T extends CollectionManager ? OramaCoreAnswerSession
124+
: OSSAnswerSession;
125+
}
126+
127+
// OramaCore - New client
128+
case this.isCore: {
129+
const p = params as OramaCoreAnswerSessionConfig;
130+
return (this.client as CollectionManager).createAnswerSession(
131+
p,
132+
) as unknown as T extends OramaClient ? CloudAnswerSession<true>
133+
: T extends CollectionManager ? OramaCoreAnswerSession
134+
: OSSAnswerSession;
135+
}
56136

57-
if (this.isOSS) {
58-
const p = params as OSSAnswerSessionConfig
59-
return new OSSAnswerSession(this.client as AnyOrama, {
137+
// OramaJS - JavaScript client
138+
case this.isJS: {
139+
const p = params as OSSAnswerSessionConfig;
140+
return new OSSAnswerSession(this.client as AnyOrama, {
60141
conversationID: p.conversationID,
61142
initialMessages: p.initialMessages,
62143
events: p.events,
63144
userContext: p.userContext,
64145
systemPrompt: p.systemPrompt,
65-
}) as unknown as T extends OramaClient ? CloudAnswerSession<true> : OSSAnswerSession
66-
}
146+
}) as unknown as T extends OramaClient ? CloudAnswerSession<true>
147+
: T extends CollectionManager ? OramaCoreAnswerSession
148+
: OSSAnswerSession;
149+
}
67150

68-
throw new Error(this.invalidClientError)
151+
default:
152+
throw new Error(this.invalidClientError);
153+
}
69154
}
70-
}
155+
}

‎packages/switch/tests/answer-session.test.ts

Lines changed: 85 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,71 +1,113 @@
1-
import { test } from 'node:test'
2-
import assert from 'node:assert/strict'
3-
import { OramaClient } from '@oramacloud/client'
4-
import { pluginSecureProxy } from '@orama/plugin-secure-proxy'
5-
import { create, insertMultiple } from '@orama/orama'
6-
import { Switch } from '../src/index.js'
1+
import type { AnswerSession as OramaCoreAnswerSession } from "@orama/core";
2+
import { test } from "node:test";
3+
import assert from "node:assert/strict";
4+
import { OramaClient } from "@oramacloud/client";
5+
import { pluginSecureProxy } from "@orama/plugin-secure-proxy";
6+
import { AnswerSession, create, insertMultiple } from "@orama/orama";
7+
import { Switch } from "../src/index.js";
8+
import { CollectionManager } from "@orama/core";
79

8-
const CLOUD_URL = process.env.ORAMA_CLOUD_E2E_URL
9-
const CLOUD_API_KEY = process.env.ORAMA_CLOUD_E2E_API_KEY
10-
const SECURE_PROXY_API_KEY = process.env.ORAMA_SECURE_PROXY_API_KEY
10+
const CLOUD_URL = process.env.ORAMA_CLOUD_E2E_URL;
11+
const CLOUD_API_KEY = process.env.ORAMA_CLOUD_E2E_API_KEY;
12+
const SECURE_PROXY_API_KEY = process.env.ORAMA_SECURE_PROXY_API_KEY;
13+
14+
const ORAMACORE_URL = process.env.ORAMACORE_E2E_URL;
15+
const ORAMACORE_COLLECTION_ID = process.env.ORAMACORE_E2E_COLLECTION_ID;
16+
const ORAMACORE_READ_API_KEY = process.env.ORAMACORE_E2E_READ_API_KEY;
1117

1218
if (!CLOUD_URL || !CLOUD_API_KEY) {
1319
console.log(
14-
'Skipping Orama Switch remote client test since ORAMA_CLOUD_E2E_URL and ORAMA_CLOUD_E2E_API_KEY are not set'
15-
)
16-
process.exit(0)
20+
"Skipping Orama Switch remote client test since ORAMA_CLOUD_E2E_URL and ORAMA_CLOUD_E2E_API_KEY are not set",
21+
);
22+
process.exit(0);
23+
}
24+
25+
if (
26+
!ORAMACORE_URL ||
27+
!ORAMACORE_COLLECTION_ID ||
28+
!ORAMACORE_READ_API_KEY
29+
) {
30+
console.log(
31+
"Skipping Orama Switch remote client test since ORAMACORE_E2E_URL, ORAMACORE_E2E_COLLECTION_ID and ORAMACORE_E2E_READ_API_KEY are not set",
32+
);
33+
process.exit(0);
1734
}
1835

19-
test('local client', async () => {
20-
const db = await create({
36+
test("local client", async () => {
37+
const db = create({
2138
schema: {
22-
name: 'string',
23-
age: 'number',
24-
embeddings: 'vector[384]'
39+
name: "string",
40+
age: "number",
41+
embeddings: "vector[384]",
2542
} as const,
2643
plugins: [
2744
pluginSecureProxy({
2845
apiKey: SECURE_PROXY_API_KEY!,
2946
chat: {
30-
model: 'openai/gpt-4'
47+
model: "openai/gpt-4",
3148
},
3249
embeddings: {
33-
defaultProperty: 'embeddings',
34-
model: 'openai/text-embedding-3-small',
50+
defaultProperty: "embeddings",
51+
model: "openai/text-embedding-3-small",
3552
},
36-
})
37-
]
38-
})
53+
}),
54+
],
55+
});
3956

4057
await insertMultiple(db, [
41-
{ name: 'Alice', age: 30 },
42-
{ name: 'Bob', age: 40 },
43-
{ name: 'Charlie', age: 50 }
44-
])
58+
{ name: "Alice", age: 30 },
59+
{ name: "Bob", age: 40 },
60+
{ name: "Charlie", age: 50 },
61+
]);
4562

4663
const answerSession = new Switch(db).createAnswerSession({
47-
systemPrompt: `You're an AI agent used to greet people. You will receive a name and will have to proceed generating a greeting message for that name. Use your fantasy.`,
48-
})
64+
systemPrompt:
65+
`You're an AI agent used to greet people. You will receive a name and will have to proceed generating a greeting message for that name. Use your fantasy.`,
66+
});
4967

50-
await answerSession.ask({ term: 'Bob', where: { age: { eq: 40 } } })
68+
await (answerSession as AnswerSession).ask({
69+
term: "Bob",
70+
where: { age: { eq: 40 } },
71+
});
5172

52-
const state = answerSession.state
73+
const state = answerSession.state;
5374

54-
assert(state.length === 1)
55-
assert(state[0].query === 'Bob')
56-
})
75+
assert(state.length === 1);
76+
assert(state[0].query === "Bob");
77+
});
5778

58-
test('remote client', async () => {
79+
test("Orama Cloud client", async () => {
5980
const cloudClient = new OramaClient({
6081
api_key: CLOUD_API_KEY,
61-
endpoint: CLOUD_URL
62-
})
82+
endpoint: CLOUD_URL,
83+
});
84+
85+
const answerSession = new Switch(cloudClient).createAnswerSession({});
86+
await (answerSession as AnswerSession).ask({ term: "What is Orama?" });
87+
88+
const state = answerSession.state;
89+
90+
assert(state.length === 1);
91+
assert(state[0].query === "What is Orama?");
92+
});
93+
94+
test("OramaCore client", async () => {
95+
const coreClient = new CollectionManager({
96+
url: ORAMACORE_URL,
97+
readAPIKey: ORAMACORE_READ_API_KEY,
98+
collectionID: ORAMACORE_COLLECTION_ID,
99+
});
63100

64-
const answerSession = new Switch(cloudClient).createAnswerSession({})
65-
await answerSession.ask({ term: 'What is Orama?' })
101+
const answerSession = new Switch(coreClient).createAnswerSession({});
102+
await (answerSession as OramaCoreAnswerSession).answer({
103+
query: "What is Orama?",
104+
interactionID: "1",
105+
sessionID: "1",
106+
visitorID: "1",
107+
});
66108

67-
const state = answerSession.state
109+
const state = answerSession.state;
68110

69-
assert(state.length === 1)
70-
assert(state[0].query === 'What is Orama?')
71-
})
111+
assert(state.length === 1);
112+
assert(state[0].query === "What is Orama?");
113+
});

0 commit comments

Comments
 (0)