feat: adds answer session to switch

packages/orama/src/methods/answer-session.ts

@@ -3,21 +3,21 @@ import type { AnyDocument, AnyOrama, Nullable, OramaPluginSync, SearchParams, Re
 import { createError } from "../errors.js"
 import { search } from "./search.js"
 
-type GenericContext =
+export type GenericContext =
  | string
  | object
 
-type MessageRole =
+export type MessageRole =
  | 'system'
  | 'user'
  | 'assistant'
 
-type Message = {
+export type Message = {
  role: MessageRole
  content: string
 }
 
-type Interaction<SourceT = AnyDocument> = {
+export type Interaction<SourceT = AnyDocument> = {
  interactionId: string
  query: string
  response: string
@@ -29,21 +29,21 @@ type Interaction<SourceT = AnyDocument> = {
  errorMessage: Nullable<string>
 }
 
-type AnswerSessionEvents<SourceT = AnyDocument> = {
+export type AnswerSessionEvents<SourceT = AnyDocument> = {
  onStateChange?: (state: Interaction<SourceT>[]) => void
 }
 
-type IAnswerSessionConfig<SourceT = AnyDocument> = {
+export type IAnswerSessionConfig<SourceT = AnyDocument> = {
  conversationID?: string
  systemPrompt?: string
  userContext?: GenericContext
  initialMessages?: Message[]
  events?: AnswerSessionEvents<SourceT>
 }
 
-type AskParams = SearchParams<AnyDocument>
+export type AskParams = SearchParams<AnyDocument>
 
-type RegenerateLastParams = {
+export type RegenerateLastParams = {
  stream: boolean
 }
 
@@ -60,8 +60,8 @@ export class AnswerSession<SourceT = AnyDocument> {
  private conversationID: string
  private messages: Message[] = []
  private events: AnswerSessionEvents<SourceT>
- private state: Interaction<SourceT>[] = []
  private initPromise?: Promise<true | null>
+ public state: Interaction<SourceT>[] = []
 
  constructor(db: AnyOrama, config: IAnswerSessionConfig<SourceT>) { 
  this.db = db
@@ -215,16 +215,16 @@ export class AnswerSession<SourceT = AnyDocument> {
  throw createError('PLUGIN_SECURE_PROXY_NOT_FOUND')
  }
 
- const pluginExtras = plugin.extra as { proxy: OramaProxy, pluginParams: { models: { chat: ChatModel } } }
+ const pluginExtras = plugin.extra as { proxy: OramaProxy, pluginParams: { chat: { model: ChatModel } } }
 
  this.proxy = pluginExtras.proxy
 
  if (this.config.systemPrompt) {
  this.messages.push({ role: 'system', content: this.config.systemPrompt })
  }
 
- if (pluginExtras?.pluginParams?.models?.chat) {
- this.chatModel = pluginExtras.pluginParams.models.chat
+ if (pluginExtras?.pluginParams?.chat?.model) {
+ this.chatModel = pluginExtras.pluginParams.chat.model
  } else {
  throw createError('PLUGIN_SECURE_PROXY_MISSING_CHAT_MODEL')
  }

packages/orama/src/types.ts

@@ -7,6 +7,18 @@ import { Sorter } from './components/sorter.js'
 import { Language } from './components/tokenizer/languages.js'
 import { Point } from './trees/bkd.js'
 
+export type {
+ IAnswerSessionConfig,
+ AnswerSession,
+ AnswerSessionEvents,
+ AskParams,
+ GenericContext,
+ Interaction,
+ Message,
+ MessageRole,
+ RegenerateLastParams
+} from './methods/answer-session.js'
+
 export { MODE_FULLTEXT_SEARCH, MODE_HYBRID_SEARCH, MODE_VECTOR_SEARCH } from './constants.js'
 
 export type { DefaultTokenizer } from './components/tokenizer/index.js'

packages/switch/package.json

@@ -16,7 +16,7 @@
  "files": ["dist"],
  "scripts": {
  "build": "tsup",
- "test": "tsx --test tests/search.test.ts"
+ "test": "tsx --test tests/answer-session.test.ts && tsx --test tests/search.test.ts"
  },
  "keywords": ["orama", "orama cloud"],
  "type": "module",
@@ -31,6 +31,7 @@
  },
  "license": "Apache-2.0",
  "devDependencies": {
+ "@orama/plugin-secure-proxy": "workspace:*",
  "tsup": "^7.2.0",
  "tsx": "^4.19.0"
  }

packages/switch/src/index.ts

@@ -1,6 +1,6 @@
-import type { AnyOrama, Results, SearchParams, Nullable } from '@orama/orama'
-import { search } from '@orama/orama'
-import { OramaClient, ClientSearchParams } from '@oramacloud/client'
+import type { AnyOrama, Results, SearchParams, Nullable, IAnswerSessionConfig as OSSAnswerSessionConfig } from '@orama/orama'
+import { search, AnswerSession as OSSAnswerSession } from '@orama/orama'
+import { OramaClient, ClientSearchParams, AnswerSessionParams as CloudAnswerSessionConfig, AnswerSession as CloudAnswerSession } from '@oramacloud/client'
 
 export type OramaSwitchClient = AnyOrama | OramaClient
 
@@ -13,6 +13,7 @@ export type SearchConfig = {
 }
 
 export class Switch<T = OramaSwitchClient> {
+ private invalidClientError = 'Invalid client. Expected either an OramaClient or an Orama OSS database.'
  client: OramaSwitchClient
  clientType: ClientType
  isCloud: boolean = false
@@ -28,7 +29,7 @@ export class Switch<T = OramaSwitchClient> {
  this.clientType = 'oss'
  this.isOSS = true
  } else {
- throw new Error('Invalid client. Expected either an OramaClient or an Orama OSS database.')
+ throw new Error(this.invalidClientError)
  }
  }
 
@@ -42,4 +43,24 @@ export class Switch<T = OramaSwitchClient> {
  return search(this.client as AnyOrama, params as SearchParams<AnyOrama>) as Promise<Nullable<Results<R>>>
  }
  }
-}
+
+ createAnswerSession(params: T extends OramaClient ? CloudAnswerSessionConfig : OSSAnswerSessionConfig): T extends OramaClient ? CloudAnswerSession : OSSAnswerSession {
+ if (this.isCloud) {
+ const p = params as CloudAnswerSessionConfig
+ return (this.client as OramaClient).createAnswerSession(p) as unknown as T extends OramaClient ? CloudAnswerSession : OSSAnswerSession
+ }
+
+ if (this.isOSS) {
+ const p = params as OSSAnswerSessionConfig
+ return new OSSAnswerSession(this.client as AnyOrama, {
+ conversationID: p.conversationID,
+ initialMessages: p.initialMessages,
+ events: p.events,
+ userContext: p.userContext,
+ systemPrompt: p.systemPrompt,
+ }) as unknown as T extends OramaClient ? CloudAnswerSession : OSSAnswerSession
+ }
+
+ throw new Error(this.invalidClientError)
+ }
+}

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

@@ -0,0 +1,71 @@
+import { test } from 'node:test'
+import assert from 'node:assert/strict'
+import { OramaClient } from '@oramacloud/client'
+import { pluginSecureProxy } from '@orama/plugin-secure-proxy'
+import { create, insertMultiple } from '@orama/orama'
+import { Switch } from '../src/index.js'
+
+const CLOUD_URL = process.env.ORAMA_CLOUD_E2E_URL
+const CLOUD_API_KEY = process.env.ORAMA_CLOUD_E2E_API_KEY
+const SECURE_PROXY_API_KEY = process.env.ORAMA_SECURE_PROXY_API_KEY
+
+if (!CLOUD_URL || !CLOUD_API_KEY) {
+ console.log(
+ 'Skipping Orama Switch remote client test since ORAMA_CLOUD_E2E_URL and ORAMA_CLOUD_E2E_API_KEY are not set'
+ )
+ process.exit(0)
+}
+
+test('local client', async () => {
+ const db = await create({
+ schema: {
+ name: 'string',
+ age: 'number',
+ embeddings: 'vector[384]'
+ } as const,
+ plugins: [
+ pluginSecureProxy({
+ apiKey: SECURE_PROXY_API_KEY!,
+ chat: {
+ model: 'openai/gpt-4'
+ },
+ embeddings: {
+ defaultProperty: 'embeddings',
+ model: 'openai/text-embedding-3-small',
+ },
+ })
+ ]
+ })
+
+ await insertMultiple(db, [
+ { name: 'Alice', age: 30 },
+ { name: 'Bob', age: 40 },
+ { name: 'Charlie', age: 50 }
+ ])
+
+ const answerSession = new Switch(db).createAnswerSession({
+ 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.`,
+ })
+
+ await answerSession.ask({ term: 'Bob', where: { age: { eq: 40 } } })
+
+ const state = answerSession.state
+
+ assert(state.length === 1)
+ assert(state[0].query === 'Bob')
+})
+
+test('remote client', async () => {
+ const cloudClient = new OramaClient({
+ api_key: CLOUD_API_KEY,
+ endpoint: CLOUD_URL
+ })
+
+ const answerSession = new Switch(cloudClient).createAnswerSession({})
+ await answerSession.ask({ term: 'What is Orama?' })
+
+ const state = answerSession.state
+
+ assert(state.length === 1)
+ assert(state[0].query === 'What is Orama?')
+})